增加流控中间件

gin/middleware/limit.go

@@ -0,0 +1,83 @@
+// Package middleware ...
+//
+// Description : 基于redis的流控管理
+//
+// Author : go_developer@163.com<白茶清欢>
+//
+// Date : 2021-07-30 20:52 下午
+package middleware
+
+import (
+	"context"
+	"net/http"
+	"time"
+
+	"github.com/pkg/errors"
+
+	"github.com/gin-gonic/gin"
+	"github.com/go-redis/redis/v8"
+	"github.com/go-redis/redis_rate/v9"
+)
+
+// BuildRateLimitConfig 构建流控配置
+//
+// Author : go_developer@163.com<白茶清欢>
+//
+// Date : 10:35 下午 2021/7/30
+func BuildRateLimitConfig(rate int, burst int, period time.Duration) redis_rate.Limit {
+	return redis_rate.Limit{
+		Rate:   rate,
+		Burst:  burst,
+		Period: period,
+	}
+}
+
+// RateLimit gin 可以直接引用的中间件
+//
+// Author : go_developer@163.com<白茶清欢>
+//
+// Date : 10:35 下午 2021/7/30
+func RateLimit(redisInstance *redis.Client, key string, limitRule redis_rate.Limit, responseData map[string]interface{}) gin.HandlerFunc {
+	return func(ctx *gin.Context) {
+		var (
+			err error
+		)
+
+		if _, err = Allow(redisInstance, key, limitRule); nil != err {
+			if nil == responseData {
+				responseData = make(map[string]interface{})
+				responseData["code"] = -1
+				responseData["message"] = err.Error()
+				responseData["data"] = make(map[string]interface{})
+			}
+			ctx.AbortWithStatusJSON(http.StatusOK, responseData)
+			return
+		}
+
+		ctx.Next()
+	}
+}
+
+// Allow 是否允许访问
+//
+// Author : go_developer@163.com<白茶清欢>
+//
+// Date : 11:52 下午 2021/7/30
+func Allow(redisInstance *redis.Client, key string, limitRule redis_rate.Limit) (bool, error) {
+	var (
+		limiter *redis_rate.Limiter
+		result  *redis_rate.Result
+		err     error
+	)
+
+	limiter = redis_rate.NewLimiter(redisInstance)
+	if result, err = limiter.Allow(context.Background(), key, limitRule); nil != err {
+		return false, err
+	}
+
+	if result.Allowed > 0 {
+		return true, nil
+	}
+
+	return false, errors.New("rate limit! current request is not allowed!")
+}