From 714d8239d1781012b6a30a4f1f1379f7aeb5dfde Mon Sep 17 00:00:00 2001 From: fengsilin Date: Thu, 14 May 2026 11:07:19 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20Relay=20=E8=AF=B7=E6=B1=82=E6=8A=93?= =?UTF-8?q?=E5=8C=85=E6=97=A5=E5=BF=97=20+=20=E7=94=A8=E6=88=B7=E6=A8=A1?= =?UTF-8?q?=E5=9E=8B=20RPM=20=E9=99=90=E6=B5=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 两个独立功能: 1. Relay Capture:对开启 capture_relay 的用户记录完整请求/响应到本地文件 - 单 writer goroutine + 128 槽 channel 架构,无 OOM 风险 - 支持流式 SSE 和非流式请求 - 30s 超时保护 2. User-Model RPM Rate Limiting:按 用户+模型 维度的 RPM 限流 - Redis 令牌桶 + 内存滑动窗口双路径,Redis 故障自动降级 - 请求失败时退还令牌/配额(Redis Refund Lua 脚本 + 内存 Refund) - Redis pipeline 错误处理:失败时清理缓存 key - 前端用户编辑弹窗新增模型限流配置面板 Co-Authored-By: Claude Opus 4.7 --- common/limiter/limiter.go | 28 +- common/limiter/lua/rate_refund.lua | 18 + common/rate-limit.go | 12 + controller/user.go | 2 + controller/user_rate_limit.go | 74 ++++ controller/user_rate_limit_test.go | 152 +++++++ main.go | 5 + middleware/relay_capture.go | 177 ++++++++ middleware/relay_capture_integration_test.go | 402 ++++++++++++++++++ middleware/relay_capture_test.go | 194 +++++++++ middleware/user-model-rate-limit.go | 199 +++++++++ middleware/user_model_rate_limit_test.go | 216 ++++++++++ model/main.go | 1 + model/user.go | 13 +- model/user_capture_test.go | 43 ++ model/user_model_rate_limit.go | 54 +++ model/user_model_rate_limit_test.go | 110 +++++ router/api-router.go | 4 + router/relay-router.go | 4 + .../table/users/modals/EditUserModal.jsx | 18 + .../modals/UserModelRateLimitSection.jsx | 192 +++++++++ web/src/helpers/api.js | 6 + 22 files changed, 1915 insertions(+), 9 deletions(-) create mode 100644 common/limiter/lua/rate_refund.lua create mode 100644 controller/user_rate_limit.go create mode 100644 controller/user_rate_limit_test.go create mode 100644 middleware/relay_capture.go create mode 100644 middleware/relay_capture_integration_test.go create mode 100644 middleware/relay_capture_test.go create mode 100644 middleware/user-model-rate-limit.go create mode 100644 middleware/user_model_rate_limit_test.go create mode 100644 model/user_capture_test.go create mode 100644 model/user_model_rate_limit.go create mode 100644 model/user_model_rate_limit_test.go create mode 100644 web/src/components/table/users/modals/UserModelRateLimitSection.jsx diff --git a/common/limiter/limiter.go b/common/limiter/limiter.go index 6be61bc..214158c 100644 --- a/common/limiter/limiter.go +++ b/common/limiter/limiter.go @@ -13,9 +13,13 @@ import ( //go:embed lua/rate_limit.lua var rateLimitScript string +//go:embed lua/rate_refund.lua +var rateRefundScript string + type RedisLimiter struct { - client *redis.Client - limitScriptSHA string + client *redis.Client + limitScriptSHA string + refundScriptSHA string } var ( @@ -30,9 +34,14 @@ func New(ctx context.Context, r *redis.Client) *RedisLimiter { if err != nil { common.SysLog(fmt.Sprintf("Failed to load rate limit script: %v", err)) } + refundSHA, err := r.ScriptLoad(ctx, rateRefundScript).Result() + if err != nil { + common.SysLog(fmt.Sprintf("Failed to load rate refund script: %v", err)) + } instance = &RedisLimiter{ - client: r, - limitScriptSHA: limitSHA, + client: r, + limitScriptSHA: limitSHA, + refundScriptSHA: refundSHA, } }) @@ -68,6 +77,17 @@ func (rl *RedisLimiter) Allow(ctx context.Context, key string, opts ...Option) ( return result == 1, nil } +func (rl *RedisLimiter) Refund(ctx context.Context, key string, requested, capacity int64) error { + _, err := rl.client.EvalSha( + ctx, + rl.refundScriptSHA, + []string{key}, + requested, + capacity, + ).Int() + return err +} + // Config 配置选项模式 type Config struct { Capacity int64 diff --git a/common/limiter/lua/rate_refund.lua b/common/limiter/lua/rate_refund.lua new file mode 100644 index 0000000..95c0ba9 --- /dev/null +++ b/common/limiter/lua/rate_refund.lua @@ -0,0 +1,18 @@ +-- 令牌桶退款脚本 +-- KEYS[1]: 限流器唯一标识 +-- ARGV[1]: 退还的令牌数 +-- ARGV[2]: 桶容量(上限) +local key = KEYS[1] +local refund = tonumber(ARGV[1]) +local capacity = tonumber(ARGV[2]) + +local bucket = redis.call('HMGET', key, 'tokens', 'last_time') +local tokens = tonumber(bucket[1]) + +if not tokens then + return 0 +end + +tokens = math.min(capacity, tokens + refund) +redis.call('HSET', key, 'tokens', tokens) +return 1 diff --git a/common/rate-limit.go b/common/rate-limit.go index 301c101..dd05528 100644 --- a/common/rate-limit.go +++ b/common/rate-limit.go @@ -68,3 +68,15 @@ func (l *InMemoryRateLimiter) Request(key string, maxRequestNum int, duration in } return true } + +// Refund 移除 key 最近一次请求记录(用于请求失败时退还配额) +func (l *InMemoryRateLimiter) Refund(key string) bool { + l.mutex.Lock() + defer l.mutex.Unlock() + queue, ok := l.store[key] + if !ok || len(*queue) == 0 { + return false + } + *queue = (*queue)[:len(*queue)-1] + return true +} diff --git a/controller/user.go b/controller/user.go index 9bfb254..98ca026 100644 --- a/controller/user.go +++ b/controller/user.go @@ -14,6 +14,7 @@ import ( "github.com/QuantumNous/new-api/dto" "github.com/QuantumNous/new-api/i18n" "github.com/QuantumNous/new-api/logger" + "github.com/QuantumNous/new-api/middleware" "github.com/QuantumNous/new-api/model" "github.com/QuantumNous/new-api/service" "github.com/QuantumNous/new-api/service/region_sync" @@ -578,6 +579,7 @@ func UpdateUser(c *gin.Context) { common.ApiError(c, err) return } + middleware.SetCaptureEnabled(int64(updatedUser.Id), updatedUser.CaptureRelay) if originUser.Quota != updatedUser.Quota { model.RecordLog(originUser.Id, model.LogTypeManage, fmt.Sprintf("管理员将用户额度从 %s修改为 %s", logger.LogQuota(originUser.Quota), logger.LogQuota(updatedUser.Quota))) } diff --git a/controller/user_rate_limit.go b/controller/user_rate_limit.go new file mode 100644 index 0000000..980c996 --- /dev/null +++ b/controller/user_rate_limit.go @@ -0,0 +1,74 @@ +package controller + +import ( + "strconv" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + "github.com/gin-gonic/gin" +) + +// GetUserRateLimits GET /api/user/:id/rate-limits +// 查询指定用户的所有模型 RPM 限制配置(管理员权限) +func GetUserRateLimits(c *gin.Context) { + userId, err := strconv.Atoi(c.Param("id")) + if err != nil { + common.ApiErrorMsg(c, "invalid user id") + return + } + + list, err := model.GetUserModelRateLimits(userId) + if err != nil { + common.ApiError(c, err) + return + } + + // 保证返回空列表而非 null + if list == nil { + list = []model.UserModelRateLimit{} + } + + c.JSON(200, gin.H{ + "success": true, + "message": "", + "data": list, + }) +} + +// SetUserRateLimits PUT /api/user/:id/rate-limits +// 覆盖式写入指定用户的所有模型 RPM 限制配置(管理员权限) +func SetUserRateLimits(c *gin.Context) { + userId, err := strconv.Atoi(c.Param("id")) + if err != nil { + common.ApiErrorMsg(c, "invalid user id") + return + } + + var items []model.UserModelRateLimit + if err := c.ShouldBindJSON(&items); err != nil { + common.ApiErrorMsg(c, err.Error()) + return + } + + for _, item := range items { + if item.Rpm < 0 { + common.ApiErrorMsg(c, "rpm must be >= 0") + return + } + } + + if err := model.SetUserModelRateLimits(userId, items); err != nil { + common.ApiError(c, err) + return + } + + // 删除 Redis 缓存 + if common.RedisEnabled { + _ = common.RedisDel("user_model_rate_limit:" + strconv.Itoa(userId)) + } + + c.JSON(200, gin.H{ + "success": true, + "message": "", + }) +} diff --git a/controller/user_rate_limit_test.go b/controller/user_rate_limit_test.go new file mode 100644 index 0000000..56b8048 --- /dev/null +++ b/controller/user_rate_limit_test.go @@ -0,0 +1,152 @@ +package controller + +import ( + "bytes" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + "github.com/glebarez/sqlite" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func setupUserRateLimitControllerDB(t *testing.T) { + t.Helper() + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + sqlDB, _ := db.DB() + sqlDB.SetMaxOpenConns(1) + + origDB := model.DB + model.DB = db + common.RedisEnabled = false + + require.NoError(t, db.AutoMigrate(&model.UserModelRateLimit{})) + + t.Cleanup(func() { + model.DB = origDB + sqlDB.Close() + }) +} + +func setupUserRateLimitRouter() *gin.Engine { + gin.SetMode(gin.TestMode) + r := gin.New() + g := r.Group("/api/user") + { + g.GET("/:id/rate-limits", GetUserRateLimits) + g.PUT("/:id/rate-limits", SetUserRateLimits) + } + return r +} + +func TestGetUserRateLimits_Empty(t *testing.T) { + setupUserRateLimitControllerDB(t) + router := setupUserRateLimitRouter() + + w := httptest.NewRecorder() + req, _ := http.NewRequest(http.MethodGet, "/api/user/1/rate-limits", nil) + router.ServeHTTP(w, req) + + assert.Equal(t, http.StatusOK, w.Code) + var resp map[string]interface{} + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) + assert.True(t, resp["success"].(bool)) + // data 应为空列表(非 nil) + data, ok := resp["data"] + require.True(t, ok, "response should have data key") + assert.Empty(t, data) +} + +func TestSetUserRateLimits_OK(t *testing.T) { + setupUserRateLimitControllerDB(t) + router := setupUserRateLimitRouter() + + items := []model.UserModelRateLimit{ + {Model: "gpt-4", Rpm: 60}, + } + body, _ := json.Marshal(items) + + w := httptest.NewRecorder() + req, _ := http.NewRequest(http.MethodPut, "/api/user/1/rate-limits", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + router.ServeHTTP(w, req) + + assert.Equal(t, http.StatusOK, w.Code) + var resp map[string]interface{} + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) + assert.True(t, resp["success"].(bool)) + + // 验证数据确实写入 + list, err := model.GetUserModelRateLimits(1) + require.NoError(t, err) + require.Len(t, list, 1) + assert.Equal(t, "gpt-4", list[0].Model) + assert.Equal(t, 60, list[0].Rpm) +} + +func TestSetUserRateLimits_NegativeRpm(t *testing.T) { + setupUserRateLimitControllerDB(t) + router := setupUserRateLimitRouter() + + items := []model.UserModelRateLimit{ + {Model: "gpt-4", Rpm: -1}, + } + body, _ := json.Marshal(items) + + w := httptest.NewRecorder() + req, _ := http.NewRequest(http.MethodPut, "/api/user/1/rate-limits", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + router.ServeHTTP(w, req) + + assert.Equal(t, http.StatusOK, w.Code) + var resp map[string]interface{} + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) + assert.False(t, resp["success"].(bool)) +} + +func TestGetUserRateLimits_InvalidId(t *testing.T) { + setupUserRateLimitControllerDB(t) + router := setupUserRateLimitRouter() + + w := httptest.NewRecorder() + req, _ := http.NewRequest(http.MethodGet, "/api/user/notanid/rate-limits", nil) + router.ServeHTTP(w, req) + + assert.Equal(t, http.StatusOK, w.Code) + var resp map[string]interface{} + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) + assert.False(t, resp["success"].(bool)) +} + +func TestSetUserRateLimits_EmptySlice(t *testing.T) { + setupUserRateLimitControllerDB(t) + router := setupUserRateLimitRouter() + + // 先写一条 + _ = model.SetUserModelRateLimits(2, []model.UserModelRateLimit{ + {UserId: 2, Model: "claude-3-opus", Rpm: 10}, + }) + + // 用空 slice 清空 + body, _ := json.Marshal([]model.UserModelRateLimit{}) + w := httptest.NewRecorder() + req, _ := http.NewRequest(http.MethodPut, "/api/user/2/rate-limits", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + router.ServeHTTP(w, req) + + assert.Equal(t, http.StatusOK, w.Code) + var resp map[string]interface{} + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) + assert.True(t, resp["success"].(bool)) + + list, err := model.GetUserModelRateLimits(2) + require.NoError(t, err) + assert.Empty(t, list) +} diff --git a/main.go b/main.go index 3b5ef0e..615c68a 100644 --- a/main.go +++ b/main.go @@ -293,6 +293,11 @@ func InitResources() error { // 初始化模型 model.GetPricing() + // 加载抓包用户缓存 + if err := middleware.LoadCaptureEnabledUsers(); err != nil { + common.SysError("failed to load capture enabled users: " + err.Error()) + } + // Initialize SQL Database err = model.InitLogDB() if err != nil { diff --git a/middleware/relay_capture.go b/middleware/relay_capture.go new file mode 100644 index 0000000..1f07090 --- /dev/null +++ b/middleware/relay_capture.go @@ -0,0 +1,177 @@ +package middleware + +import ( + "fmt" + "os" + "path/filepath" + "sync" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + "github.com/gin-gonic/gin" +) + +var captureEnabledUsers sync.Map // key: int64, value: struct{} + +func isCaptureEnabled(userID int64) bool { + _, ok := captureEnabledUsers.Load(userID) + return ok +} + +func SetCaptureEnabled(userID int64, enabled bool) { + if enabled { + captureEnabledUsers.Store(userID, struct{}{}) + } else { + captureEnabledUsers.Delete(userID) + } +} + +// LoadCaptureEnabledUsers 从数据库加载 capture_relay=true 的用户到内存缓存 +func LoadCaptureEnabledUsers() error { + var ids []int64 + if err := model.DB.Model(&model.User{}). + Where("capture_relay = ?", true). + Pluck("id", &ids).Error; err != nil { + return err + } + for _, id := range ids { + captureEnabledUsers.Store(id, struct{}{}) + } + return nil +} + +// GetCaptureEnabledUsersCount 仅用于测试 +func GetCaptureEnabledUsersCount() int { + count := 0 + captureEnabledUsers.Range(func(key, value interface{}) bool { + count++ + return true + }) + return count +} + +// --- captureResponseWriter --- + +type captureResponseWriter struct { + gin.ResponseWriter + ch chan<- []byte +} + +func (cw *captureResponseWriter) Write(b []byte) (int, error) { + n, err := cw.ResponseWriter.Write(b) + if n > 0 { + buf := make([]byte, n) + copy(buf, b[:n]) + cw.ch <- buf + } + return n, err +} + +func (cw *captureResponseWriter) WriteString(s string) (int, error) { + n, err := cw.ResponseWriter.WriteString(s) + if n > 0 { + buf := make([]byte, n) + copy(buf, s[:n]) + cw.ch <- buf + } + return n, err +} + +// --- startCaptureWriter --- + +// startCaptureWriter 启动单个 writer goroutine,从 chIn 顺序写入文件 +func startCaptureWriter(f *os.File, requestID string, start time.Time) (in chan<- []byte, done <-chan struct{}) { + chIn := make(chan []byte, 128) + doneCh := make(chan struct{}) + + go func() { + defer close(doneCh) + defer f.Close() + for chunk := range chIn { + f.Write(chunk) + } + fmt.Fprintf(f, "\n=== END duration_ms=%d ===\n", time.Since(start).Milliseconds()) + }() + + return chIn, doneCh +} + +// --- 辅助函数 --- + +func createCaptureFile(requestID string) (*os.File, error) { + dir := filepath.Join("./data", "relay-capture", time.Now().Format("2006-01-02")) + if err := os.MkdirAll(dir, 0755); err != nil { + return nil, err + } + return os.OpenFile( + filepath.Join(dir, requestID+".log"), + os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644, + ) +} + +func writeRequestBlock(c *gin.Context, f *os.File, userID int64) { + fmt.Fprintf(f, "=== REQUEST %s ===\n", time.Now().UTC().Format(time.RFC3339Nano)) + fmt.Fprintf(f, "%s %s\n", c.Request.Method, c.Request.URL.Path) + fmt.Fprintf(f, "user_id: %d\n", userID) + + for key, values := range c.Request.Header { + for _, v := range values { + fmt.Fprintf(f, "%s: %s\n", key, v) + } + } + fmt.Fprintf(f, "\n") + + // 读取请求体 + if c.Request.Body != nil { + if storage, err := common.GetBodyStorage(c); err == nil { + if body, err := storage.Bytes(); err == nil && len(body) > 0 { + f.Write(body) + } + } + } + fmt.Fprintf(f, "\n") +} + +// RelayCaptureMiddleware 抓包中间件 +// 对开启 capture_relay 的用户,将请求体和响应体写入本地文件 +func RelayCaptureMiddleware() gin.HandlerFunc { + return func(c *gin.Context) { + userID := int64(c.GetInt("id")) + if !isCaptureEnabled(userID) { + c.Next() + return + } + + start := time.Now() + requestID := c.GetString(common.RequestIdKey) + + f, err := createCaptureFile(requestID) + if err != nil { + c.Next() + return + } + + // 同步写 REQUEST 块 + writeRequestBlock(c, f, userID) + fmt.Fprintf(f, "\n=== RESPONSE ===\n") + + // 启动 writer goroutine + chIn, doneCh := startCaptureWriter(f, requestID, start) + + c.Writer = &captureResponseWriter{ + ResponseWriter: c.Writer, + ch: chIn, + } + + c.Next() + + // 关闭 chIn,等待 writer 完成(带超时保护) + close(chIn) + select { + case <-doneCh: + case <-time.After(30 * time.Second): + common.SysError("relay capture: timeout waiting for writer goroutine, request_id=" + requestID) + } + } +} diff --git a/middleware/relay_capture_integration_test.go b/middleware/relay_capture_integration_test.go new file mode 100644 index 0000000..52dadb7 --- /dev/null +++ b/middleware/relay_capture_integration_test.go @@ -0,0 +1,402 @@ +package middleware + +import ( + "bytes" + "fmt" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "sync/atomic" + "testing" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func setupCaptureTestRouter() *gin.Engine { + gin.SetMode(gin.TestMode) + return gin.New() +} + +// captureTestCaptureDir 集成测试共享的捕获文件目录 +const captureTestDir = "./data/relay-capture" + +// --- 非流式请求测试 --- + +func TestRelayCapture_NonStreaming_EnabledUser(t *testing.T) { + clearCaptureCache() + + SetCaptureEnabled(42, true) + defer SetCaptureEnabled(42, false) + + router := setupCaptureTestRouter() + router.Use(func(c *gin.Context) { + c.Set("id", 42) // must be int, not int64 — matches production TokenAuth behavior + c.Set(common.RequestIdKey, "test-non-stream-"+time.Now().Format("150405")) + c.Next() + }) + router.Use(RelayCaptureMiddleware()) + router.POST("/v1/chat/completions", func(c *gin.Context) { + c.JSON(200, gin.H{ + "id": "chatcmpl-test", + "object": "chat.completion", + "choices": []gin.H{{"message": gin.H{"role": "assistant", "content": "Hello!"}}}, + }) + }) + + body := `{"model":"gpt-4","messages":[{"role":"user","content":"hi"}]}` + req, _ := http.NewRequest("POST", "/v1/chat/completions", bytes.NewReader([]byte(body))) + req.Header.Set("Authorization", "Bearer sk-test") + req.Header.Set("Content-Type", "application/json") + + rec := httptest.NewRecorder() + router.ServeHTTP(rec, req) + + assert.Equal(t, 200, rec.Code) + + today := time.Now().Format("2006-01-02") + dateDir := filepath.Join(captureTestDir, today) + files, err := os.ReadDir(dateDir) + require.NoError(t, err) + require.True(t, len(files) > 0, "should have capture log files") + + logPath := filepath.Join(dateDir, files[len(files)-1].Name()) + content, err := os.ReadFile(logPath) + require.NoError(t, err) + contentStr := string(content) + + t.Logf("Capture file content:\n%s", contentStr) + + assert.Contains(t, contentStr, "=== REQUEST") + assert.Contains(t, contentStr, "POST /v1/chat/completions") + assert.Contains(t, contentStr, "user_id: 42") + assert.Contains(t, contentStr, "Authorization: Bearer sk-test") + assert.Contains(t, contentStr, "=== RESPONSE") + assert.Contains(t, contentStr, "chatcmpl-test") + assert.Contains(t, contentStr, "=== END") + assert.Contains(t, contentStr, "duration_ms=") +} + +// --- 流式 SSE 请求测试 --- + +func TestRelayCapture_Streaming_EnabledUser(t *testing.T) { + clearCaptureCache() + t.Cleanup(func() { os.RemoveAll(captureTestDir) }) + + SetCaptureEnabled(42, true) + defer SetCaptureEnabled(42, false) + + router := setupCaptureTestRouter() + router.Use(func(c *gin.Context) { + c.Set("id", 42) + c.Set(common.RequestIdKey, "test-stream-"+time.Now().Format("150405")) + c.Next() + }) + router.Use(RelayCaptureMiddleware()) + router.POST("/v1/chat/completions", func(c *gin.Context) { + c.Header("Content-Type", "text/event-stream") + c.Header("Cache-Control", "no-cache") + c.Header("Connection", "keep-alive") + + flusher := c.Writer.(http.Flusher) + + chunks := []string{ + "data: {\"id\":\"1\",\"choices\":[{\"delta\":{\"content\":\"Hello \"}}]}\n\n", + "data: {\"id\":\"1\",\"choices\":[{\"delta\":{\"content\":\"world!\"}}]}\n\n", + "data: [DONE]\n\n", + } + + for _, chunk := range chunks { + c.Writer.Write([]byte(chunk)) + flusher.Flush() + time.Sleep(5 * time.Millisecond) + } + }) + + body := `{"model":"gpt-4","messages":[{"role":"user","content":"hi"}],"stream":true}` + req, _ := http.NewRequest("POST", "/v1/chat/completions", bytes.NewReader([]byte(body))) + req.Header.Set("Authorization", "Bearer sk-test") + req.Header.Set("Content-Type", "application/json") + + rec := httptest.NewRecorder() + router.ServeHTTP(rec, req) + + assert.Equal(t, 200, rec.Code) + + today := time.Now().Format("2006-01-02") + dateDir := filepath.Join(captureTestDir, today) + files, err := os.ReadDir(dateDir) + require.NoError(t, err) + require.True(t, len(files) > 0, "should have capture log files") + + logPath := filepath.Join(dateDir, files[len(files)-1].Name()) + content, err := os.ReadFile(logPath) + require.NoError(t, err) + contentStr := string(content) + + t.Logf("Stream capture file content:\n%s", contentStr) + + assert.Contains(t, contentStr, "=== REQUEST") + assert.Contains(t, contentStr, "POST /v1/chat/completions") + assert.Contains(t, contentStr, "user_id: 42") + assert.Contains(t, contentStr, "=== RESPONSE") + assert.Contains(t, contentStr, "Hello ") + assert.Contains(t, contentStr, "world!") + assert.Contains(t, contentStr, "[DONE]") + assert.Contains(t, contentStr, "=== END") + assert.Contains(t, contentStr, "duration_ms=") +} + +// --- 未开启抓包用户测试 --- + +func TestRelayCapture_DisabledUser(t *testing.T) { + clearCaptureCache() + + router := setupCaptureTestRouter() + router.Use(func(c *gin.Context) { + c.Set("id", 99) + c.Set(common.RequestIdKey, "test-disabled-"+time.Now().Format("150405")) + c.Next() + }) + router.Use(RelayCaptureMiddleware()) + router.POST("/v1/chat/completions", func(c *gin.Context) { + c.JSON(200, gin.H{"result": "ok"}) + }) + + req, _ := http.NewRequest("POST", "/v1/chat/completions", bytes.NewReader([]byte(`{}`))) + req.Header.Set("Authorization", "Bearer sk-test") + + rec := httptest.NewRecorder() + router.ServeHTTP(rec, req) + + assert.Equal(t, 200, rec.Code) + + today := time.Now().Format("2006-01-02") + dateDir := filepath.Join(captureTestDir, today) + if _, err := os.Stat(dateDir); err == nil { + files, _ := os.ReadDir(dateDir) + for _, f := range files { + assert.NotContains(t, f.Name(), "test-disabled", "disabled user should not generate capture files") + } + } +} + +// --- 大量数据流测试 --- + +func TestRelayCapture_LargeStream(t *testing.T) { + clearCaptureCache() + t.Cleanup(func() { os.RemoveAll(captureTestDir) }) + + SetCaptureEnabled(42, true) + defer SetCaptureEnabled(42, false) + + router := setupCaptureTestRouter() + router.Use(func(c *gin.Context) { + c.Set("id", 42) + c.Set(common.RequestIdKey, "test-large-"+time.Now().Format("150405")) + c.Next() + }) + router.Use(RelayCaptureMiddleware()) + router.POST("/v1/chat/completions", func(c *gin.Context) { + c.Header("Content-Type", "text/event-stream") + flusher := c.Writer.(http.Flusher) + + for i := 0; i < 50; i++ { + chunk := fmt.Sprintf("data: {\"id\":\"%d\",\"choices\":[{\"delta\":{\"content\":\"chunk%d \"}}]}\n\n", i, i) + c.Writer.Write([]byte(chunk)) + flusher.Flush() + } + c.Writer.Write([]byte("data: [DONE]\n\n")) + flusher.Flush() + }) + + req, _ := http.NewRequest("POST", "/v1/chat/completions", bytes.NewReader([]byte(`{"model":"gpt-4","stream":true}`))) + req.Header.Set("Authorization", "Bearer sk-test") + + rec := httptest.NewRecorder() + router.ServeHTTP(rec, req) + + assert.Equal(t, 200, rec.Code) + + today := time.Now().Format("2006-01-02") + dateDir := filepath.Join(captureTestDir, today) + files, err := os.ReadDir(dateDir) + require.NoError(t, err) + require.True(t, len(files) > 0) + + logPath := filepath.Join(dateDir, files[len(files)-1].Name()) + content, err := os.ReadFile(logPath) + require.NoError(t, err) + contentStr := string(content) + + assert.Contains(t, contentStr, "chunk0") + assert.Contains(t, contentStr, "chunk49") + assert.Contains(t, contentStr, "[DONE]") +} + +// --- GET 请求测试 --- + +func TestRelayCapture_GetRequest(t *testing.T) { + clearCaptureCache() + t.Cleanup(func() { os.RemoveAll(captureTestDir) }) + + SetCaptureEnabled(42, true) + defer SetCaptureEnabled(42, false) + + router := setupCaptureTestRouter() + router.Use(func(c *gin.Context) { + c.Set("id", 42) + c.Set(common.RequestIdKey, "test-get-"+time.Now().Format("150405")) + c.Next() + }) + router.Use(RelayCaptureMiddleware()) + router.GET("/v1/models", func(c *gin.Context) { + c.JSON(200, gin.H{ + "data": []gin.H{{"id": "gpt-4", "object": "model"}}, + }) + }) + + req, _ := http.NewRequest("GET", "/v1/models", nil) + req.Header.Set("Authorization", "Bearer sk-test") + + rec := httptest.NewRecorder() + router.ServeHTTP(rec, req) + + assert.Equal(t, 200, rec.Code) + + today := time.Now().Format("2006-01-02") + dateDir := filepath.Join(captureTestDir, today) + files, err := os.ReadDir(dateDir) + require.NoError(t, err) + require.True(t, len(files) > 0) + + logPath := filepath.Join(dateDir, files[len(files)-1].Name()) + content, err := os.ReadFile(logPath) + require.NoError(t, err) + contentStr := string(content) + + assert.Contains(t, contentStr, "=== REQUEST") + assert.Contains(t, contentStr, "GET /v1/models") + assert.Contains(t, contentStr, "=== RESPONSE") + assert.Contains(t, contentStr, "gpt-4") + assert.Contains(t, contentStr, "=== END") +} + +// --- 零 ID 用户测试 --- + +func TestRelayCapture_ZeroUserID(t *testing.T) { + clearCaptureCache() + + router := setupCaptureTestRouter() + router.Use(func(c *gin.Context) { + c.Set(common.RequestIdKey, "test-zero-"+time.Now().Format("150405")) + c.Next() + }) + router.Use(RelayCaptureMiddleware()) + router.GET("/v1/models", func(c *gin.Context) { + c.JSON(200, gin.H{"data": []gin.H{}}) + }) + + req, _ := http.NewRequest("GET", "/v1/models", nil) + rec := httptest.NewRecorder() + router.ServeHTTP(rec, req) + + assert.Equal(t, 200, rec.Code) +} + +// --- 验证响应不被篡改 --- + +func TestRelayCapture_ResponseUnchanged(t *testing.T) { + clearCaptureCache() + t.Cleanup(func() { os.RemoveAll(captureTestDir) }) + + SetCaptureEnabled(42, true) + defer SetCaptureEnabled(42, false) + + router := setupCaptureTestRouter() + router.Use(func(c *gin.Context) { + c.Set("id", 42) + c.Set(common.RequestIdKey, "test-unchanged-"+time.Now().Format("150405")) + c.Next() + }) + router.Use(RelayCaptureMiddleware()) + router.POST("/v1/chat/completions", func(c *gin.Context) { + c.JSON(200, gin.H{ + "id": "chatcmpl-123", + "choices": []gin.H{{"message": gin.H{"content": "exact response"}}}, + }) + }) + + req, _ := http.NewRequest("POST", "/v1/chat/completions", bytes.NewReader([]byte(`{"model":"gpt-4"}`))) + req.Header.Set("Content-Type", "application/json") + + rec := httptest.NewRecorder() + router.ServeHTTP(rec, req) + + assert.Equal(t, 200, rec.Code) + assert.Contains(t, rec.Body.String(), "exact response") + assert.Contains(t, rec.Body.String(), "chatcmpl-123") +} + +// --- 并发安全测试 --- + +func TestRelayCapture_ConcurrentRequests(t *testing.T) { + clearCaptureCache() + t.Cleanup(func() { os.RemoveAll(captureTestDir) }) + + SetCaptureEnabled(42, true) + defer SetCaptureEnabled(42, false) + + router := setupCaptureTestRouter() + var reqCounter int64 = 0 + router.Use(func(c *gin.Context) { + c.Set("id", 42) + n := atomic.AddInt64(&reqCounter, 1) + c.Set(common.RequestIdKey, fmt.Sprintf("concurrent-%d-%d", n, time.Now().UnixNano())) + c.Next() + }) + router.Use(RelayCaptureMiddleware()) + router.POST("/v1/chat/completions", func(c *gin.Context) { + c.JSON(200, gin.H{"status": "ok"}) + }) + + done := make(chan bool, 10) + for i := 0; i < 10; i++ { + go func(idx int) { + req, _ := http.NewRequest("POST", "/v1/chat/completions", bytes.NewReader([]byte(fmt.Sprintf(`{"req":%d}`, idx)))) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + router.ServeHTTP(rec, req) + assert.Equal(t, 200, rec.Code) + done <- true + }(i) + } + + for i := 0; i < 10; i++ { + select { + case <-done: + case <-time.After(10 * time.Second): + t.Fatal("timeout waiting for concurrent request") + } + } + + time.Sleep(100 * time.Millisecond) + today := time.Now().Format("2006-01-02") + dateDir := filepath.Join(captureTestDir, today) + files, err := os.ReadDir(dateDir) + require.NoError(t, err) + assert.Equal(t, 10, len(files), "should have 10 capture log files") + + for _, f := range files { + content, err := os.ReadFile(filepath.Join(dateDir, f.Name())) + require.NoError(t, err) + contentStr := string(content) + assert.Contains(t, contentStr, "=== REQUEST") + assert.Contains(t, contentStr, "=== RESPONSE") + assert.Contains(t, contentStr, "=== END") + } +} diff --git a/middleware/relay_capture_test.go b/middleware/relay_capture_test.go new file mode 100644 index 0000000..8fad565 --- /dev/null +++ b/middleware/relay_capture_test.go @@ -0,0 +1,194 @@ +package middleware + +import ( + "fmt" + "io" + "net/http/httptest" + "os" + "path/filepath" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func clearCaptureCache() { + captureEnabledUsers.Range(func(key, value interface{}) bool { + captureEnabledUsers.Delete(key) + return true + }) +} + +func TestCaptureEnabled_NotFound(t *testing.T) { + clearCaptureCache() + assert.False(t, isCaptureEnabled(1)) +} + +func TestSetCaptureEnabled_Enable(t *testing.T) { + clearCaptureCache() + + SetCaptureEnabled(1, true) + assert.True(t, isCaptureEnabled(1)) + + SetCaptureEnabled(1, false) +} + +func TestSetCaptureEnabled_Disable(t *testing.T) { + clearCaptureCache() + + SetCaptureEnabled(1, true) + SetCaptureEnabled(1, false) + assert.False(t, isCaptureEnabled(1)) +} + +func TestSetCaptureEnabled_DifferentUser(t *testing.T) { + clearCaptureCache() + + SetCaptureEnabled(1, true) + assert.False(t, isCaptureEnabled(2)) + + SetCaptureEnabled(1, false) +} + +func TestGetCaptureEnabledUsersCount(t *testing.T) { + clearCaptureCache() + + assert.Equal(t, 0, GetCaptureEnabledUsersCount()) + + SetCaptureEnabled(1, true) + SetCaptureEnabled(2, true) + assert.Equal(t, 2, GetCaptureEnabledUsersCount()) + + SetCaptureEnabled(1, false) + assert.Equal(t, 1, GetCaptureEnabledUsersCount()) + + SetCaptureEnabled(2, false) +} + +// --- captureResponseWriter 测试 --- + +func TestCaptureResponseWriter_Write(t *testing.T) { + gin.SetMode(gin.TestMode) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + ch := make(chan []byte, 16) + + cw := &captureResponseWriter{ + ResponseWriter: c.Writer, + ch: ch, + } + + n, err := cw.Write([]byte("hello")) + assert.NoError(t, err) + assert.Equal(t, 5, n) + assert.Equal(t, "hello", rec.Body.String()) + + select { + case data := <-ch: + assert.Equal(t, []byte("hello"), data) + default: + t.Fatal("expected data in channel") + } +} + +func TestCaptureResponseWriter_WriteString(t *testing.T) { + gin.SetMode(gin.TestMode) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + ch := make(chan []byte, 16) + + cw := &captureResponseWriter{ + ResponseWriter: c.Writer, + ch: ch, + } + + n, err := cw.WriteString("test-string") + assert.NoError(t, err) + assert.Equal(t, len("test-string"), n) + assert.Equal(t, "test-string", rec.Body.String()) + + select { + case data := <-ch: + assert.Equal(t, []byte("test-string"), data) + default: + t.Fatal("expected data in channel") + } +} + +// --- startCaptureWriter 测试 --- + +func TestStartCaptureWriter_WritesAllChunks(t *testing.T) { + f, err := os.CreateTemp("", "test-capture-*.log") + require.NoError(t, err) + filePath := f.Name() + defer os.Remove(filePath) + + start := time.Now() + chIn, doneCh := startCaptureWriter(f, "test-req", start) + + // 发送 130 个 chunk(大于 channel buffer 128) + for i := 0; i < 130; i++ { + chIn <- []byte(fmt.Sprintf("data: %d\n", i)) + } + close(chIn) + + select { + case <-doneCh: + case <-time.After(5 * time.Second): + t.Fatal("timeout waiting for doneCh") + } + + content, err := os.ReadFile(filePath) + require.NoError(t, err) + contentStr := string(content) + assert.Contains(t, contentStr, "data: 0") + assert.Contains(t, contentStr, "data: 19") + assert.Contains(t, contentStr, "=== END") + assert.Contains(t, contentStr, "duration_ms=") +} + +// --- 辅助函数测试 --- + +func TestCreateCaptureFile(t *testing.T) { + f, err := createCaptureFile("test-req-123") + require.NoError(t, err) + defer os.RemoveAll(filepath.Dir(f.Name())) + + assert.Contains(t, f.Name(), "relay-capture") + assert.Contains(t, f.Name(), "test-req-123.log") + + dir := filepath.Dir(f.Name()) + info, err := os.Stat(dir) + require.NoError(t, err) + assert.True(t, info.IsDir()) + + f.Close() +} + +func TestWriteRequestBlock(t *testing.T) { + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + + c.Set("id", int64(42)) + c.Request = httptest.NewRequest("POST", "/v1/chat/completions", nil) + c.Request.Header.Set("Authorization", "Bearer sk-test") + c.Request.Header.Set("Content-Type", "application/json") + + f, err := os.CreateTemp("", "test-req-block-*.log") + require.NoError(t, err) + filePath := f.Name() + defer os.Remove(filePath) + + writeRequestBlock(c, f, 42) + + f.Seek(0, io.SeekStart) + content, err := io.ReadAll(f) + require.NoError(t, err) + contentStr := string(content) + assert.Contains(t, contentStr, "=== REQUEST") + assert.Contains(t, contentStr, "POST /v1/chat/completions") + assert.Contains(t, contentStr, "user_id: 42") +} diff --git a/middleware/user-model-rate-limit.go b/middleware/user-model-rate-limit.go new file mode 100644 index 0000000..eb23a93 --- /dev/null +++ b/middleware/user-model-rate-limit.go @@ -0,0 +1,199 @@ +package middleware + +import ( + "context" + "fmt" + "net/http" + "strconv" + "sync" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/common/limiter" + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/setting" + + "github.com/gin-gonic/gin" +) + +const ( + UserModelRateLimitMark = "UMRL" +) + +var userModelRateLimiter common.InMemoryRateLimiter + +// userModelRpmCache 非 Redis 模式下的进程内缓存,避免每次请求打 DB +var userModelRpmCache sync.Map // key: "userId:model" -> rpmCacheEntry + +type rpmCacheEntry struct { + rpm int + expiresAt int64 +} + +const rpmCacheTTLSeconds = 300 // 5 分钟,与 Redis 路径一致 + +// loadUserModelRpm 查询指定用户+模型的 RPM 限制。 +// Redis 可用时使用 Hash 缓存批量回填;Redis 不可用时直接查 DB。 +// 返回 0 表示不限制(无配置)。 +func loadUserModelRpm(userId int, modelName string) int { + if common.RedisEnabled { + return loadUserModelRpmFromRedis(userId, modelName) + } + return loadUserModelRpmFromMemory(userId, modelName) +} + +func loadUserModelRpmFromMemory(userId int, modelName string) int { + cacheKey := fmt.Sprintf("%d:%s", userId, modelName) + now := time.Now().Unix() + + if v, ok := userModelRpmCache.Load(cacheKey); ok { + entry := v.(rpmCacheEntry) + if now < entry.expiresAt { + return entry.rpm + } + userModelRpmCache.Delete(cacheKey) + } + + rpm, _ := model.GetUserModelRpm(userId, modelName) + userModelRpmCache.Store(cacheKey, rpmCacheEntry{rpm: rpm, expiresAt: now + rpmCacheTTLSeconds}) + return rpm +} + +func loadUserModelRpmFromRedis(userId int, modelName string) int { + ctx := context.Background() + rdb := common.RDB + hashKey := fmt.Sprintf("user_model_rate_limit:%d", userId) + + // 单次 HMGet 同时查目标模型和哨兵字段 + vals, err := rdb.HMGet(ctx, hashKey, modelName, "__loaded").Result() + if err == nil && len(vals) == 2 { + // vals[0] = 模型 RPM,vals[1] = 哨兵 + if vals[0] != nil { + if s, ok := vals[0].(string); ok { + if rpm, parseErr := strconv.Atoi(s); parseErr == nil { + return rpm + } + } + } + // 哨兵存在说明已加载过,该模型无配置 + if vals[1] != nil { + return 0 + } + } + + // 首次查询:从 DB 加载并批量回填(含哨兵字段,防止未配置用户每次打 DB) + limits, dbErr := model.GetUserModelRateLimits(userId) + if dbErr != nil { + return 0 + } + + pipe := rdb.Pipeline() + for _, l := range limits { + pipe.HSet(ctx, hashKey, l.Model, l.Rpm) + } + pipe.HSet(ctx, hashKey, "__loaded", "1") + pipe.Expire(ctx, hashKey, 5*time.Minute) + if _, err := pipe.Exec(ctx); err != nil { + common.SysError(fmt.Sprintf("user model rate limit: failed to cache rpm for user %d: %v", userId, err)) + rdb.Del(ctx, hashKey) + } + + for _, l := range limits { + if l.Model == modelName { + return l.Rpm + } + } + return 0 +} + +// UserModelRateLimit 用户-模型维度的 RPM 速率限制中间件。 +// 依赖 context 中的 "id"(userId int)和 "original_model"(modelName string)。 +func UserModelRateLimit() gin.HandlerFunc { + duration := time.Duration(setting.ModelRequestRateLimitDurationMinutes) * time.Minute + userModelRateLimiter.Init(duration) + + return func(c *gin.Context) { + // 1. 全局开关 + if !setting.ModelRequestRateLimitEnabled { + c.Next() + return + } + + // 2. 读取 userId / modelName + userId := c.GetInt("id") + modelName := c.GetString("original_model") + if userId == 0 || modelName == "" { + c.Next() + return + } + + // 3. 查 RPM 配置 + rpm := loadUserModelRpm(userId, modelName) + if rpm == 0 { + c.Next() + return + } + + duration := time.Duration(setting.ModelRequestRateLimitDurationMinutes) * time.Minute + + if common.RedisEnabled { + userModelRedisRateLimit(c, userId, modelName, rpm, duration) + } else { + userModelMemoryRateLimit(c, userId, modelName, rpm, duration) + } + } +} + +// userModelRedisRateLimit Redis 令牌桶限流 +func userModelRedisRateLimit(c *gin.Context, userId int, modelName string, rpm int, duration time.Duration) { + ctx := context.Background() + rdb := common.RDB + key := fmt.Sprintf("rateLimit:%s:%d:%s", UserModelRateLimitMark, userId, modelName) + durationSeconds := int64(duration.Seconds()) + + tb := limiter.New(ctx, rdb) + allowed, err := tb.Allow(ctx, key, + limiter.WithCapacity(int64(rpm)*durationSeconds), + limiter.WithRate(int64(rpm)), + limiter.WithRequested(durationSeconds), + ) + if err != nil { + // Redis 异常时降级到内存限流 + userModelMemoryRateLimit(c, userId, modelName, rpm, duration) + return + } + + if !allowed { + abortWithOpenAiMessage(c, http.StatusTooManyRequests, + fmt.Sprintf("用户模型速率限制:每 %d 分钟最多 %d 次请求 (%s)", + setting.ModelRequestRateLimitDurationMinutes, rpm, modelName)) + return + } + + c.Next() + + // 请求失败时退还令牌 + if c.Writer.Status() >= 400 { + _ = tb.Refund(ctx, key, durationSeconds, int64(rpm)*durationSeconds) + } +} + +// userModelMemoryRateLimit 内存滑动窗口限流(参考 memoryRateLimitHandler) +func userModelMemoryRateLimit(c *gin.Context, userId int, modelName string, rpm int, duration time.Duration) { + key := fmt.Sprintf("%s%d:%s", UserModelRateLimitMark, userId, modelName) + durationSeconds := int64(duration.Seconds()) + + if !userModelRateLimiter.Request(key, rpm, durationSeconds) { + abortWithOpenAiMessage(c, http.StatusTooManyRequests, + fmt.Sprintf("用户模型速率限制:每 %d 分钟最多 %d 次请求 (%s)", + setting.ModelRequestRateLimitDurationMinutes, rpm, modelName)) + return + } + + c.Next() + + // 请求失败时退还配额 + if c.Writer.Status() >= 400 { + userModelRateLimiter.Refund(key) + } +} diff --git a/middleware/user_model_rate_limit_test.go b/middleware/user_model_rate_limit_test.go new file mode 100644 index 0000000..779cce4 --- /dev/null +++ b/middleware/user_model_rate_limit_test.go @@ -0,0 +1,216 @@ +package middleware + +import ( + "net/http" + "net/http/httptest" + "sync" + "sync/atomic" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/setting" + "github.com/gin-gonic/gin" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func setupTestDB(t *testing.T) { + t.Helper() + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + sqlDB, _ := db.DB() + sqlDB.SetMaxOpenConns(1) + + origDB := model.DB + model.DB = db + require.NoError(t, db.AutoMigrate(&model.UserModelRateLimit{})) + + t.Cleanup(func() { + model.DB = origDB + sqlDB.Close() + }) +} + +// buildTestRouter 创建用于测试的 gin Engine,前置中间件注入 userId 和 modelName +func buildTestRouter(userId int, modelName string) *gin.Engine { + gin.SetMode(gin.TestMode) + r := gin.New() + r.POST("/test", + func(c *gin.Context) { + c.Set("id", userId) + c.Set("original_model", modelName) + c.Next() + }, + UserModelRateLimit(), + func(c *gin.Context) { + c.JSON(http.StatusOK, gin.H{"ok": true}) + }, + ) + return r +} + +func doRequest(r *gin.Engine) *httptest.ResponseRecorder { + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/test", nil) + r.ServeHTTP(w, req) + return w +} + +// TestUserModelRateLimit_Disabled:全局开关关闭时即使有限流配置也应放行 +func TestUserModelRateLimit_Disabled(t *testing.T) { + setupTestDB(t) + common.RedisEnabled = false + setting.ModelRequestRateLimitEnabled = false + + // 插入 rpm=1 的严格限制 + require.NoError(t, model.DB.Create(&model.UserModelRateLimit{ + UserId: 42, Model: "gpt-4", Rpm: 1, + }).Error) + + r := buildTestRouter(42, "gpt-4") + + // 连续发两次请求,都应该 200 + w1 := doRequest(r) + assert.Equal(t, http.StatusOK, w1.Code) + w2 := doRequest(r) + assert.Equal(t, http.StatusOK, w2.Code) +} + +// TestUserModelRateLimit_NoModel:original_model 为空时放行 +func TestUserModelRateLimit_NoModel(t *testing.T) { + setupTestDB(t) + common.RedisEnabled = false + setting.ModelRequestRateLimitEnabled = true + + r := buildTestRouter(42, "") + + w := doRequest(r) + assert.Equal(t, http.StatusOK, w.Code) +} + +// TestUserModelRateLimit_NoConfig:DB 中无该用户+模型的配置时放行 +func TestUserModelRateLimit_NoConfig(t *testing.T) { + setupTestDB(t) + common.RedisEnabled = false + setting.ModelRequestRateLimitEnabled = true + setting.ModelRequestRateLimitDurationMinutes = 1 + + // DB 中没有任何配置 + r := buildTestRouter(99, "gpt-4") + + w := doRequest(r) + assert.Equal(t, http.StatusOK, w.Code) +} + +// TestUserModelRateLimit_UnderLimit:rpm=2,发 1 次请求,放行 +func TestUserModelRateLimit_UnderLimit(t *testing.T) { + setupTestDB(t) + common.RedisEnabled = false + setting.ModelRequestRateLimitEnabled = true + setting.ModelRequestRateLimitDurationMinutes = 1 + + require.NoError(t, model.DB.Create(&model.UserModelRateLimit{ + UserId: 10, Model: "claude-3", Rpm: 2, + }).Error) + + r := buildTestRouter(10, "claude-3") + + w := doRequest(r) + assert.Equal(t, http.StatusOK, w.Code) +} + +// TestUserModelRateLimit_OverLimit:rpm=2,连续发 3 次请求,第 3 次被 429 +func TestUserModelRateLimit_OverLimit(t *testing.T) { + setupTestDB(t) + common.RedisEnabled = false + setting.ModelRequestRateLimitEnabled = true + setting.ModelRequestRateLimitDurationMinutes = 1 + + require.NoError(t, model.DB.Create(&model.UserModelRateLimit{ + UserId: 20, Model: "gpt-4o", Rpm: 2, + }).Error) + + r := buildTestRouter(20, "gpt-4o") + + w1 := doRequest(r) + assert.Equal(t, http.StatusOK, w1.Code, "第 1 次请求应放行") + + w2 := doRequest(r) + assert.Equal(t, http.StatusOK, w2.Code, "第 2 次请求应放行") + + w3 := doRequest(r) + assert.Equal(t, http.StatusTooManyRequests, w3.Code, "第 3 次请求应被限流") +} + +// TestUserModelRateLimit_ConcurrentMemory:内存路径并发安全验证。 +// RPM=2,并发 20 个请求,预期最多 2 个 200,其余 429。 +func TestUserModelRateLimit_ConcurrentMemory(t *testing.T) { + setupTestDB(t) + common.RedisEnabled = false + setting.ModelRequestRateLimitEnabled = true + setting.ModelRequestRateLimitDurationMinutes = 1 + + require.NoError(t, model.DB.Create(&model.UserModelRateLimit{ + UserId: 30, Model: "deepseek-r1", Rpm: 2, + }).Error) + + r := buildTestRouter(30, "deepseek-r1") + + var successCount int64 + var wg sync.WaitGroup + totalRequests := 20 + + for i := 0; i < totalRequests; i++ { + wg.Add(1) + go func() { + defer wg.Done() + w := doRequest(r) + if w.Code == http.StatusOK { + atomic.AddInt64(&successCount, 1) + } + }() + } + wg.Wait() + + sc := int(atomic.LoadInt64(&successCount)) + assert.LessOrEqual(t, sc, 2, "并发请求中最多 %d 个应成功,实际 %d 个", 2, sc) + assert.GreaterOrEqual(t, sc, 1, "至少应有 1 个请求成功") +} + +// TestUserModelRateLimit_UserIsolation:不同用户之间的限流互不影响 +func TestUserModelRateLimit_UserIsolation(t *testing.T) { + setupTestDB(t) + common.RedisEnabled = false + setting.ModelRequestRateLimitEnabled = true + setting.ModelRequestRateLimitDurationMinutes = 1 + + // 两个用户都设置 RPM=2 + require.NoError(t, model.DB.Create(&model.UserModelRateLimit{ + UserId: 100, Model: "gpt-4", Rpm: 2, + }).Error) + require.NoError(t, model.DB.Create(&model.UserModelRateLimit{ + UserId: 101, Model: "gpt-4", Rpm: 2, + }).Error) + + r1 := buildTestRouter(100, "gpt-4") + r2 := buildTestRouter(101, "gpt-4") + + // 用户 100 用完配额 + w1 := doRequest(r1) + assert.Equal(t, http.StatusOK, w1.Code) + w2 := doRequest(r1) + assert.Equal(t, http.StatusOK, w2.Code) + w3 := doRequest(r1) + assert.Equal(t, http.StatusTooManyRequests, w3.Code) + + // 用户 101 应该不受影响 + w4 := doRequest(r2) + assert.Equal(t, http.StatusOK, w4.Code) + w5 := doRequest(r2) + assert.Equal(t, http.StatusOK, w5.Code) + w6 := doRequest(r2) + assert.Equal(t, http.StatusTooManyRequests, w6.Code) +} diff --git a/model/main.go b/model/main.go index c832546..eac2b49 100644 --- a/model/main.go +++ b/model/main.go @@ -290,6 +290,7 @@ func migrateDB() error { &PendingSyncRecord{}, &QuotaSyncLog{}, &EmailQuotaRule{}, + &UserModelRateLimit{}, ) if err != nil { return err diff --git a/model/user.go b/model/user.go index eeec72c..8e0d5a9 100644 --- a/model/user.go +++ b/model/user.go @@ -65,6 +65,8 @@ type User struct { RemoteUserId int `json:"remote_user_id" gorm:"type:int;default:0;column:remote_user_id"` SyncedQuota int `json:"synced_quota" gorm:"type:int;default:0;column:synced_quota"` LastSyncAt int64 `json:"last_sync_at" gorm:"type:bigint;default:0;column:last_sync_at"` + // Relay 抓包开关 + CaptureRelay bool `json:"capture_relay" gorm:"default:false"` } // IsSyncedUser 判断是否为国内同步用户 @@ -580,11 +582,12 @@ func (user *User) Edit(updatePassword bool) error { newUser := *user updates := map[string]interface{}{ - "username": newUser.Username, - "display_name": newUser.DisplayName, - "group": newUser.Group, - "quota": newUser.Quota, - "remark": newUser.Remark, + "username": newUser.Username, + "display_name": newUser.DisplayName, + "group": newUser.Group, + "quota": newUser.Quota, + "remark": newUser.Remark, + "capture_relay": newUser.CaptureRelay, } if updatePassword { updates["password"] = newUser.Password diff --git a/model/user_capture_test.go b/model/user_capture_test.go new file mode 100644 index 0000000..a748d97 --- /dev/null +++ b/model/user_capture_test.go @@ -0,0 +1,43 @@ +package model + +import ( + "testing" + + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func TestUser_CaptureRelayField_DB(t *testing.T) { + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + defer func() { + sqlDB, _ := db.DB() + sqlDB.Close() + }() + + err = db.AutoMigrate(&User{}) + require.NoError(t, err) + + user := &User{ + Username: "capture-test-user", + Password: "12345678", + DisplayName: "capture test", + CaptureRelay: false, + } + err = db.Create(user).Error + require.NoError(t, err) + + var retrieved User + err = db.First(&retrieved, user.Id).Error + require.NoError(t, err) + assert.False(t, retrieved.CaptureRelay) + + // 测试设置为 true + err = db.Model(&retrieved).Update("capture_relay", true).Error + require.NoError(t, err) + err = db.First(&retrieved, user.Id).Error + require.NoError(t, err) + assert.True(t, retrieved.CaptureRelay) +} diff --git a/model/user_model_rate_limit.go b/model/user_model_rate_limit.go new file mode 100644 index 0000000..67b42ba --- /dev/null +++ b/model/user_model_rate_limit.go @@ -0,0 +1,54 @@ +package model + +import ( + "github.com/QuantumNous/new-api/common" + "gorm.io/gorm" +) + +type UserModelRateLimit struct { + Id int `json:"id" gorm:"primaryKey"` + UserId int `json:"user_id" gorm:"uniqueIndex:idx_user_model_rate_limit"` + Model string `json:"model" gorm:"size:255;uniqueIndex:idx_user_model_rate_limit"` + Rpm int `json:"rpm"` + CreatedAt int64 `json:"created_at"` + UpdatedAt int64 `json:"updated_at"` +} + +// GetUserModelRateLimits 查询指定用户的所有 RPM 配置 +func GetUserModelRateLimits(userId int) ([]UserModelRateLimit, error) { + var list []UserModelRateLimit + err := DB.Where("user_id = ?", userId).Find(&list).Error + return list, err +} + +// SetUserModelRateLimits 覆盖式写入:事务中先硬删再批量插入。 +// 空 slice 时仅删除该用户所有配置。 +func SetUserModelRateLimits(userId int, items []UserModelRateLimit) error { + return DB.Transaction(func(tx *gorm.DB) error { + // 先硬删该用户所有配置 + if err := tx.Where("user_id = ?", userId).Delete(&UserModelRateLimit{}).Error; err != nil { + return err + } + if len(items) == 0 { + return nil + } + now := common.GetTimestamp() + for i := range items { + items[i].UserId = userId + items[i].Id = 0 // 让数据库自增 + items[i].CreatedAt = now + items[i].UpdatedAt = now + } + return tx.Create(&items).Error + }) +} + +// GetUserModelRpm 查单个用户+模型的 RPM,不命中返回 (0, false) +func GetUserModelRpm(userId int, model string) (int, bool) { + var item UserModelRateLimit + err := DB.Where("user_id = ? AND model = ?", userId, model).First(&item).Error + if err != nil { + return 0, false + } + return item.Rpm, true +} diff --git a/model/user_model_rate_limit_test.go b/model/user_model_rate_limit_test.go new file mode 100644 index 0000000..1f727f3 --- /dev/null +++ b/model/user_model_rate_limit_test.go @@ -0,0 +1,110 @@ +package model + +import ( + "testing" + + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func setupUserModelRateLimitDB(t *testing.T) { + t.Helper() + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + sqlDB, _ := db.DB() + sqlDB.SetMaxOpenConns(1) + + origDB := DB + DB = db + + require.NoError(t, db.AutoMigrate(&UserModelRateLimit{})) + + t.Cleanup(func() { + DB = origDB + sqlDB.Close() + }) +} + +func TestGetUserModelRateLimits_Empty(t *testing.T) { + setupUserModelRateLimitDB(t) + + list, err := GetUserModelRateLimits(1) + require.NoError(t, err) + assert.Empty(t, list) +} + +func TestSetUserModelRateLimits_CreateAndOverwrite(t *testing.T) { + setupUserModelRateLimitDB(t) + + // 首次写入 2 条 + items := []UserModelRateLimit{ + {UserId: 1, Model: "gpt-4", Rpm: 60}, + {UserId: 1, Model: "gpt-3.5-turbo", Rpm: 120}, + } + err := SetUserModelRateLimits(1, items) + require.NoError(t, err) + + list, err := GetUserModelRateLimits(1) + require.NoError(t, err) + assert.Len(t, list, 2) + + // 再覆盖写入 1 条(只保留 gpt-4,rpm 改为 30) + items2 := []UserModelRateLimit{ + {UserId: 1, Model: "gpt-4", Rpm: 30}, + } + err = SetUserModelRateLimits(1, items2) + require.NoError(t, err) + + list2, err := GetUserModelRateLimits(1) + require.NoError(t, err) + assert.Len(t, list2, 1) + assert.Equal(t, "gpt-4", list2[0].Model) + assert.Equal(t, 30, list2[0].Rpm) +} + +func TestSetUserModelRateLimits_EmptySlice(t *testing.T) { + setupUserModelRateLimitDB(t) + + // 先写 1 条 + items := []UserModelRateLimit{ + {UserId: 2, Model: "claude-3-opus", Rpm: 10}, + } + err := SetUserModelRateLimits(2, items) + require.NoError(t, err) + + list, err := GetUserModelRateLimits(2) + require.NoError(t, err) + assert.Len(t, list, 1) + + // 再用空 slice 清除 + err = SetUserModelRateLimits(2, []UserModelRateLimit{}) + require.NoError(t, err) + + list2, err := GetUserModelRateLimits(2) + require.NoError(t, err) + assert.Empty(t, list2) +} + +func TestGetUserModelRpm_Found(t *testing.T) { + setupUserModelRateLimitDB(t) + + items := []UserModelRateLimit{ + {UserId: 3, Model: "gemini-pro", Rpm: 20}, + } + err := SetUserModelRateLimits(3, items) + require.NoError(t, err) + + rpm, found := GetUserModelRpm(3, "gemini-pro") + assert.True(t, found) + assert.Equal(t, 20, rpm) +} + +func TestGetUserModelRpm_NotFound(t *testing.T) { + setupUserModelRateLimitDB(t) + + rpm, found := GetUserModelRpm(999, "nonexistent-model") + assert.False(t, found) + assert.Equal(t, 0, rpm) +} diff --git a/router/api-router.go b/router/api-router.go index dcbf578..158c079 100644 --- a/router/api-router.go +++ b/router/api-router.go @@ -142,6 +142,10 @@ func SetApiRouter(router *gin.Engine) { // Admin 2FA routes adminRoute.GET("/2fa/stats", controller.Admin2FAStats) adminRoute.DELETE("/:id/2fa", controller.AdminDisable2FA) + + // User model RPM rate limit routes + adminRoute.GET("/:id/rate-limits", controller.GetUserRateLimits) + adminRoute.PUT("/:id/rate-limits", controller.SetUserRateLimits) } } diff --git a/router/relay-router.go b/router/relay-router.go index dcec439..7ecd05e 100644 --- a/router/relay-router.go +++ b/router/relay-router.go @@ -66,6 +66,7 @@ func SetRelayRouter(router *gin.Engine) { relayV1Router.Use(middleware.SystemPerformanceCheck()) relayV1Router.Use(middleware.TokenAuth()) relayV1Router.Use(middleware.ModelRequestRateLimit()) + relayV1Router.Use(middleware.RelayCaptureMiddleware()) { // WebSocket 路由(统一到 Relay) wsRouter := relayV1Router.Group("") @@ -78,6 +79,7 @@ func SetRelayRouter(router *gin.Engine) { //http router httpRouter := relayV1Router.Group("") httpRouter.Use(middleware.Distribute()) + httpRouter.Use(middleware.UserModelRateLimit()) // claude related routes httpRouter.POST("/messages", func(c *gin.Context) { @@ -182,7 +184,9 @@ func SetRelayRouter(router *gin.Engine) { relayGeminiRouter.Use(middleware.SystemPerformanceCheck()) relayGeminiRouter.Use(middleware.TokenAuth()) relayGeminiRouter.Use(middleware.ModelRequestRateLimit()) + relayGeminiRouter.Use(middleware.RelayCaptureMiddleware()) relayGeminiRouter.Use(middleware.Distribute()) + relayGeminiRouter.Use(middleware.UserModelRateLimit()) { // Gemini API 路径格式: /v1beta/models/{model_name}:{action} relayGeminiRouter.POST("/models/*path", func(c *gin.Context) { diff --git a/web/src/components/table/users/modals/EditUserModal.jsx b/web/src/components/table/users/modals/EditUserModal.jsx index 1cd7b74..ebd936e 100644 --- a/web/src/components/table/users/modals/EditUserModal.jsx +++ b/web/src/components/table/users/modals/EditUserModal.jsx @@ -57,6 +57,7 @@ import { } from '@douyinfe/semi-icons'; import UserBindingManagementModal from './UserBindingManagementModal'; import UserRatioSection from './UserRatioSection'; +import UserModelRateLimitSection from './UserModelRateLimitSection'; const { Text, Title } = Typography; @@ -88,6 +89,7 @@ const EditUserModal = (props) => { quota: 0, group: 'default', remark: '', + capture_relay: false, }); const fetchGroups = async () => { @@ -267,6 +269,16 @@ const EditUserModal = (props) => { showClear /> + + {userId && ( + + + + )} @@ -335,6 +347,12 @@ const EditUserModal = (props) => { )} + {/* 模型速率限制 */} + {userId && ( + + + + )} {/* 绑定信息入口 */} {userId && ( diff --git a/web/src/components/table/users/modals/UserModelRateLimitSection.jsx b/web/src/components/table/users/modals/UserModelRateLimitSection.jsx new file mode 100644 index 0000000..6e169be --- /dev/null +++ b/web/src/components/table/users/modals/UserModelRateLimitSection.jsx @@ -0,0 +1,192 @@ +import React, { useEffect, useState, useRef } from 'react'; +import { useTranslation } from 'react-i18next'; +import { userRateLimitApi, showError, showSuccess } from '../../../../helpers'; +import { + Button, + Table, + Modal, + Form, + Avatar, + Typography, + Popconfirm, +} from '@douyinfe/semi-ui'; +import { IconDelete, IconPlus, IconClock } from '@douyinfe/semi-icons'; + +const { Text } = Typography; + +const UserModelRateLimitSection = ({ userId }) => { + const { t } = useTranslation(); + const [limits, setLimits] = useState([]); + const [loading, setLoading] = useState(false); + const [addModalVisible, setAddModalVisible] = useState(false); + const formApiRef = useRef(null); + + const loadLimits = async () => { + if (!userId) return; + setLoading(true); + try { + const res = await userRateLimitApi.get(userId); + const { success, data, message } = res.data; + if (success) { + setLimits(data || []); + } else { + showError(message); + } + } catch (e) { + showError(e.message); + } + setLoading(false); + }; + + useEffect(() => { + loadLimits(); + }, [userId]); + + const saveLimits = async (newLimits) => { + try { + const res = await userRateLimitApi.set(userId, newLimits); + const { success, message } = res.data; + if (success) { + showSuccess(t('保存成功')); + return true; + } else { + showError(message); + return false; + } + } catch (e) { + showError(e.message); + return false; + } + }; + + const handleAdd = async (values) => { + const rpm = parseInt(values.rpm) || 0; + const newItem = { model: values.model, rpm }; + const filtered = limits.filter((item) => item.model !== values.model); + const newLimits = [...filtered, newItem]; + const ok = await saveLimits(newLimits); + if (ok) { + setLimits(newLimits); + setAddModalVisible(false); + formApiRef.current?.reset(); + } + }; + + const handleDelete = async (model) => { + const newLimits = limits.filter((item) => item.model !== model); + const ok = await saveLimits(newLimits); + if (ok) { + setLimits(newLimits); + } + }; + + const columns = [ + { + title: t('模型'), + dataIndex: 'model', + key: 'model', + }, + { + title: t('RPM 限制'), + dataIndex: 'rpm', + key: 'rpm', + render: (rpm) => + rpm === 0 ? ( + {t('不限制')} + ) : ( + {rpm} + ), + }, + { + title: t('操作'), + key: 'action', + render: (_, record) => ( + handleDelete(record.model)} + > + + + ), + }, + ]; + + return ( + <> +
+
+ + + +
+ + {t('模型速率限制(RPM)')} + +
+ {t('为该用户配置各模型每分钟请求数上限,0 表示不限制')} +
+
+
+ +
+ + + + formApiRef.current?.submitForm()} + onCancel={() => { + setAddModalVisible(false); + formApiRef.current?.reset(); + }} + > +
(formApiRef.current = api)} + onSubmit={handleAdd} + > + + + +
+ + ); +}; + +export default UserModelRateLimitSection; diff --git a/web/src/helpers/api.js b/web/src/helpers/api.js index 5116e55..0b73271 100644 --- a/web/src/helpers/api.js +++ b/web/src/helpers/api.js @@ -418,3 +418,9 @@ export const pricingTagApi = { update: (id, data) => API.put(`/api/pricing_tag/${id}`, data), delete: (id) => API.delete(`/api/pricing_tag/${id}`), }; + +// 用户模型限流 API +export const userRateLimitApi = { + get: (userId) => API.get(`/api/user/${userId}/rate-limits`), + set: (userId, items) => API.put(`/api/user/${userId}/rate-limits`, items), +};