Browse Source

feat:优化初始化limit对象

main
songzhibin97 3 years ago
parent
commit
c6501c2d0a
  1. 46
      server/middleware/limit_ip.go

46
server/middleware/limit_ip.go

@ -17,21 +17,21 @@ type LimitConfig struct {
// GenerationKey 根据业务生成key 下面CheckOrMark查询生成 // GenerationKey 根据业务生成key 下面CheckOrMark查询生成
GenerationKey func(c *gin.Context) string GenerationKey func(c *gin.Context) string
// 检查函数,用户可修改具体逻辑,更加灵活 // 检查函数,用户可修改具体逻辑,更加灵活
CheckOrMark func(key string, Expire int) bool
CheckOrMark func(key string, expire int, limit int) error
// Expire key 过期时间 // Expire key 过期时间
Expire int Expire int
// Limit 周期时间 // Limit 周期时间
Limit int Limit int
} }
func (l *LimitConfig) LimitWithTime() gin.HandlerFunc {
func (l LimitConfig) LimitWithTime() gin.HandlerFunc {
return func(c *gin.Context) { 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": "操作频繁,请稍后再试"})
if err := l.CheckOrMark(l.GenerationKey(c), l.Expire, l.Limit); err != nil {
c.JSON(http.StatusOK, gin.H{"code": response.ERROR, "msg": err})
c.Abort() c.Abort()
return return
} else {
c.Next()
} }
} }
} }
@ -41,37 +41,29 @@ func DefaultGenerationKey(c *gin.Context) string {
return "GVA_Limit" + c.ClientIP() return "GVA_Limit" + c.ClientIP()
} }
func DefaultCheckOrMark(key string, expire int, limit int) bool {
func DefaultCheckOrMark(key string, expire int, limit int) (err error) {
// 判断是否开启redis // 判断是否开启redis
if global.GVA_REDIS == nil { if global.GVA_REDIS == nil {
return true
return err
} }
if err := SetLimitWithTime(key, limit, time.Duration(expire)*time.Second); err != nil {
if err = SetLimitWithTime(key, limit, time.Duration(expire)*time.Second); err != nil {
global.GVA_LOG.Error("limit", zap.Error(err)) global.GVA_LOG.Error("limit", zap.Error(err))
return false
} }
return true
return err
} }
// IPLimit ip限制
func IPLimit() gin.HandlerFunc {
return func(c *gin.Context) {
key := "RequestClientIPLimit===" + c.ClientIP()
limit := global.GVA_CONFIG.System.LimitCountIP
second := global.GVA_CONFIG.System.LimitTimeIP
expiration := time.Duration(second) * time.Second
if err := SetLimitWithTime(key, limit, expiration); err != nil {
response.FailWithMessage(err.Error(), c)
c.Abort()
return
}
// 继续往下处理
c.Next()
}
func DefaultLimit() gin.HandlerFunc {
return LimitConfig{
GenerationKey: DefaultGenerationKey,
CheckOrMark: DefaultCheckOrMark,
Expire: global.GVA_CONFIG.System.LimitTimeIP,
Limit: global.GVA_CONFIG.System.LimitCountIP,
}.LimitWithTime()
} }
// 设置访问次数
// SetLimitWithTime 设置访问次数
func SetLimitWithTime(key string, limit int, expiration time.Duration) error { func SetLimitWithTime(key string, limit int, expiration time.Duration) error {
count, err := global.GVA_REDIS.Exists(context.Background(), key).Result() count, err := global.GVA_REDIS.Exists(context.Background(), key).Result()
if err != nil { if err != nil {

Loading…
Cancel
Save