🎨 添加中间件 && 完善预设注入功能 && 新增流式传输

server/middleware/limit_ip.go

@@ -0,0 +1,66 @@
+package middleware
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"net/http"
+	"time"
+
+	"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
+}
+
+func (l LimitConfig) LimitKey(c *gin.Context) string {
+	return "GVA_Limit" + c.ClientIP()
+}
+
+func (l LimitConfig) GetLimit(c *gin.Context) int {
+	return l.Limit
+}
+
+func (l LimitConfig) Reached(c *gin.Context) response.Response {
+	return response.Response{Code: response.ERROR, Data: nil, Msg: "操作过于频繁，请稍后再试"}
+}
+
+// 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))
+			}
+		}
+		c.Next()
+	}
+}