diff --git a/controller/auth_session.go b/controller/auth_session.go index 0f1599c5..3921c779 100644 --- a/controller/auth_session.go +++ b/controller/auth_session.go @@ -2,9 +2,11 @@ package controller import ( "errors" + "fmt" "net/http" "strings" + "github.com/QuantumNous/new-api/logger" "github.com/QuantumNous/new-api/middleware" "github.com/QuantumNous/new-api/model" "github.com/QuantumNous/new-api/service" @@ -164,6 +166,12 @@ func writeAuthSessionError(c *gin.Context, err error) { if errors.Is(err, gorm.ErrRecordNotFound) { status, code = http.StatusUnauthorized, "AUTH_UNAUTHORIZED" } + if status == http.StatusInternalServerError { + // The response body only carries the generic AUTH_INTERNAL_ERROR + // code; without this log the underlying Redis/database/session + // failure is indistinguishable from the client side. + logger.LogError(c.Request.Context(), fmt.Sprintf("auth session internal error (%s %s): %v", c.Request.Method, c.Request.URL.Path, err)) + } c.JSON(status, gin.H{"success": false, "code": code, "message": http.StatusText(status)}) } diff --git a/middleware/rate-limit.go b/middleware/rate-limit.go index a6389161..350ee8ec 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/logger" "github.com/gin-gonic/gin" ) @@ -106,33 +107,43 @@ func redisFixedWindowTake(ctx context.Context, key string, maxRequestNum int, du } func redisRateLimiter(c *gin.Context, maxRequestNum int, duration int64, mark string) { - allowed, _, _, err := redisFixedWindowTake( + allowed, _, ttlSeconds, err := redisFixedWindowTake( c.Request.Context(), redisIPRateLimitKey(mark, c.ClientIP()), maxRequestNum, duration, ) if err != nil { - fmt.Println(err.Error()) + logger.LogError(c.Request.Context(), fmt.Sprintf("rate limit check failed (mark=%s): %v", mark, err)) c.Status(http.StatusInternalServerError) c.Abort() return } if !allowed { - c.Status(http.StatusTooManyRequests) - c.Abort() + writeRateLimited(c, ttlSeconds) } } func memoryRateLimiter(c *gin.Context, maxRequestNum int, duration int64, mark string) { key := mark + c.ClientIP() if !inMemoryRateLimiter.Request(key, maxRequestNum, duration) { - c.Status(http.StatusTooManyRequests) - c.Abort() + writeRateLimited(c, duration) return } } +// 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. +// 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() +} + func rateLimitFactory(maxRequestNum int, duration int64, mark string) func(c *gin.Context) { if common.RedisEnabled { return func(c *gin.Context) { @@ -201,8 +212,7 @@ func userRateLimitFactory(maxRequestNum int, duration int64, mark string) func(c } key := fmt.Sprintf("%s:user:%d", mark, userID) if !inMemoryRateLimiter.Request(key, maxRequestNum, duration) { - c.Status(http.StatusTooManyRequests) - c.Abort() + writeRateLimited(c, duration) return } } @@ -211,16 +221,15 @@ func userRateLimitFactory(maxRequestNum int, duration int64, mark string) func(c // userRedisRateLimiter is like redisRateLimiter but accepts a pre-built key // (to support user-ID-based keys). func userRedisRateLimiter(c *gin.Context, maxRequestNum int, duration int64, key string) { - allowed, _, _, err := redisFixedWindowTake(c.Request.Context(), key, maxRequestNum, duration) + allowed, _, ttlSeconds, err := redisFixedWindowTake(c.Request.Context(), key, maxRequestNum, duration) if err != nil { - fmt.Println(err.Error()) + logger.LogError(c.Request.Context(), fmt.Sprintf("rate limit check failed (key=%s): %v", key, err)) c.Status(http.StatusInternalServerError) c.Abort() return } if !allowed { - c.Status(http.StatusTooManyRequests) - c.Abort() + writeRateLimited(c, ttlSeconds) } } diff --git a/middleware/rate_limit_test.go b/middleware/rate_limit_test.go index b955753b..1ca48372 100644 --- a/middleware/rate_limit_test.go +++ b/middleware/rate_limit_test.go @@ -61,7 +61,9 @@ func TestRedisIPRateLimiterThresholdTTLAndNamespace(t *testing.T) { require.NoError(t, err) assert.Equal(t, http.StatusNoContent, performRateLimitRequest(router, "/limited", remoteAddr).Code) assert.Equal(t, http.StatusNoContent, performRateLimitRequest(router, "/limited", remoteAddr).Code) - assert.Equal(t, http.StatusTooManyRequests, performRateLimitRequest(router, "/limited", remoteAddr).Code) + limitedResponse := performRateLimitRequest(router, "/limited", remoteAddr) + assert.Equal(t, http.StatusTooManyRequests, limitedResponse.Code) + assert.Equal(t, "37", limitedResponse.Header().Get("Retry-After")) key := redisIPRateLimitKey("TEST", "192.0.2.10") count, err := redisServer.Get(key) diff --git a/web/src/lib/auth-session.test.ts b/web/src/lib/auth-session.test.ts index 7797684b..4a97d3c2 100644 --- a/web/src/lib/auth-session.test.ts +++ b/web/src/lib/auth-session.test.ts @@ -147,6 +147,30 @@ describe('authentication session coordination', () => { assert.equal(transientCount, 1) }) + test('a rate limited refresh remains retryable without clearing the session', async () => { + let transientCount = 0 + let clearCount = 0 + const runtime: AuthRefreshRuntime = { + request: async () => ({ status: 429 }), + getExpectedSID: () => bundle.session.sid, + parseBundle: () => null, + acceptBundle: () => undefined, + clear: () => { + clearCount += 1 + }, + markTransient: () => { + transientCount += 1 + }, + wait: async () => undefined, + } + + const outcome = await createRefreshRunner(runtime)() + + assert.equal(outcome.kind, 'transient_error') + assert.equal(clearCount, 0) + assert.equal(transientCount, 1) + }) + test('an exhausted refresh race clears the unusable local session', async () => { const requestedDelays: number[] = [] const clears: Array<[boolean, string | undefined]> = [] diff --git a/web/src/lib/auth-session.ts b/web/src/lib/auth-session.ts index 3921af6b..67c5fe8b 100644 --- a/web/src/lib/auth-session.ts +++ b/web/src/lib/auth-session.ts @@ -253,7 +253,7 @@ export function createRefreshRunner( return { kind: 'anonymous' } } - if (!response.status || response.status >= 500) { + if (!response.status || response.status >= 500 || response.status === 429) { runtime.markTransient() return { kind: 'transient_error',