feat:优化限流中间件

server/middleware/limit_ip.go

@@ -3,13 +3,58 @@ package middleware
 import (
 	"context"
 	"errors"
+	"net/http"
+	"time"
+
+	"go.uber.org/zap"
+
 	"github.com/flipped-aurora/gin-vue-admin/server/global"
 	"github.com/flipped-aurora/gin-vue-admin/server/model/common/response"
 	"github.com/gin-gonic/gin"
-	"time"
 )
 
-// ip限制
+type LimitConfig struct {
+	// GenerationKey 根据业务生成key 下面CheckOrMark查询生成
+	GenerationKey func(c *gin.Context) string
+	// 检查函数,用户可修改具体逻辑,更加灵活
+	CheckOrMark func(key string, Expire int) bool
+	// Expire key 过期时间
+	Expire int
+	// Limit 周期时间
+	Limit int
+}
+
+func (l *LimitConfig) LimitWithTime() gin.HandlerFunc {
+	return func(c *gin.Context) {
+		if l.CheckOrMark(l.GenerationKey(c), l.Expire) {
+			c.Next()
+		} else {
+			c.JSON(http.StatusOK, gin.H{"code": response.ERROR, "msg": "操作频繁,请稍后再试"})
+			c.Abort()
+			return
+		}
+	}
+}
+
+// DefaultGenerationKey 默认生成key
+func DefaultGenerationKey(c *gin.Context) string {
+	return "GVA_Limit" + c.ClientIP()
+}
+
+func DefaultCheckOrMark(key string, expire int, limit int) bool {
+	// 判断是否开启redis
+	if global.GVA_REDIS == nil {
+		return true
+	}
+	if err := SetLimitWithTime(key, limit, time.Duration(expire)*time.Second); err != nil {
+		global.GVA_LOG.Error("limit", zap.Error(err))
+		return false
+	}
+	return true
+
+}
+
+// IPLimit ip限制
 func IPLimit() gin.HandlerFunc {
 	return func(c *gin.Context) {
 		key := "RequestClientIPLimit===" + c.ClientIP()
@@ -29,7 +74,10 @@ func IPLimit() gin.HandlerFunc {
 // 设置访问次数
 func SetLimitWithTime(key string, limit int, expiration time.Duration) error {
 	count, err := global.GVA_REDIS.Exists(context.Background(), key).Result()
-	if err != nil || count == 0 {
+	if err != nil {
+		return err
+	}
+	if count == 0 {
 		pipe := global.GVA_REDIS.TxPipeline()
 		pipe.Incr(context.Background(), key)
 		pipe.Expire(context.Background(), key, expiration)