🎨 优化项目结构 && 完善ai配置

server/middleware/limit_ip.go

@@ -3,64 +3,90 @@ package middleware
 import (
 	"context"
 	"errors"
-	"fmt"
 	"net/http"
 	"time"
 
+	"go.uber.org/zap"
+
 	"git.echol.cn/loser/ai_proxy/server/global"
 	"git.echol.cn/loser/ai_proxy/server/model/common/response"
 	"github.com/gin-gonic/gin"
-	"go.uber.org/zap"
 )
 
 type LimitConfig struct {
-	GenerationDuration time.Duration
-	Limit              int
+	// GenerationKey 根据业务生成key 下面CheckOrMark查询生成
+	GenerationKey func(c *gin.Context) string
+	// 检查函数,用户可修改具体逻辑,更加灵活
+	CheckOrMark func(key string, expire int, limit int) error
+	// Expire key 过期时间
+	Expire int
+	// Limit 周期时间
+	Limit int
 }
 
-func (l LimitConfig) LimitKey(c *gin.Context) string {
+func (l LimitConfig) LimitWithTime() gin.HandlerFunc {
+	return func(c *gin.Context) {
+		if err := l.CheckOrMark(l.GenerationKey(c), l.Expire, l.Limit); err != nil {
+			c.JSON(http.StatusOK, gin.H{"code": response.ERROR, "msg": err.Error()})
+			c.Abort()
+			return
+		} else {
+			c.Next()
+		}
+	}
+}
+
+// DefaultGenerationKey 默认生成key
+func DefaultGenerationKey(c *gin.Context) string {
 	return "GVA_Limit" + c.ClientIP()
 }
 
-func (l LimitConfig) GetLimit(c *gin.Context) int {
-	return l.Limit
+func DefaultCheckOrMark(key string, expire int, limit int) (err error) {
+	// 判断是否开启redis
+	if global.GVA_REDIS == nil {
+		return err
+	}
+	if err = SetLimitWithTime(key, limit, time.Duration(expire)*time.Second); err != nil {
+		global.GVA_LOG.Error("limit", zap.Error(err))
+	}
+	return err
 }
 
-func (l LimitConfig) Reached(c *gin.Context) response.Response {
-	return response.Response{Code: response.ERROR, Data: nil, Msg: "操作过于频繁，请稍后再试"}
+func DefaultLimit() gin.HandlerFunc {
+	return LimitConfig{
+		GenerationKey: DefaultGenerationKey,
+		CheckOrMark:   DefaultCheckOrMark,
+		Expire:        global.GVA_CONFIG.System.LimitTimeIP,
+		Limit:         global.GVA_CONFIG.System.LimitCountIP,
+	}.LimitWithTime()
 }
 
-// IPLimit IP 限流中间件
-func IPLimit() gin.HandlerFunc {
-	return func(c *gin.Context) {
-		if global.GVA_CONFIG.System.UseRedis {
-			key := "GVA_Limit" + c.ClientIP()
-			limit := global.GVA_CONFIG.System.IpLimitCount
-			limitTime := global.GVA_CONFIG.System.IpLimitTime
-
-			ctx := context.Background()
-			count, err := global.GVA_REDIS.Get(ctx, key).Int()
-			if err != nil && !errors.Is(err, context.DeadlineExceeded) {
-				global.GVA_LOG.Error("get redis key error", zap.Error(err))
-			}
-
-			if count >= limit {
-				c.JSON(http.StatusOK, gin.H{
-					"code": response.ERROR,
-					"msg":  fmt.Sprintf("操作过于频繁，请在 %d 秒后再试", limitTime),
-				})
-				c.Abort()
-				return
-			}
-
-			pipe := global.GVA_REDIS.Pipeline()
-			pipe.Incr(ctx, key)
-			pipe.Expire(ctx, key, time.Second*time.Duration(limitTime))
-			_, err = pipe.Exec(ctx)
-			if err != nil {
-				global.GVA_LOG.Error("redis pipeline error", zap.Error(err))
+// SetLimitWithTime 设置访问次数
+func SetLimitWithTime(key string, limit int, expiration time.Duration) error {
+	count, err := global.GVA_REDIS.Exists(context.Background(), key).Result()
+	if err != nil {
+		return err
+	}
+	if count == 0 {
+		pipe := global.GVA_REDIS.TxPipeline()
+		pipe.Incr(context.Background(), key)
+		pipe.Expire(context.Background(), key, expiration)
+		_, err = pipe.Exec(context.Background())
+		return err
+	} else {
+		// 次数
+		if times, err := global.GVA_REDIS.Get(context.Background(), key).Int(); err != nil {
+			return err
+		} else {
+			if times >= limit {
+				if t, err := global.GVA_REDIS.PTTL(context.Background(), key).Result(); err != nil {
+					return errors.New("请求太过频繁，请稍后再试")
+				} else {
+					return errors.New("请求太过频繁, 请 " + t.String() + " 秒后尝试")
+				}
+			} else {
+				return global.GVA_REDIS.Incr(context.Background(), key).Err()
 			}
 		}
-		c.Next()
 	}
 }