goframe framework adapter for sentinel

pkg/adapters/goframe/doc.go

@@ -0,0 +1,19 @@
+// Package goframe provides Sentinel middleware for GoFrame.
+//
+// Users may register SentinelMiddleware to the GoFrame server, like:
+//
+//	import (
+//		sentinelPlugin "github.com/your-repo/goframe-sentinel-adapter"
+//		"github.com/gogf/gf/v2/frame/g"
+//		"github.com/gogf/gf/v2/net/ghttp"
+//	)
+//
+//	s := g.Server()
+//	s.Use(ghttp.MiddlewareHandlerFunc(sentinelPlugin.SentinelMiddleware()))
+//
+// The plugin extracts "HttpMethod:FullPath" as the resource name by default (e.g. GET:/foo/:id).
+// Users may provide a customized resource name extractor when creating new SentinelMiddleware (via options).
+//
+// Fallback logic: the plugin will return "429 Too Many Requests" status code if the current request is blocked by Sentinel rules.
+// Users may also provide customized fallback logic via WithBlockFallback(handler) options.
+package goframe

pkg/adapters/goframe/go.mod

@@ -0,0 +1,52 @@
+module github.com/alibaba/sentinel-golang/pkg/adapters/goframe
+
+go 1.18
+
+require (
+	github.com/alibaba/sentinel-golang v1.0.4
+	github.com/gogf/gf/v2 v2.6.4
+	github.com/stretchr/testify v1.8.2
+)
+
+require (
+	github.com/BurntSushi/toml v1.2.0 // indirect
+	github.com/StackExchange/wmi v0.0.0-20190523213315-cbe66965904d // indirect
+	github.com/beorn7/perks v1.0.1 // indirect
+	github.com/cespare/xxhash/v2 v2.1.1 // indirect
+	github.com/clbanning/mxj/v2 v2.7.0 // indirect
+	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/fatih/color v1.15.0 // indirect
+	github.com/fsnotify/fsnotify v1.7.0 // indirect
+	github.com/go-logr/logr v1.2.4 // indirect
+	github.com/go-logr/stdr v1.2.2 // indirect
+	github.com/go-ole/go-ole v1.2.4 // indirect
+	github.com/golang/protobuf v1.4.3 // indirect
+	github.com/google/uuid v1.1.1 // indirect
+	github.com/gorilla/websocket v1.5.0 // indirect
+	github.com/grokify/html-strip-tags-go v0.0.1 // indirect
+	github.com/magiconair/properties v1.8.6 // indirect
+	github.com/mattn/go-colorable v0.1.13 // indirect
+	github.com/mattn/go-isatty v0.0.20 // indirect
+	github.com/mattn/go-runewidth v0.0.15 // indirect
+	github.com/matttproud/golang_protobuf_extensions v1.0.1 // indirect
+	github.com/olekukonko/tablewriter v0.0.5 // indirect
+	github.com/pkg/errors v0.9.1 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
+	github.com/prometheus/client_golang v1.9.0 // indirect
+	github.com/prometheus/client_model v0.2.0 // indirect
+	github.com/prometheus/common v0.15.0 // indirect
+	github.com/prometheus/procfs v0.2.0 // indirect
+	github.com/rivo/uniseg v0.4.4 // indirect
+	github.com/shirou/gopsutil/v3 v3.21.6 // indirect
+	github.com/tklauser/go-sysconf v0.3.6 // indirect
+	github.com/tklauser/numcpus v0.2.2 // indirect
+	go.opentelemetry.io/otel v1.14.0 // indirect
+	go.opentelemetry.io/otel/sdk v1.14.0 // indirect
+	go.opentelemetry.io/otel/trace v1.14.0 // indirect
+	golang.org/x/net v0.17.0 // indirect
+	golang.org/x/sys v0.13.0 // indirect
+	golang.org/x/text v0.13.0 // indirect
+	google.golang.org/protobuf v1.23.0 // indirect
+	gopkg.in/yaml.v2 v2.3.0 // indirect
+	gopkg.in/yaml.v3 v3.0.1 // indirect
+)

pkg/adapters/goframe/middleware.go

@@ -0,0 +1,43 @@
+package goframe
+
+import (
+	"github.com/alibaba/sentinel-golang/api"
+	"github.com/alibaba/sentinel-golang/core/base"
+	"github.com/gogf/gf/v2/net/ghttp"
+	"net/http"
+)
+
+// SentinelMiddleware returns new ghttp.HandlerFunc
+// Default resource name is {method}:{path}, such as "GET:/api/users/:id"
+// Default block fallback is returning 429 status code
+// Define your own behavior by setting options
+func SentinelMiddleware(opts ...Option) ghttp.HandlerFunc {
+	options := evaluateOptions(opts)
+	return func(r *ghttp.Request) {
+		resourceName := r.Method + ":" + r.URL.Path
+
+		if options.resourceExtract != nil {
+			resourceName = options.resourceExtract(r)
+		}
+
+		entry, err := api.Entry(
+			resourceName,
+			api.WithResourceType(base.ResTypeWeb),
+			api.WithTrafficType(base.Inbound),
+		)
+
+		if err != nil {
+			if options.blockFallback != nil {
+				options.blockFallback(r)
+			} else {
+				r.Response.WriteHeader(http.StatusTooManyRequests)
+				r.Response.Writeln("Too Many Requests")
+			}
+			return
+		}
+
+		defer entry.Exit()
+
+		r.Middleware.Next()
+	}
+}

pkg/adapters/goframe/middleware_example_test.go

@@ -0,0 +1,37 @@
+package goframe
+
+import (
+	"github.com/gogf/gf/v2/frame/g"
+	"github.com/gogf/gf/v2/net/ghttp"
+)
+
+func Example() {
+	s := g.Server()
+	s.Use(
+		SentinelMiddleware(
+			// customize resource extractor if required
+			WithResourceExtractor(func(r *ghttp.Request) string {
+				if res, ok := r.Header["X-Real-IP"]; ok && len(res) > 0 {
+					return res[0]
+				}
+				return ""
+			}),
+			// customize block fallback if required
+			WithBlockFallback(func(r *ghttp.Request) {
+				r.Response.WriteHeader(400)
+				r.Response.WriteJson(map[string]interface{}{
+					"err":  "too many requests; the quota used up",
+					"code": 10222,
+				})
+			}),
+		),
+	)
+
+	s.Group("/", func(group *ghttp.RouterGroup) {
+		group.GET("/test", func(r *ghttp.Request) {
+			r.Response.Write("hello sentinel")
+		})
+	})
+
+	s.SetPort(8199)
+}

pkg/adapters/goframe/middleware_test.go

@@ -0,0 +1,235 @@
+package goframe
+
+import (
+	sentinel "github.com/alibaba/sentinel-golang/api"
+	"github.com/alibaba/sentinel-golang/core/flow"
+	"github.com/gogf/gf/v2/frame/g"
+	"github.com/gogf/gf/v2/net/ghttp"
+	"github.com/stretchr/testify/assert"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+)
+
+func initSentinel(t *testing.T) {
+	err := sentinel.InitDefault()
+	if err != nil {
+		t.Fatalf("Unexpected error: %+v", err)
+	}
+
+	_, err = flow.LoadRules([]*flow.Rule{
+		{
+			Resource:               "GET:/",
+			Threshold:              1.0,
+			TokenCalculateStrategy: flow.Direct,
+			ControlBehavior:        flow.Reject,
+			StatIntervalInMs:       1000,
+		},
+		{
+			Resource:               "/api/users/:id",
+			Threshold:              0.0,
+			TokenCalculateStrategy: flow.Direct,
+			ControlBehavior:        flow.Reject,
+			StatIntervalInMs:       1000,
+		},
+		{
+			Resource:               "GET:/ping",
+			Threshold:              0.0,
+			TokenCalculateStrategy: flow.Direct,
+			ControlBehavior:        flow.Reject,
+			StatIntervalInMs:       1000,
+		},
+	})
+	if err != nil {
+		t.Fatalf("Unexpected error: %+v", err)
+		return
+	}
+}
+
+// go test -run ^TestSentinelMiddlewareDefault -v
+func TestSentinelMiddlewareDefault(t *testing.T) {
+	type args struct {
+		opts    []Option
+		method  string
+		path    string
+		reqPath string
+		handler func(r *ghttp.Request)
+		body    io.Reader
+	}
+	type want struct {
+		code int
+	}
+	var (
+		tests = []struct {
+			name string
+			args args
+			want want
+		}{
+			{
+				name: "default get",
+				args: args{
+					opts: []Option{
+						WithResourceExtractor(func(r *ghttp.Request) string {
+							return r.Router.Uri
+						}),
+					},
+					method:  http.MethodPost,
+					path:    "/",
+					reqPath: "/",
+					handler: func(r *ghttp.Request) {
+						r.Response.WriteStatusExit(http.StatusOK, "/")
+					},
+					body: nil,
+				},
+				want: want{
+					code: http.StatusOK,
+				},
+			},
+		}
+	)
+
+	initSentinel(t)
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			s := g.Server()
+			s.SetRouteOverWrite(true)
+			s.Group("/", func(group *ghttp.RouterGroup) {
+				group.Middleware(SentinelMiddleware(tt.args.opts...))
+				group.ALL(tt.args.path, tt.args.handler)
+			})
+			s.Start()
+
+			r := httptest.NewRequest(tt.args.method, tt.args.reqPath, tt.args.body)
+			w := httptest.NewRecorder()
+			s.ServeHTTP(w, r)
+
+			assert.Equal(t, tt.want.code, w.Code)
+		})
+	}
+}
+
+// go test -run ^TestSentinelMiddlewareExtractor -v
+func TestSentinelMiddlewareExtractor(t *testing.T) {
+	type args struct {
+		opts    []Option
+		method  string
+		path    string
+		reqPath string
+		handler func(r *ghttp.Request)
+		body    io.Reader
+	}
+	type want struct {
+		code int
+	}
+	var (
+		tests = []struct {
+			name string
+			args args
+			want want
+		}{
+			{
+				name: "customize resource extract",
+				args: args{
+					opts: []Option{
+						WithResourceExtractor(func(r *ghttp.Request) string {
+							return r.Router.Uri
+						}),
+					},
+					method:  http.MethodPost,
+					path:    "/api/users/:id",
+					reqPath: "/api/users/123",
+					handler: func(r *ghttp.Request) {
+						r.Response.WriteStatusExit(http.StatusOK, "/api/users/123")
+					},
+					body: nil,
+				},
+				want: want{
+					code: http.StatusTooManyRequests,
+				},
+			},
+		}
+	)
+
+	initSentinel(t)
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			s := g.Server()
+			s.SetRouteOverWrite(true)
+			s.Group("/", func(group *ghttp.RouterGroup) {
+				group.Middleware(SentinelMiddleware(tt.args.opts...))
+				group.ALL(tt.args.path, tt.args.handler)
+			})
+			s.Start()
+
+			r := httptest.NewRequest(tt.args.method, tt.args.reqPath, tt.args.body)
+			w := httptest.NewRecorder()
+			s.ServeHTTP(w, r)
+			assert.Equal(t, tt.want.code, w.Code)
+		})
+	}
+}
+
+// go test -run ^TestSentinelMiddlewareFallback -v
+func TestSentinelMiddlewareFallback(t *testing.T) {
+	type args struct {
+		opts    []Option
+		method  string
+		path    string
+		reqPath string
+		handler func(r *ghttp.Request)
+		body    io.Reader
+	}
+	type want struct {
+		code int
+	}
+	var (
+		tests = []struct {
+			name string
+			args args
+			want want
+		}{
+			{
+				name: "customize block fallback",
+				args: args{
+					opts: []Option{
+						WithBlockFallback(func(r *ghttp.Request) {
+							r.Response.WriteStatus(http.StatusBadRequest, "/ping")
+						}),
+					},
+					method:  http.MethodGet,
+					path:    "/ping",
+					reqPath: "/ping",
+					handler: func(r *ghttp.Request) {
+						r.Response.WriteStatus(http.StatusOK, "ping")
+					},
+					body: nil,
+				},
+				want: want{
+					code: http.StatusBadRequest,
+				},
+			},
+		}
+	)
+
+	initSentinel(t)
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			s := g.Server()
+			s.SetRouteOverWrite(true)
+			s.Group("/", func(group *ghttp.RouterGroup) {
+				group.Middleware(SentinelMiddleware(tt.args.opts...))
+				group.ALL(tt.args.path, tt.args.handler)
+			})
+			s.Start()
+
+			r := httptest.NewRequest(tt.args.method, tt.args.reqPath, tt.args.body)
+			w := httptest.NewRecorder()
+			s.ServeHTTP(w, r)
+			assert.Equal(t, tt.want.code, w.Code)
+		})
+	}
+}