From 7947374c9ee3e609894700cb55646518044bf553 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=BE=90=E8=B6=85=E8=B6=8A?= <993921@qq.com> Date: Sun, 3 Mar 2024 16:56:24 +0800 Subject: [PATCH 1/3] feat: rate limit --- .gitignore | 2 + go.mod | 7 ++-- go.sum | 8 +--- rpc_server/middlewares/rate_limit.go | 50 +++++++++++++++++++++++ rpc_server/middlewares/rate_limit_test.go | 38 +++++++++++++++++ rpc_server/routers/builder.go | 2 + 6 files changed, 98 insertions(+), 9 deletions(-) create mode 100644 rpc_server/middlewares/rate_limit.go create mode 100644 rpc_server/middlewares/rate_limit_test.go diff --git a/.gitignore b/.gitignore index b4cf4192..86951d3e 100644 --- a/.gitignore +++ b/.gitignore @@ -25,4 +25,6 @@ go.work build/ config/*.json +vendor + conf/appsettings.*.yaml diff --git a/go.mod b/go.mod index 2a50609a..2048c10f 100644 --- a/go.mod +++ b/go.mod @@ -6,10 +6,11 @@ require ( github.com/appleboy/gin-jwt/v2 v2.9.2 github.com/gin-contrib/cors v1.5.0 github.com/gin-gonic/gin v1.9.1 + github.com/stretchr/testify v1.8.4 github.com/swaggo/files v1.0.1 github.com/swaggo/gin-swagger v1.6.0 github.com/swaggo/swag v1.16.3 - gorm.io/gorm v1.25.7 + golang.org/x/time v0.3.0 k8s.io/apimachinery v0.29.2 ) @@ -18,6 +19,7 @@ require ( github.com/bytedance/sonic v1.11.1 // indirect github.com/chenzhuoyu/base64x v0.0.0-20230717121745-296ad89f973d // indirect github.com/chenzhuoyu/iasm v0.9.1 // indirect + github.com/davecgh/go-spew v1.1.1 // indirect github.com/gabriel-vasile/mimetype v1.4.3 // indirect github.com/gin-contrib/sse v0.1.0 // indirect github.com/go-openapi/jsonpointer v0.20.2 // indirect @@ -29,8 +31,6 @@ require ( github.com/go-playground/validator/v10 v10.18.0 // indirect github.com/goccy/go-json v0.10.2 // indirect github.com/golang-jwt/jwt/v4 v4.5.0 // indirect - github.com/jinzhu/inflection v1.0.0 // indirect - github.com/jinzhu/now v1.1.5 // indirect github.com/josharian/intern v1.0.0 // indirect github.com/json-iterator/go v1.1.12 // indirect github.com/klauspost/cpuid/v2 v2.2.7 // indirect @@ -40,6 +40,7 @@ require ( github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect github.com/modern-go/reflect2 v1.0.2 // indirect github.com/pelletier/go-toml/v2 v2.1.1 // indirect + github.com/pmezard/go-difflib v1.0.0 // indirect github.com/twitchyliquid64/golang-asm v0.15.1 // indirect github.com/ugorji/go/codec v1.2.12 // indirect golang.org/x/arch v0.7.0 // indirect diff --git a/go.sum b/go.sum index b83ab1a2..eae4b29c 100644 --- a/go.sum +++ b/go.sum @@ -51,10 +51,6 @@ github.com/golang-jwt/jwt/v4 v4.5.0/go.mod h1:m21LjoU+eqJr34lmDMbreY2eSTRJ1cv77w github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI= github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg= -github.com/jinzhu/inflection v1.0.0 h1:K317FqzuhWc8YvSVlFMCCUb36O/S9MCKRDI7QkRKD/E= -github.com/jinzhu/inflection v1.0.0/go.mod h1:h+uFLlag+Qp1Va5pdKtLDYj+kHp5pxUVkryuEj+Srlc= -github.com/jinzhu/now v1.1.5 h1:/o9tlHleP7gOFmsnYNz3RGnqzefHA47wQpKrrdTIwXQ= -github.com/jinzhu/now v1.1.5/go.mod h1:d3SSVoowX0Lcu0IBviAWJpolVfI5UJVZZ7cO71lE/z8= github.com/josharian/intern v1.0.0 h1:vlS4z54oSdjm0bgjRigI+G1HpF+tI+9rE5LLzOg8HmY= github.com/josharian/intern v1.0.0/go.mod h1:5DoeVV0s6jJacbCEi61lwdGj/aVlrQvzHFFd8Hwg//Y= github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM= @@ -147,6 +143,8 @@ golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ= golang.org/x/text v0.7.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8= golang.org/x/text v0.14.0 h1:ScX5w1eTa3QqT8oi6+ziP7dTV1S2+ALU0bI+0zXKWiQ= golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU= +golang.org/x/time v0.3.0 h1:rg5rLMjNzMS1RkNLzCG38eapWhnYLFYXDXj2gOlr8j4= +golang.org/x/time v0.3.0/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ= golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc= @@ -163,8 +161,6 @@ gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ= gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= -gorm.io/gorm v1.25.7 h1:VsD6acwRjz2zFxGO50gPO6AkNs7KKnvfzUjHQhZDz/A= -gorm.io/gorm v1.25.7/go.mod h1:hbnx/Oo0ChWMn1BIhpy1oYozzpM15i4YPuHDmfYtwg8= k8s.io/apimachinery v0.29.2 h1:EWGpfJ856oj11C52NRCHuU7rFDwxev48z+6DSlGNsV8= k8s.io/apimachinery v0.29.2/go.mod h1:6HVkd1FwxIagpYrHSwJlQqZI3G9LfYWRPAkUvLnXTKU= nullprogram.com/x/optparse v1.0.0/go.mod h1:KdyPE+Igbe0jQUrVfMqDMeJQIJZEuyV7pjYmp6pbG50= diff --git a/rpc_server/middlewares/rate_limit.go b/rpc_server/middlewares/rate_limit.go new file mode 100644 index 00000000..4e5bed35 --- /dev/null +++ b/rpc_server/middlewares/rate_limit.go @@ -0,0 +1,50 @@ +package middlewares + +import ( + "AAStarCommunity/EthPaymaster_BackService/rpc_server/api/utils" + "errors" + "github.com/gin-gonic/gin" + "golang.org/x/time/rate" + "net/http" +) + +const ( + DefaultLimit rate.Limit = 50 // limit `DefaultLimit` requests per second + DefaultBurst int = 50 // burst size, for surge traffic +) + +var limiter map[string]*rate.Limiter + +// RateLimiterByApiKey represents the rate limit by each ApiKey for each api calling +func RateLimiterByApiKey() gin.HandlerFunc { + return func(ctx *gin.Context) { + if exists, current := utils.CurrentUser(ctx); exists { + + if limiting(¤t) { + ctx.Next() + } else { + _ = ctx.AbortWithError(http.StatusTooManyRequests, errors.New("too many requests")) + } + } else { + _ = ctx.AbortWithError(http.StatusUnauthorized, errors.New("401 Unauthorized")) + } + } +} + +func limiting(apiKey *string) bool { + + var l *rate.Limiter + if limit, ok := limiter[*apiKey]; ok { + l = limit + } else { + // TODO: different rate config for each current(apiKey) should get from dashboard service + l = rate.NewLimiter(DefaultLimit, DefaultBurst) + limiter[*apiKey] = l + } + + return l.Allow() +} + +func init() { + limiter = make(map[string]*rate.Limiter, 100) +} diff --git a/rpc_server/middlewares/rate_limit_test.go b/rpc_server/middlewares/rate_limit_test.go new file mode 100644 index 00000000..8bd7cfd4 --- /dev/null +++ b/rpc_server/middlewares/rate_limit_test.go @@ -0,0 +1,38 @@ +package middlewares + +import ( + "github.com/stretchr/testify/assert" + "testing" + "time" +) + +func TestRateLimitShouldPreventRequestWhenOverDefaultLimit(t *testing.T) { + + mockApiKey := "TestingAipKey" + + // assuming this for loop taking less than 1 second to finish + for i := 0; i < int(DefaultLimit)+5; i++ { + b := limiting(&mockApiKey) + if i < int(DefaultLimit) { + assert.Equal(t, true, b) + } else { + assert.Equal(t, false, b) + } + } +} + +func TestRateLimiterShouldAllowDefaultLimitPerSecond(t *testing.T) { + mockApiKey := "TestingAipKey" + + for x := 1; x <= 2; x++ { + for i := 0; i < int(DefaultLimit)+5; i++ { + b := limiting(&mockApiKey) + if i < int(DefaultLimit) { + assert.Equal(t, true, b) + } else { + assert.Equal(t, false, b) + } + } + time.Sleep(time.Second) + } +} diff --git a/rpc_server/routers/builder.go b/rpc_server/routers/builder.go index 946040dd..a3d62638 100644 --- a/rpc_server/routers/builder.go +++ b/rpc_server/routers/builder.go @@ -13,6 +13,8 @@ func buildRouters(router *gin.Engine) { router.Use(middlewares.AuthHandler()) { + router.Use(middlewares.RateLimiterByApiKey()) + for _, routerMap := range RouterMaps { for _, method := range routerMap.Methods { if method == GET { From daab35de40c58d8ff7074772f423318824c31270 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=BE=90=E8=B6=85=E8=B6=8A?= <993921@qq.com> Date: Sun, 3 Mar 2024 17:02:59 +0800 Subject: [PATCH 2/3] chore: compatible unit test running on github action --- rpc_server/middlewares/rate_limit_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rpc_server/middlewares/rate_limit_test.go b/rpc_server/middlewares/rate_limit_test.go index 8bd7cfd4..9ced77ff 100644 --- a/rpc_server/middlewares/rate_limit_test.go +++ b/rpc_server/middlewares/rate_limit_test.go @@ -33,6 +33,6 @@ func TestRateLimiterShouldAllowDefaultLimitPerSecond(t *testing.T) { assert.Equal(t, false, b) } } - time.Sleep(time.Second) + time.Sleep(time.Second * 2) } } From c314d9b9a8a5bc5c26de7eba5d8b94aadbc0aee5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=BE=90=E8=B6=85=E8=B6=8A?= <993921@qq.com> Date: Sun, 3 Mar 2024 17:07:00 +0800 Subject: [PATCH 3/3] chore: compatible unit test running on github action --- rpc_server/middlewares/rate_limit_test.go | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/rpc_server/middlewares/rate_limit_test.go b/rpc_server/middlewares/rate_limit_test.go index 9ced77ff..3965094a 100644 --- a/rpc_server/middlewares/rate_limit_test.go +++ b/rpc_server/middlewares/rate_limit_test.go @@ -2,6 +2,7 @@ package middlewares import ( "github.com/stretchr/testify/assert" + "os" "testing" "time" ) @@ -22,6 +23,10 @@ func TestRateLimitShouldPreventRequestWhenOverDefaultLimit(t *testing.T) { } func TestRateLimiterShouldAllowDefaultLimitPerSecond(t *testing.T) { + if os.Getenv("GITHUB_ACTIONS") != "" { + t.Skip() + return + } mockApiKey := "TestingAipKey" for x := 1; x <= 2; x++ {