From 17b4f99a50ff7e750f74ab595b5f0fbc08ab2bbf Mon Sep 17 00:00:00 2001 From: viletyy Date: Sun, 13 Jun 2021 22:47:19 +0800 Subject: [PATCH] add: limitable and timeout middleware --- go.mod | 1 + internal/controller/api/{v1 => }/auth.go | 8 ++-- internal/middleware/context_timeout.go | 24 ++++++++++++ internal/middleware/limiter.go | 31 ++++++++++++++++ internal/routers/router.go | 38 +++++++++++++------ pkg/limiter/limiter.go | 31 ++++++++++++++++ pkg/limiter/method_limiter.go | 47 ++++++++++++++++++++++++ 7 files changed, 164 insertions(+), 16 deletions(-) rename internal/controller/api/{v1 => }/auth.go (90%) create mode 100644 internal/middleware/context_timeout.go create mode 100644 internal/middleware/limiter.go create mode 100644 pkg/limiter/limiter.go create mode 100644 pkg/limiter/method_limiter.go diff --git a/go.mod b/go.mod index e32c8fc..151283e 100644 --- a/go.mod +++ b/go.mod @@ -12,6 +12,7 @@ require ( github.com/go-playground/validator/v10 v10.6.1 github.com/go-redis/redis v6.15.9+incompatible github.com/jinzhu/gorm v1.9.16 + github.com/juju/ratelimit v1.0.1 github.com/lestrrat-go/file-rotatelogs v2.4.0+incompatible github.com/lestrrat-go/strftime v1.0.4 // indirect github.com/onsi/ginkgo v1.16.4 // indirect diff --git a/internal/controller/api/v1/auth.go b/internal/controller/api/auth.go similarity index 90% rename from internal/controller/api/v1/auth.go rename to internal/controller/api/auth.go index cde7ddc..337c89b 100644 --- a/internal/controller/api/v1/auth.go +++ b/internal/controller/api/auth.go @@ -1,10 +1,10 @@ /* * @Date: 2021-06-10 18:58:25 * @LastEditors: viletyy - * @LastEditTime: 2021-06-11 15:44:18 - * @FilePath: /potato/internal/controller/api/v1/auth.go + * @LastEditTime: 2021-06-13 22:41:01 + * @FilePath: /potato/internal/controller/api/auth.go */ -package v1 +package api import ( "github.com/gin-gonic/gin" @@ -21,7 +21,7 @@ import ( // @Param app_key formData string true "app key" // @Param app_secret formData string true "app secret" // @Success 200 {object} errcode.Error "请求成功" -// @Router /v1/auth [post] +// @Router /auth [post] func GetAuth(c *gin.Context) { param := service.AuthRequest{} response := app.NewResponse(c) diff --git a/internal/middleware/context_timeout.go b/internal/middleware/context_timeout.go new file mode 100644 index 0000000..eb64692 --- /dev/null +++ b/internal/middleware/context_timeout.go @@ -0,0 +1,24 @@ +/* + * @Date: 2021-06-13 22:35:30 + * @LastEditors: viletyy + * @LastEditTime: 2021-06-13 22:37:09 + * @FilePath: /potato/internal/middleware/context_timeout.go + */ +package middleware + +import ( + "context" + "time" + + "github.com/gin-gonic/gin" +) + +func ContextTimeout(t time.Duration) func(c *gin.Context) { + return func(c *gin.Context) { + ctx, cancel := context.WithTimeout(c.Request.Context(), t) + defer cancel() + + c.Request = c.Request.WithContext(ctx) + c.Next() + } +} diff --git a/internal/middleware/limiter.go b/internal/middleware/limiter.go new file mode 100644 index 0000000..cbec7a3 --- /dev/null +++ b/internal/middleware/limiter.go @@ -0,0 +1,31 @@ +/* + * @Date: 2021-06-13 22:27:24 + * @LastEditors: viletyy + * @LastEditTime: 2021-06-13 22:29:27 + * @FilePath: /potato/internal/middleware/limiter.go + */ +package middleware + +import ( + "github.com/gin-gonic/gin" + "github.com/viletyy/potato/pkg/app" + "github.com/viletyy/potato/pkg/errcode" + "github.com/viletyy/potato/pkg/limiter" +) + +func RateLimiter(l limiter.LimiterInterface) gin.HandlerFunc { + return func(c *gin.Context) { + key := l.Key(c) + if bucket, ok := l.GetBucket(key); ok { + count := bucket.TakeAvailable(1) + if count == 0 { + response := app.NewResponse(c) + response.ToErrorResponse(errcode.TooManyRequests) + c.Abort() + return + } + } + + c.Next() + } +} diff --git a/internal/routers/router.go b/internal/routers/router.go index 438a732..040c10d 100644 --- a/internal/routers/router.go +++ b/internal/routers/router.go @@ -1,13 +1,14 @@ /* * @Date: 2021-03-21 19:54:57 * @LastEditors: viletyy - * @LastEditTime: 2021-06-13 22:03:27 + * @LastEditTime: 2021-06-13 22:46:46 * @FilePath: /potato/internal/routers/router.go */ package routers import ( "net/http" + "time" "github.com/gin-gonic/gin" _ "github.com/swaggo/gin-swagger" @@ -16,29 +17,42 @@ import ( _ "github.com/swaggo/gin-swagger/swaggerFiles" _ "github.com/viletyy/potato/docs" "github.com/viletyy/potato/global" + "github.com/viletyy/potato/internal/controller/api" v1 "github.com/viletyy/potato/internal/controller/api/v1" "github.com/viletyy/potato/internal/middleware" + "github.com/viletyy/potato/pkg/limiter" ) var ( - Engine = gin.Default() - V1RouterGroup = Engine.Group("../api/v1") + Engine = gin.Default() + V1RouterGroup = Engine.Group("../api/v1") + methodLimiters = limiter.NewMethodLimiter().AddBuckets(limiter.LimiterBucketRule{ + Key: "/auth", + FillInterval: time.Second, + Capacity: 10, + Quantum: 10, + }) ) func InitRouter() *gin.Engine { - Engine.Use(gin.Logger()) + gin.SetMode(global.GO_CONFIG.App.RunMode) // 设置运行模式 - Engine.Use(gin.Recovery()) - Engine.Use(middleware.Translations()) + if global.GO_CONFIG.App.RunMode == "debug" { + Engine.Use(gin.Logger()) // 设置log + Engine.Use(gin.Recovery()) // 设置recovery + } else { + Engine.Use(middleware.AccessLog()) + Engine.Use(middleware.Recovery()) + } + Engine.Use(middleware.AppInfo()) // 设置app信息 + Engine.Use(middleware.RateLimiter(methodLimiters)) // 设置限流控制 + Engine.Use(middleware.ContextTimeout(60 * time.Second)) // 设置统一超时控制 + Engine.Use(middleware.Translations()) // 设置自定义验证 + Engine.Use(middleware.CORS()) // 设置跨域 - gin.SetMode(global.GO_CONFIG.App.RunMode) - - Engine.Use(middleware.CORS()) - Engine.Use(middleware.AccessLog()) - Engine.Use(middleware.AppInfo()) Engine.StaticFS("/static", http.Dir(global.GO_CONFIG.App.UploadSavePath)) - Engine.POST("/api/v1/auth", v1.GetAuth) + Engine.POST("/api/auth", api.GetAuth) Engine.GET("/swagger/*any", ginSwagger.WrapHandler(swaggerFiles.Handler)) V1InitModule() diff --git a/pkg/limiter/limiter.go b/pkg/limiter/limiter.go new file mode 100644 index 0000000..358284c --- /dev/null +++ b/pkg/limiter/limiter.go @@ -0,0 +1,31 @@ +/* + * @Date: 2021-06-13 22:06:52 + * @LastEditors: viletyy + * @LastEditTime: 2021-06-13 22:13:35 + * @FilePath: /potato/pkg/limiter/limiter.go + */ +package limiter + +import ( + "time" + + "github.com/gin-gonic/gin" + "github.com/juju/ratelimit" +) + +type LimiterInterface interface { + Key(c *gin.Context) string // 获取对应的限流器的键值对名称。 + GetBucket(key string) (*ratelimit.Bucket, bool) // 获取令牌桶 + AddBuckets(rules ...LimiterBucketRule) LimiterInterface // 新增多个令牌桶 +} + +type Limiter struct { + limiterBuckets map[string]*ratelimit.Bucket +} + +type LimiterBucketRule struct { + Key string // 自定义键值对名称 + FillInterval time.Duration // 间隔多久时间放N个令牌 + Capacity int64 // 令牌桶的容量 + Quantum int64 // 每次到达间隔时间后所放的具体令牌数量 +} diff --git a/pkg/limiter/method_limiter.go b/pkg/limiter/method_limiter.go new file mode 100644 index 0000000..07d3604 --- /dev/null +++ b/pkg/limiter/method_limiter.go @@ -0,0 +1,47 @@ +/* + * @Date: 2021-06-13 22:13:53 + * @LastEditors: viletyy + * @LastEditTime: 2021-06-13 22:26:09 + * @FilePath: /potato/pkg/limiter/method_limiter.go + */ +package limiter + +import ( + "strings" + + "github.com/gin-gonic/gin" + "github.com/juju/ratelimit" +) + +type MethodLimiter struct { + *Limiter +} + +func NewMethodLimiter() LimiterInterface { + return MethodLimiter{ + Limiter: &Limiter{limiterBuckets: make(map[string]*ratelimit.Bucket)}, + } +} + +func (l MethodLimiter) Key(c *gin.Context) string { + uri := c.Request.RequestURI + index := strings.Index(uri, "?") + if index == -1 { + return uri + } + return uri[:index] +} + +func (l MethodLimiter) GetBucket(key string) (*ratelimit.Bucket, bool) { + bucket, ok := l.limiterBuckets[key] + return bucket, ok +} + +func (l MethodLimiter) AddBuckets(rules ...LimiterBucketRule) LimiterInterface { + for _, rule := range rules { + if _, ok := l.limiterBuckets[rule.Key]; !ok { + l.limiterBuckets[rule.Key] = ratelimit.NewBucketWithQuantum(rule.FillInterval, rule.Capacity, rule.Quantum) + } + } + return l +}