From fb220ee910ed171aea69addd5637f2bbde312bc8 Mon Sep 17 00:00:00 2001 From: tangfangchun Date: Wed, 18 Dec 2024 13:12:09 +0800 Subject: [PATCH] =?UTF-8?q?feat(sentinel):=20=E6=B7=BB=E5=8A=A0=20goframe?= =?UTF-8?q?=20=E6=A1=86=E6=9E=B6=E9=80=82=E9=85=8D=E5=99=A8?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- pkg/adapters/goframe/doc.go | 19 ++ pkg/adapters/goframe/go.mod | 54 ++++ pkg/adapters/goframe/middleware.go | 47 ++++ .../goframe/middleware_example_test.go | 37 +++ pkg/adapters/goframe/middleware_test.go | 235 ++++++++++++++++++ pkg/adapters/goframe/option.go | 35 +++ 6 files changed, 427 insertions(+) create mode 100644 pkg/adapters/goframe/doc.go create mode 100644 pkg/adapters/goframe/go.mod create mode 100644 pkg/adapters/goframe/middleware.go create mode 100644 pkg/adapters/goframe/middleware_example_test.go create mode 100644 pkg/adapters/goframe/middleware_test.go create mode 100644 pkg/adapters/goframe/option.go diff --git a/pkg/adapters/goframe/doc.go b/pkg/adapters/goframe/doc.go new file mode 100644 index 00000000..a0ba743b --- /dev/null +++ b/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 diff --git a/pkg/adapters/goframe/go.mod b/pkg/adapters/goframe/go.mod new file mode 100644 index 00000000..a86b6e72 --- /dev/null +++ b/pkg/adapters/goframe/go.mod @@ -0,0 +1,54 @@ +module github.com/tangfc/goframe-sentinel-adapter + +go 1.20 + +require ( + github.com/alibaba/sentinel-golang v1.0.4 + github.com/gogf/gf/v2 v2.8.2 + github.com/stretchr/testify v1.8.4 +) + +require ( + github.com/BurntSushi/toml v1.4.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/emirpasic/gods v1.18.1 // indirect + github.com/fatih/color v1.18.0 // indirect + github.com/fsnotify/fsnotify v1.7.0 // indirect + github.com/go-logr/logr v1.4.2 // 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.3 // indirect + github.com/grokify/html-strip-tags-go v0.1.0 // indirect + github.com/magiconair/properties v1.8.9 // 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.16 // 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.7 // 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.24.0 // indirect + go.opentelemetry.io/otel/metric v1.24.0 // indirect + go.opentelemetry.io/otel/sdk v1.24.0 // indirect + go.opentelemetry.io/otel/trace v1.24.0 // indirect + golang.org/x/net v0.32.0 // indirect + golang.org/x/sys v0.28.0 // indirect + golang.org/x/text v0.21.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 +) diff --git a/pkg/adapters/goframe/middleware.go b/pkg/adapters/goframe/middleware.go new file mode 100644 index 00000000..4133000c --- /dev/null +++ b/pkg/adapters/goframe/middleware.go @@ -0,0 +1,47 @@ +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 { + extractedName := options.resourceExtract(r) + if extractedName == "" { + extractedName = resourceName + } + resourceName = extractedName + } + + 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() + } +} diff --git a/pkg/adapters/goframe/middleware_example_test.go b/pkg/adapters/goframe/middleware_example_test.go new file mode 100644 index 00000000..7531dbf8 --- /dev/null +++ b/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( + // 自定义资源提取器 + WithResourceExtractor(func(r *ghttp.Request) string { + if res, ok := r.Header["X-Real-IP"]; ok && len(res) > 0 { + return res[0] + } + return "" + }), + // 自定义阻塞回退 + 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) +} diff --git a/pkg/adapters/goframe/middleware_test.go b/pkg/adapters/goframe/middleware_test.go new file mode 100644 index 00000000..57cf88fc --- /dev/null +++ b/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) + }) + } +} diff --git a/pkg/adapters/goframe/option.go b/pkg/adapters/goframe/option.go new file mode 100644 index 00000000..7818db38 --- /dev/null +++ b/pkg/adapters/goframe/option.go @@ -0,0 +1,35 @@ +package goframe + +import "github.com/gogf/gf/v2/net/ghttp" + +type ( + Option func(*options) + options struct { + resourceExtract func(*ghttp.Request) string + blockFallback func(*ghttp.Request) + } +) + +// evaluateOptions 评估选项 +func evaluateOptions(opts []Option) *options { + optCopy := &options{} + for _, opt := range opts { + opt(optCopy) + } + + return optCopy +} + +// WithResourceExtractor 设置资源提取器 +func WithResourceExtractor(fn func(*ghttp.Request) string) Option { + return func(opts *options) { + opts.resourceExtract = fn + } +} + +// WithBlockFallback 设置被流控的回退处理函数 +func WithBlockFallback(fn func(r *ghttp.Request)) Option { + return func(opts *options) { + opts.blockFallback = fn + } +}