diff --git a/server/middleware/limit_ip.go b/server/middleware/limit_ip.go index 0e18cac9..2bbdcb4f 100644 --- a/server/middleware/limit_ip.go +++ b/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)