|
|
@ -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) |
|
|
|