Skip to content

Commit

Permalink
feat(goframe): sentinel adapter
Browse files Browse the repository at this point in the history
  • Loading branch information
tangfc committed Dec 18, 2024
1 parent b8d94d1 commit c91a893
Show file tree
Hide file tree
Showing 7 changed files with 908 additions and 0 deletions.
19 changes: 19 additions & 0 deletions pkg/adapters/goframe/doc.go
Original file line number Diff line number Diff line change
@@ -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
54 changes: 54 additions & 0 deletions pkg/adapters/goframe/go.mod
Original file line number Diff line number Diff line change
@@ -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
)
481 changes: 481 additions & 0 deletions pkg/adapters/goframe/go.sum

Large diffs are not rendered by default.

47 changes: 47 additions & 0 deletions pkg/adapters/goframe/middleware.go
Original file line number Diff line number Diff line change
@@ -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()
}
}
37 changes: 37 additions & 0 deletions pkg/adapters/goframe/middleware_example_test.go
Original file line number Diff line number Diff line change
@@ -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)
}
235 changes: 235 additions & 0 deletions pkg/adapters/goframe/middleware_test.go
Original file line number Diff line number Diff line change
@@ -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)
})
}
}
Loading

0 comments on commit c91a893

Please sign in to comment.