From c6501c2d0a505c22c0896735f96587df3b662228 Mon Sep 17 00:00:00 2001
From: songzhibin97 <718428482@qq.com>
Date: Tue, 12 Oct 2021 20:37:51 +0800
Subject: [PATCH] =?UTF-8?q?feat:=E4=BC=98=E5=8C=96=E5=88=9D=E5=A7=8B?=
 =?UTF-8?q?=E5=8C=96limit=E5=AF=B9=E8=B1=A1?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 server/middleware/limit_ip.go | 46 +++++++++++++++--------------------
 1 file changed, 19 insertions(+), 27 deletions(-)

diff --git a/server/middleware/limit_ip.go b/server/middleware/limit_ip.go
index 2bbdcb4f..f98037f3 100644
--- a/server/middleware/limit_ip.go
+++ b/server/middleware/limit_ip.go
@@ -17,21 +17,21 @@ type LimitConfig struct {
 	// GenerationKey 根据业务生成key 下面CheckOrMark查询生成
 	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 int
 	// Limit 周期时间
 	Limit int
 }
 
-func (l *LimitConfig) LimitWithTime() gin.HandlerFunc {
+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": "操作频繁,请稍后再试"})
+		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()
 			return
+		} else {
+			c.Next()
 		}
 	}
 }
@@ -41,37 +41,29 @@ func DefaultGenerationKey(c *gin.Context) string {
 	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
 	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))
-		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 {
 	count, err := global.GVA_REDIS.Exists(context.Background(), key).Result()
 	if err != nil {