| @@ -13,9 +13,13 @@ import ( | |||||
| //go:embed lua/rate_limit.lua | //go:embed lua/rate_limit.lua | ||||
| var rateLimitScript string | var rateLimitScript string | ||||
| //go:embed lua/rate_refund.lua | |||||
| var rateRefundScript string | |||||
| type RedisLimiter struct { | type RedisLimiter struct { | ||||
| client *redis.Client | |||||
| limitScriptSHA string | |||||
| client *redis.Client | |||||
| limitScriptSHA string | |||||
| refundScriptSHA string | |||||
| } | } | ||||
| var ( | var ( | ||||
| @@ -30,9 +34,14 @@ func New(ctx context.Context, r *redis.Client) *RedisLimiter { | |||||
| if err != nil { | if err != nil { | ||||
| common.SysLog(fmt.Sprintf("Failed to load rate limit script: %v", err)) | 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{ | 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 | 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 配置选项模式 | // Config 配置选项模式 | ||||
| type Config struct { | type Config struct { | ||||
| Capacity int64 | Capacity int64 | ||||
| @@ -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 | |||||
| @@ -68,3 +68,15 @@ func (l *InMemoryRateLimiter) Request(key string, maxRequestNum int, duration in | |||||
| } | } | ||||
| return true | 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 | |||||
| } | |||||
| @@ -14,6 +14,7 @@ import ( | |||||
| "github.com/QuantumNous/new-api/dto" | "github.com/QuantumNous/new-api/dto" | ||||
| "github.com/QuantumNous/new-api/i18n" | "github.com/QuantumNous/new-api/i18n" | ||||
| "github.com/QuantumNous/new-api/logger" | "github.com/QuantumNous/new-api/logger" | ||||
| "github.com/QuantumNous/new-api/middleware" | |||||
| "github.com/QuantumNous/new-api/model" | "github.com/QuantumNous/new-api/model" | ||||
| "github.com/QuantumNous/new-api/service" | "github.com/QuantumNous/new-api/service" | ||||
| "github.com/QuantumNous/new-api/service/region_sync" | "github.com/QuantumNous/new-api/service/region_sync" | ||||
| @@ -578,6 +579,7 @@ func UpdateUser(c *gin.Context) { | |||||
| common.ApiError(c, err) | common.ApiError(c, err) | ||||
| return | return | ||||
| } | } | ||||
| middleware.SetCaptureEnabled(int64(updatedUser.Id), updatedUser.CaptureRelay) | |||||
| if originUser.Quota != updatedUser.Quota { | if originUser.Quota != updatedUser.Quota { | ||||
| model.RecordLog(originUser.Id, model.LogTypeManage, fmt.Sprintf("管理员将用户额度从 %s修改为 %s", logger.LogQuota(originUser.Quota), logger.LogQuota(updatedUser.Quota))) | model.RecordLog(originUser.Id, model.LogTypeManage, fmt.Sprintf("管理员将用户额度从 %s修改为 %s", logger.LogQuota(originUser.Quota), logger.LogQuota(updatedUser.Quota))) | ||||
| } | } | ||||
| @@ -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": "", | |||||
| }) | |||||
| } | |||||
| @@ -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) | |||||
| } | |||||
| @@ -293,6 +293,11 @@ func InitResources() error { | |||||
| // 初始化模型 | // 初始化模型 | ||||
| model.GetPricing() | model.GetPricing() | ||||
| // 加载抓包用户缓存 | |||||
| if err := middleware.LoadCaptureEnabledUsers(); err != nil { | |||||
| common.SysError("failed to load capture enabled users: " + err.Error()) | |||||
| } | |||||
| // Initialize SQL Database | // Initialize SQL Database | ||||
| err = model.InitLogDB() | err = model.InitLogDB() | ||||
| if err != nil { | if err != nil { | ||||
| @@ -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) | |||||
| } | |||||
| } | |||||
| } | |||||
| @@ -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") | |||||
| } | |||||
| } | |||||
| @@ -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") | |||||
| } | |||||
| @@ -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) | |||||
| } | |||||
| } | |||||
| @@ -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) | |||||
| } | |||||
| @@ -290,6 +290,7 @@ func migrateDB() error { | |||||
| &PendingSyncRecord{}, | &PendingSyncRecord{}, | ||||
| &QuotaSyncLog{}, | &QuotaSyncLog{}, | ||||
| &EmailQuotaRule{}, | &EmailQuotaRule{}, | ||||
| &UserModelRateLimit{}, | |||||
| ) | ) | ||||
| if err != nil { | if err != nil { | ||||
| return err | return err | ||||
| @@ -65,6 +65,8 @@ type User struct { | |||||
| RemoteUserId int `json:"remote_user_id" gorm:"type:int;default:0;column:remote_user_id"` | 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"` | 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"` | 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 判断是否为国内同步用户 | // IsSyncedUser 判断是否为国内同步用户 | ||||
| @@ -580,11 +582,12 @@ func (user *User) Edit(updatePassword bool) error { | |||||
| newUser := *user | newUser := *user | ||||
| updates := map[string]interface{}{ | 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 { | if updatePassword { | ||||
| updates["password"] = newUser.Password | updates["password"] = newUser.Password | ||||
| @@ -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) | |||||
| } | |||||
| @@ -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 | |||||
| } | |||||
| @@ -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) | |||||
| } | |||||
| @@ -142,6 +142,10 @@ func SetApiRouter(router *gin.Engine) { | |||||
| // Admin 2FA routes | // Admin 2FA routes | ||||
| adminRoute.GET("/2fa/stats", controller.Admin2FAStats) | adminRoute.GET("/2fa/stats", controller.Admin2FAStats) | ||||
| adminRoute.DELETE("/:id/2fa", controller.AdminDisable2FA) | 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) | |||||
| } | } | ||||
| } | } | ||||
| @@ -66,6 +66,7 @@ func SetRelayRouter(router *gin.Engine) { | |||||
| relayV1Router.Use(middleware.SystemPerformanceCheck()) | relayV1Router.Use(middleware.SystemPerformanceCheck()) | ||||
| relayV1Router.Use(middleware.TokenAuth()) | relayV1Router.Use(middleware.TokenAuth()) | ||||
| relayV1Router.Use(middleware.ModelRequestRateLimit()) | relayV1Router.Use(middleware.ModelRequestRateLimit()) | ||||
| relayV1Router.Use(middleware.RelayCaptureMiddleware()) | |||||
| { | { | ||||
| // WebSocket 路由(统一到 Relay) | // WebSocket 路由(统一到 Relay) | ||||
| wsRouter := relayV1Router.Group("") | wsRouter := relayV1Router.Group("") | ||||
| @@ -78,6 +79,7 @@ func SetRelayRouter(router *gin.Engine) { | |||||
| //http router | //http router | ||||
| httpRouter := relayV1Router.Group("") | httpRouter := relayV1Router.Group("") | ||||
| httpRouter.Use(middleware.Distribute()) | httpRouter.Use(middleware.Distribute()) | ||||
| httpRouter.Use(middleware.UserModelRateLimit()) | |||||
| // claude related routes | // claude related routes | ||||
| httpRouter.POST("/messages", func(c *gin.Context) { | httpRouter.POST("/messages", func(c *gin.Context) { | ||||
| @@ -182,7 +184,9 @@ func SetRelayRouter(router *gin.Engine) { | |||||
| relayGeminiRouter.Use(middleware.SystemPerformanceCheck()) | relayGeminiRouter.Use(middleware.SystemPerformanceCheck()) | ||||
| relayGeminiRouter.Use(middleware.TokenAuth()) | relayGeminiRouter.Use(middleware.TokenAuth()) | ||||
| relayGeminiRouter.Use(middleware.ModelRequestRateLimit()) | relayGeminiRouter.Use(middleware.ModelRequestRateLimit()) | ||||
| relayGeminiRouter.Use(middleware.RelayCaptureMiddleware()) | |||||
| relayGeminiRouter.Use(middleware.Distribute()) | relayGeminiRouter.Use(middleware.Distribute()) | ||||
| relayGeminiRouter.Use(middleware.UserModelRateLimit()) | |||||
| { | { | ||||
| // Gemini API 路径格式: /v1beta/models/{model_name}:{action} | // Gemini API 路径格式: /v1beta/models/{model_name}:{action} | ||||
| relayGeminiRouter.POST("/models/*path", func(c *gin.Context) { | relayGeminiRouter.POST("/models/*path", func(c *gin.Context) { | ||||
| @@ -57,6 +57,7 @@ import { | |||||
| } from '@douyinfe/semi-icons'; | } from '@douyinfe/semi-icons'; | ||||
| import UserBindingManagementModal from './UserBindingManagementModal'; | import UserBindingManagementModal from './UserBindingManagementModal'; | ||||
| import UserRatioSection from './UserRatioSection'; | import UserRatioSection from './UserRatioSection'; | ||||
| import UserModelRateLimitSection from './UserModelRateLimitSection'; | |||||
| const { Text, Title } = Typography; | const { Text, Title } = Typography; | ||||
| @@ -88,6 +89,7 @@ const EditUserModal = (props) => { | |||||
| quota: 0, | quota: 0, | ||||
| group: 'default', | group: 'default', | ||||
| remark: '', | remark: '', | ||||
| capture_relay: false, | |||||
| }); | }); | ||||
| const fetchGroups = async () => { | const fetchGroups = async () => { | ||||
| @@ -267,6 +269,16 @@ const EditUserModal = (props) => { | |||||
| showClear | showClear | ||||
| /> | /> | ||||
| </Col> | </Col> | ||||
| {userId && ( | |||||
| <Col span={24}> | |||||
| <Form.Switch | |||||
| field='capture_relay' | |||||
| label={t('Relay 抓包')} | |||||
| helpText={t('开启后该用户所有 relay 请求的请求体和响应体将保存到服务器本地文件,仅用于调试')} | |||||
| /> | |||||
| </Col> | |||||
| )} | |||||
| </Row> | </Row> | ||||
| </Card> | </Card> | ||||
| @@ -335,6 +347,12 @@ const EditUserModal = (props) => { | |||||
| <UserRatioSection userId={userId} /> | <UserRatioSection userId={userId} /> | ||||
| </Card> | </Card> | ||||
| )} | )} | ||||
| {/* 模型速率限制 */} | |||||
| {userId && ( | |||||
| <Card className='!rounded-2xl shadow-sm border-0'> | |||||
| <UserModelRateLimitSection userId={userId} /> | |||||
| </Card> | |||||
| )} | |||||
| {/* 绑定信息入口 */} | {/* 绑定信息入口 */} | ||||
| {userId && ( | {userId && ( | ||||
| <Card className='!rounded-2xl shadow-sm border-0'> | <Card className='!rounded-2xl shadow-sm border-0'> | ||||
| @@ -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 ? ( | |||||
| <Text type='tertiary'>{t('不限制')}</Text> | |||||
| ) : ( | |||||
| <Text>{rpm}</Text> | |||||
| ), | |||||
| }, | |||||
| { | |||||
| title: t('操作'), | |||||
| key: 'action', | |||||
| render: (_, record) => ( | |||||
| <Popconfirm | |||||
| title={t('确认删除该限制?')} | |||||
| onConfirm={() => handleDelete(record.model)} | |||||
| > | |||||
| <Button | |||||
| type='danger' | |||||
| size='small' | |||||
| icon={<IconDelete />} | |||||
| > | |||||
| {t('删除')} | |||||
| </Button> | |||||
| </Popconfirm> | |||||
| ), | |||||
| }, | |||||
| ]; | |||||
| return ( | |||||
| <> | |||||
| <div className='flex items-center justify-between mb-2'> | |||||
| <div className='flex items-center'> | |||||
| <Avatar size='small' color='cyan' className='mr-2 shadow-md'> | |||||
| <IconClock size={16} /> | |||||
| </Avatar> | |||||
| <div> | |||||
| <Text className='text-lg font-medium'> | |||||
| {t('模型速率限制(RPM)')} | |||||
| </Text> | |||||
| <div className='text-xs text-gray-600'> | |||||
| {t('为该用户配置各模型每分钟请求数上限,0 表示不限制')} | |||||
| </div> | |||||
| </div> | |||||
| </div> | |||||
| <Button | |||||
| size='small' | |||||
| icon={<IconPlus />} | |||||
| onClick={() => setAddModalVisible(true)} | |||||
| > | |||||
| {t('添加限制')} | |||||
| </Button> | |||||
| </div> | |||||
| <Table | |||||
| columns={columns} | |||||
| dataSource={limits} | |||||
| loading={loading} | |||||
| rowKey='model' | |||||
| size='small' | |||||
| pagination={false} | |||||
| empty={t('暂无速率限制')} | |||||
| /> | |||||
| <Modal | |||||
| title={t('添加模型速率限制')} | |||||
| visible={addModalVisible} | |||||
| onOk={() => formApiRef.current?.submitForm()} | |||||
| onCancel={() => { | |||||
| setAddModalVisible(false); | |||||
| formApiRef.current?.reset(); | |||||
| }} | |||||
| > | |||||
| <Form | |||||
| getFormApi={(api) => (formApiRef.current = api)} | |||||
| onSubmit={handleAdd} | |||||
| > | |||||
| <Form.Input | |||||
| field='model' | |||||
| label={t('模型名称')} | |||||
| placeholder={t('请输入精确的模型名,如 gpt-4o')} | |||||
| rules={[{ required: true, message: t('请输入模型名称') }]} | |||||
| style={{ width: '100%' }} | |||||
| /> | |||||
| <Form.InputNumber | |||||
| field='rpm' | |||||
| label={t('RPM 限制')} | |||||
| placeholder='60' | |||||
| initValue={60} | |||||
| min={0} | |||||
| step={10} | |||||
| extraText={t('0 表示不限制')} | |||||
| rules={[{ required: true, message: t('请输入 RPM 限制') }]} | |||||
| style={{ width: '100%' }} | |||||
| /> | |||||
| </Form> | |||||
| </Modal> | |||||
| </> | |||||
| ); | |||||
| }; | |||||
| export default UserModelRateLimitSection; | |||||
| @@ -418,3 +418,9 @@ export const pricingTagApi = { | |||||
| update: (id, data) => API.put(`/api/pricing_tag/${id}`, data), | update: (id, data) => API.put(`/api/pricing_tag/${id}`, data), | ||||
| delete: (id) => API.delete(`/api/pricing_tag/${id}`), | 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), | |||||
| }; | |||||