diff --git a/i18n/keys.go b/i18n/keys.go index 8e9a4b5694d4..3adf67601711 100644 --- a/i18n/keys.go +++ b/i18n/keys.go @@ -234,6 +234,7 @@ const ( // Rate limit related messages const ( + MsgRateLimitExceeded = "rate_limit.exceeded" MsgRateLimitReached = "rate_limit.reached" MsgRateLimitTotalReached = "rate_limit.total_reached" ) diff --git a/i18n/locales/en.yaml b/i18n/locales/en.yaml index 3f1fd03cb090..53839c45cdb7 100644 --- a/i18n/locales/en.yaml +++ b/i18n/locales/en.yaml @@ -198,6 +198,7 @@ twofa.record_id_empty: "2FA record ID cannot be empty" twofa.code_invalid: "Verification code or backup code is incorrect" # Rate limit messages +rate_limit.exceeded: "Too many requests. Please try again later." rate_limit.reached: "You have reached the request limit: maximum {{.Max}} requests in {{.Minutes}} minutes" rate_limit.total_reached: "You have reached the total request limit: maximum {{.Max}} requests in {{.Minutes}} minutes, including failed attempts" diff --git a/i18n/locales/zh-CN.yaml b/i18n/locales/zh-CN.yaml index fe982e59a0f3..7ae1bb3f07de 100644 --- a/i18n/locales/zh-CN.yaml +++ b/i18n/locales/zh-CN.yaml @@ -199,6 +199,7 @@ twofa.record_id_empty: "2FA记录ID不能为空" twofa.code_invalid: "验证码或备用码不正确" # Rate limit messages +rate_limit.exceeded: "请求过于频繁，请稍后再试" rate_limit.reached: "您已达到请求数限制：{{.Minutes}}分钟内最多请求{{.Max}}次" rate_limit.total_reached: "您已达到总请求数限制：{{.Minutes}}分钟内最多请求{{.Max}}次，包括失败次数" diff --git a/i18n/locales/zh-TW.yaml b/i18n/locales/zh-TW.yaml index 27759d07f378..147086c7310b 100644 --- a/i18n/locales/zh-TW.yaml +++ b/i18n/locales/zh-TW.yaml @@ -199,6 +199,7 @@ twofa.record_id_empty: "2FA記錄ID不能為空" twofa.code_invalid: "驗證碼或備用碼不正確" # Rate limit messages +rate_limit.exceeded: "請求過於頻繁，請稍後再試" rate_limit.reached: "您已達到請求數限制：{{.Minutes}}分鐘內最多請求{{.Max}}次" rate_limit.total_reached: "您已達到總請求數限制：{{.Minutes}}分鐘內最多請求{{.Max}}次，包括失敗次數" diff --git a/middleware/rate-limit.go b/middleware/rate-limit.go index 350ee8ec0718..a357a4237586 100644 --- a/middleware/rate-limit.go +++ b/middleware/rate-limit.go @@ -8,6 +8,7 @@ import ( "strconv" "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/i18n" "github.com/QuantumNous/new-api/logger" "github.com/gin-gonic/gin" ) @@ -124,6 +125,7 @@ func redisRateLimiter(c *gin.Context, maxRequestNum int, duration int64, mark st } } +// memoryRateLimiter enforces an IP-based limit using process-local state. func memoryRateLimiter(c *gin.Context, maxRequestNum int, duration int64, mark string) { key := mark + c.ClientIP() if !inMemoryRateLimiter.Request(key, maxRequestNum, duration) { @@ -132,18 +134,21 @@ func memoryRateLimiter(c *gin.Context, maxRequestNum int, duration int64, mark s } } -// writeRateLimited rejects the request with 429 and a Retry-After hint so -// clients can back off instead of treating the rejection as a fatal error. +// writeRateLimited rejects the request with a structured 429 response and a +// Retry-After hint so clients can report the error and back off safely. // The in-memory limiter cannot report the remaining window, so callers // without a TTL pass the full window duration as a conservative upper bound. func writeRateLimited(c *gin.Context, retryAfterSeconds int64) { if retryAfterSeconds > 0 { c.Header("Retry-After", strconv.FormatInt(retryAfterSeconds, 10)) } - c.Status(http.StatusTooManyRequests) - c.Abort() + c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{ + "success": false, + "message": common.TranslateMessage(c, i18n.MsgRateLimitExceeded), + }) } +// rateLimitFactory selects the configured backend for an IP-based rate limiter. func rateLimitFactory(maxRequestNum int, duration int64, mark string) func(c *gin.Context) { if common.RedisEnabled { return func(c *gin.Context) { diff --git a/middleware/rate_limit_test.go b/middleware/rate_limit_test.go index 1ca48372b494..6f0698b0dd5e 100644 --- a/middleware/rate_limit_test.go +++ b/middleware/rate_limit_test.go @@ -2,6 +2,7 @@ package middleware import ( "context" + "fmt" "net/http" "net/http/httptest" "sync" @@ -10,6 +11,7 @@ import ( "time" "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/i18n" "github.com/alicebob/miniredis/v2" "github.com/gin-gonic/gin" "github.com/go-redis/redis/v8" @@ -17,6 +19,9 @@ import ( "github.com/stretchr/testify/require" ) +var memoryCriticalRateLimitTestRun atomic.Uint64 + +// useRateLimitMiniRedis configures an isolated Redis backend for rate-limit tests. func useRateLimitMiniRedis(t *testing.T) (*miniredis.Miniredis, *redis.Client) { t.Helper() @@ -37,6 +42,7 @@ func useRateLimitMiniRedis(t *testing.T) (*miniredis.Miniredis, *redis.Client) { return redisServer, redisClient } +// performRateLimitRequest records a GET request with a controlled client address. func performRateLimitRequest(router http.Handler, path string, remoteAddr string) *httptest.ResponseRecorder { recorder := httptest.NewRecorder() request := httptest.NewRequest(http.MethodGet, path, nil) @@ -45,18 +51,69 @@ func performRateLimitRequest(router http.Handler, path string, remoteAddr string return recorder } -func TestRedisIPRateLimiterThresholdTTLAndNamespace(t *testing.T) { +// TestMemoryCriticalRateLimiterReturnsJSONWithRetryAfter verifies the in-memory 429 response contract. +func TestMemoryCriticalRateLimiterReturnsJSONWithRetryAfter(t *testing.T) { + gin.SetMode(gin.TestMode) + require.NoError(t, i18n.Init()) + + previousRedisEnabled := common.RedisEnabled + previousEnabled := common.CriticalRateLimitEnable + previousMaxRequests := common.CriticalRateLimitNum + previousDuration := common.CriticalRateLimitDuration + common.RedisEnabled = false + common.CriticalRateLimitEnable = true + common.CriticalRateLimitNum = 2 + common.CriticalRateLimitDuration = 43 + t.Cleanup(func() { + common.RedisEnabled = previousRedisEnabled + common.CriticalRateLimitEnable = previousEnabled + common.CriticalRateLimitNum = previousMaxRequests + common.CriticalRateLimitDuration = previousDuration + }) + + router := gin.New() + require.NoError(t, router.SetTrustedProxies(nil)) + router.GET("/limited", CriticalRateLimit(), func(c *gin.Context) { + c.Status(http.StatusNoContent) + }) + + runID := memoryCriticalRateLimitTestRun.Add(1) + remoteAddr := fmt.Sprintf("[2001:db8::%x]:12345", runID) + assert.Equal(t, http.StatusNoContent, performRateLimitRequest(router, "/limited", remoteAddr).Code) + assert.Equal(t, http.StatusNoContent, performRateLimitRequest(router, "/limited", remoteAddr).Code) + limitedResponse := performRateLimitRequest(router, "/limited", remoteAddr) + assert.Equal(t, http.StatusTooManyRequests, limitedResponse.Code) + assert.Equal(t, "43", limitedResponse.Header().Get("Retry-After")) + assert.Equal(t, "application/json; charset=utf-8", limitedResponse.Header().Get("Content-Type")) + assert.JSONEq(t, `{"success":false,"message":"Too many requests. Please try again later."}`, limitedResponse.Body.String()) +} + +// TestRedisCriticalRateLimiterReturnsJSONWithRetryAfter verifies the Redis-backed 429 response contract. +func TestRedisCriticalRateLimiterReturnsJSONWithRetryAfter(t *testing.T) { gin.SetMode(gin.TestMode) + require.NoError(t, i18n.Init()) redisServer, _ := useRateLimitMiniRedis(t) + previousEnabled := common.CriticalRateLimitEnable + previousMaxRequests := common.CriticalRateLimitNum + previousDuration := common.CriticalRateLimitDuration + common.CriticalRateLimitEnable = true + common.CriticalRateLimitNum = 2 + common.CriticalRateLimitDuration = 37 + t.Cleanup(func() { + common.CriticalRateLimitEnable = previousEnabled + common.CriticalRateLimitNum = previousMaxRequests + common.CriticalRateLimitDuration = previousDuration + }) + router := gin.New() require.NoError(t, router.SetTrustedProxies(nil)) - router.GET("/limited", rateLimitFactory(2, 37, "TEST"), func(c *gin.Context) { + router.GET("/limited", CriticalRateLimit(), func(c *gin.Context) { c.Status(http.StatusNoContent) }) remoteAddr := "192.0.2.10:12345" - legacyKey := "rateLimit:TEST192.0.2.10" + legacyKey := "rateLimit:CT192.0.2.10" _, err := redisServer.Push(legacyKey, "legacy-list-entry") require.NoError(t, err) assert.Equal(t, http.StatusNoContent, performRateLimitRequest(router, "/limited", remoteAddr).Code) @@ -64,8 +121,10 @@ func TestRedisIPRateLimiterThresholdTTLAndNamespace(t *testing.T) { limitedResponse := performRateLimitRequest(router, "/limited", remoteAddr) assert.Equal(t, http.StatusTooManyRequests, limitedResponse.Code) assert.Equal(t, "37", limitedResponse.Header().Get("Retry-After")) + assert.Equal(t, "application/json; charset=utf-8", limitedResponse.Header().Get("Content-Type")) + assert.JSONEq(t, `{"success":false,"message":"Too many requests. Please try again later."}`, limitedResponse.Body.String()) - key := redisIPRateLimitKey("TEST", "192.0.2.10") + key := redisIPRateLimitKey("CT", "192.0.2.10") count, err := redisServer.Get(key) require.NoError(t, err) assert.Equal(t, "3", count)