pkg/adapters/go-zero/routing_middleware.go

package go_zero

/*
This is an example based on go-zero routing middleware tempalte,
which is generated by goctl via

		goctl api go -api xxx.api

The xxx.api file defines routing middlewares by

		@server(
			middleware: SentinelRoute
		)

For more information, please refer to go-zero's documentation.
*/

import (
	"net/http"

	sentinel "github.com/alibaba/sentinel-golang/api"
	"github.com/alibaba/sentinel-golang/core/base"
)

type SentinelRouteMiddleware struct {
}

func NewSentinelRouteMiddleware() *SentinelRouteMiddleware {
	return &SentinelRouteMiddleware{}
}

func (m *SentinelRouteMiddleware) Handle(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		resourceName := r.Method + ":" + r.URL.Path
		entry, blockErr := sentinel.Entry(
			resourceName,
			sentinel.WithResourceType(base.ResTypeWeb),
			sentinel.WithTrafficType(base.Inbound),
		)
		if blockErr != nil {
			http.Error(w, "Blocked by Sentinel", http.StatusTooManyRequests)
			return
		}
		defer entry.Exit()
		next(w, r)
	}
}