fix: harden relay reliability and billing integrity

.env.example

@@ -60,6 +60,8 @@
 # RELAY_IDLE_CONN_TIMEOUT=90
 # 流模式无响应超时时间，单位秒，如果出现空补全可以尝试改为更大值
 # STREAMING_TIMEOUT=300
+# 流式请求等待上游首字节的独立超时，单位秒；0 表示不限制且不会影响后续流持续时间
+# STREAMING_FIRST_BYTE_TIMEOUT=0
 
 # TLS / HTTP 跳过验证设置
 # TLS_INSECURE_SKIP_VERIFY=false

README.en.md

@@ -311,6 +311,7 @@ docker run --name new-api -d --restart always \
 | `SQL_DSN` | Database connection string | - |
 | `REDIS_CONN_STRING` | Redis connection string | - |
 | `STREAMING_TIMEOUT` | Streaming timeout (seconds) | `300` |
+| `STREAMING_FIRST_BYTE_TIMEOUT` | Maximum wait for streaming response headers/first byte without limiting the stream duration (seconds, `0` disables) | `0` |
 | `STREAM_SCANNER_MAX_BUFFER_MB` | Max per-line buffer (MB) for the stream scanner; increase when upstream sends huge image/base64 payloads | `64` |
 | `MAX_REQUEST_BODY_MB` | Max request body size (MB, counted **after decompression**; prevents huge requests/zip bombs from exhausting memory). Exceeding it returns `413` | `32` |
 | `AZURE_DEFAULT_API_VERSION` | Azure API version | `2025-04-01-preview` |

README.fr.md

@@ -318,6 +318,7 @@ docker run --name new-api -d --restart always \
 | `SQL_DSN` | Chaine de connexion à la base de données | - |
 | `REDIS_CONN_STRING` | Chaine de connexion Redis | - |
 | `STREAMING_TIMEOUT` | Délai d'expiration du streaming (secondes) | `300` |
+| `STREAMING_FIRST_BYTE_TIMEOUT` | Délai maximal indépendant avant les en-têtes/le premier octet du flux, sans limiter sa durée (secondes, `0` désactive) | `0` |
 | `STREAM_SCANNER_MAX_BUFFER_MB` | Taille max du buffer par ligne (Mo) pour le scanner SSE ; à augmenter quand les sorties image/base64 sont très volumineuses (ex. images 4K) | `64` |
 | `MAX_REQUEST_BODY_MB` | Taille maximale du corps de requête (Mo, comptée **après décompression** ; évite les requêtes énormes/zip bombs qui saturent la mémoire). Dépassement ⇒ `413` | `32` |
 | `AZURE_DEFAULT_API_VERSION` | Version de l'API Azure | `2025-04-01-preview` |

README.ja.md

@@ -320,6 +320,7 @@ docker run --name new-api -d --restart always \
 | `SQL_DSN** | データベース接続文字列 | - |
 | `REDIS_CONN_STRING` | Redis接続文字列 | - |
 | `STREAMING_TIMEOUT` | ストリーミング応答のタイムアウト時間（秒） | `300` |
+| `STREAMING_FIRST_BYTE_TIMEOUT` | ストリーミング応答ヘッダー/最初のバイトを待つ独立タイムアウト。後続のストリーム時間は制限しない（秒、`0` で無効） | `0` |
 | `STREAM_SCANNER_MAX_BUFFER_MB` | ストリームスキャナの1行あたりバッファ上限（MB）。4K画像など巨大なbase64 `data:` ペイロードを扱う場合は値を増加させてください | `64` |
 | `MAX_REQUEST_BODY_MB` | リクエストボディ最大サイズ（MB、**解凍後**に計測。巨大リクエスト/zip bomb によるメモリ枯渇を防止）。超過時は `413` | `32` |
 | `AZURE_DEFAULT_API_VERSION` | Azure APIバージョン | `2025-04-01-preview` |

README.md

@@ -319,6 +319,7 @@ docker run --name new-api -d --restart always \
 | `REDIS_CONN_STRING` | Redis connection string | - |
 | `RELAY_IDLE_CONN_TIMEOUT` | Idle keep-alive timeout for relay HTTP clients, seconds. Defaults to Go standard library behavior; set `0` to disable | `90` |
 | `STREAMING_TIMEOUT` | Streaming timeout (seconds) | `300` |
+| `STREAMING_FIRST_BYTE_TIMEOUT` | Maximum wait for streaming response headers/first byte without limiting the stream duration (seconds, `0` disables) | `0` |
 | `STREAM_SCANNER_MAX_BUFFER_MB` | Max per-line buffer (MB) for the stream scanner; increase when upstream sends huge image/base64 payloads | `64` |
 | `MAX_REQUEST_BODY_MB` | Max request body size (MB, counted **after decompression**; prevents huge requests/zip bombs from exhausting memory). Exceeding it returns `413` | `32` |
 | `AZURE_DEFAULT_API_VERSION` | Azure API version | `2025-04-01-preview` |

README.zh_CN.md

@@ -318,6 +318,7 @@ docker run --name new-api -d --restart always \
 | `SQL_DSN` | 数据库连接字符串                                                     | - |
 | `REDIS_CONN_STRING` | Redis 连接字符串                                                  | - |
 | `STREAMING_TIMEOUT` | 流式超时时间（秒）                                                    | `300` |
+| `STREAMING_FIRST_BYTE_TIMEOUT` | 等待流式响应头/首字节的独立超时，不限制后续流持续时间（秒，`0` 关闭） | `0` |
 | `STREAM_SCANNER_MAX_BUFFER_MB` | 流式扫描器单行最大缓冲（MB），图像生成等超大 `data:` 片段（如 4K 图片 base64）需适当调大 | `64` |
 | `MAX_REQUEST_BODY_MB` | 请求体最大大小（MB，**解压后**计；防止超大请求/zip bomb 导致内存暴涨），超过将返回 `413` | `32` |
 | `AZURE_DEFAULT_API_VERSION` | Azure API 版本                                                 | `2025-04-01-preview` |

README.zh_TW.md

@@ -318,6 +318,7 @@ docker run --name new-api -d --restart always \
 | `SQL_DSN` | 資料庫連接字符串                                                     | - |
 | `REDIS_CONN_STRING` | Redis 連接字符串                                                  | - |
 | `STREAMING_TIMEOUT` | 流式超時時間（秒）                                                    | `300` |
+| `STREAMING_FIRST_BYTE_TIMEOUT` | 等待串流回應標頭/首位元組的獨立逾時，不限制後續串流時間（秒，`0` 關閉） | `0` |
 | `STREAM_SCANNER_MAX_BUFFER_MB` | 流式掃描器單行最大緩衝（MB），圖像生成等超大 `data:` 片段（如 4K 圖片 base64）需適當調大 | `64` |
 | `MAX_REQUEST_BODY_MB` | 請求體最大大小（MB，**解壓縮後**計；防止超大請求/zip bomb 導致記憶體暴漲），超過將返回 `413` | `32` |
 | `AZURE_DEFAULT_API_VERSION` | Azure API 版本                                                 | `2025-04-01-preview` |

common/constants.go

@@ -183,7 +183,8 @@ var SyncFrequency int // unit is second
 var BatchUpdateEnabled = false
 var BatchUpdateInterval int
 
-var RelayTimeout int // unit is second
+var RelayTimeout int              // unit is second
+var StreamingFirstByteTimeout int // unit is second; 0 disables the independent header wait deadline
 
 var RelayIdleConnTimeout int // unit is second
 var RelayMaxIdleConns int

common/init.go

@@ -28,6 +28,8 @@ func printHelp() {
 	fmt.Println("Usage: newapi [--port <port>] [--log-dir <log directory>] [--version] [--help]")
 }
 
+// InitEnv applies process-level environment overrides before dependent settings
+// are initialized, including relay retry and first-byte timeout controls.
 func InitEnv() {
 	flag.Parse()
 
@@ -107,6 +109,7 @@ func InitEnv() {
 	SyncFrequency = GetEnvOrDefault("SYNC_FREQUENCY", 60)
 	BatchUpdateInterval = GetEnvOrDefault("BATCH_UPDATE_INTERVAL", 5)
 	RelayTimeout = GetEnvOrDefault("RELAY_TIMEOUT", 0)
+	StreamingFirstByteTimeout = GetEnvOrDefault("STREAMING_FIRST_BYTE_TIMEOUT", 0)
 	RelayIdleConnTimeout = GetEnvOrDefault("RELAY_IDLE_CONN_TIMEOUT", 90)
 	RelayMaxIdleConns = GetEnvOrDefault("RELAY_MAX_IDLE_CONNS", 500)
 	RelayMaxIdleConnsPerHost = GetEnvOrDefault("RELAY_MAX_IDLE_CONNS_PER_HOST", 100)

constant/context_key.go

@@ -66,6 +66,9 @@ const (
 	// ContextKeyLanguage stores the user's language preference for i18n
 	ContextKeyLanguage ContextKey = "language"
 	ContextKeyIsStream ContextKey = "is_stream"
+	// ContextKeyRelayPingBytes tracks synthetic SSE keepalive bytes so retry
+	// logic can distinguish them from actual upstream response data.
+	ContextKeyRelayPingBytes ContextKey = "relay_ping_bytes"
 
 	// ContextKeyAuditLogged marks that the current request has already recorded
 	// a manage/operation audit log inside the handler. When set, the admin-audit

controller/channel-billing.go

@@ -136,7 +136,9 @@ func GetClaudeAuthHeader(token string) http.Header {
 	return h
 }
 
-func GetResponseBody(method, url string, channel *model.Channel, headers http.Header) ([]byte, error) {
+// GetResponseBody executes a channel-authenticated request and closes the body
+// on every status path, preserving an earlier request or read error over a close error.
+func GetResponseBody(method, url string, channel *model.Channel, headers http.Header) (body []byte, err error) {
 	req, err := http.NewRequest(method, url, nil)
 	if err != nil {
 		return nil, err
@@ -152,14 +154,15 @@ func GetResponseBody(method, url string, channel *model.Channel, headers http.He
 	if err != nil {
 		return nil, err
 	}
+	defer func() {
+		if closeErr := res.Body.Close(); err == nil && closeErr != nil {
+			err = closeErr
+		}
+	}()
 	if res.StatusCode != http.StatusOK {
 		return nil, fmt.Errorf("status code: %d", res.StatusCode)
 	}
-	body, err := io.ReadAll(res.Body)
-	if err != nil {
-		return nil, err
-	}
-	err = res.Body.Close()
+	body, err = io.ReadAll(res.Body)
 	if err != nil {
 		return nil, err
 	}

controller/channel.go

@@ -1156,6 +1156,8 @@ func equalStringPtr(a, b *string) bool {
 	return *a == *b
 }
 
+// FetchModels proxies a channel model-list request and always releases the
+// upstream response body, including non-200 responses.
 func FetchModels(c *gin.Context) {
 	var req struct {
 		BaseURL string `json:"base_url"`
@@ -1241,6 +1243,7 @@ func FetchModels(c *gin.Context) {
 		})
 		return
 	}
+	defer response.Body.Close()
 	//check status code
 	if response.StatusCode != http.StatusOK {
 		c.JSON(http.StatusInternalServerError, gin.H{
@@ -1249,8 +1252,6 @@ func FetchModels(c *gin.Context) {
 		})
 		return
 	}
-	defer response.Body.Close()
-
 	var result struct {
 		Data []struct {
 			ID string `json:"id"`

controller/login_2fa_gate_test.go

@@ -0,0 +1,224 @@
+package controller
+
+import (
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/QuantumNous/new-api/common"
+	appI18n "github.com/QuantumNous/new-api/i18n"
+	"github.com/QuantumNous/new-api/model"
+	"github.com/gin-contrib/sessions"
+	"github.com/gin-contrib/sessions/cookie"
+	"github.com/gin-gonic/gin"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// TestCompleteLoginWithTwoFAKeepsSessionPending verifies that the shared gate
+// never creates an authenticated session before the second factor succeeds.
+func TestCompleteLoginWithTwoFAKeepsSessionPending(t *testing.T) {
+	require.NoError(t, appI18n.Init())
+	db := setupModelListControllerTestDB(t)
+	require.NoError(t, db.AutoMigrate(&model.TwoFA{}))
+
+	user := &model.User{
+		Username: "external-login-twofa-user",
+		Role:     common.RoleCommonUser,
+		Status:   common.UserStatusEnabled,
+		Group:    "default",
+	}
+	require.NoError(t, db.Create(user).Error)
+	require.NoError(t, db.Create(&model.TwoFA{
+		UserId:    user.Id,
+		Secret:    "test-secret",
+		IsEnabled: true,
+	}).Error)
+
+	router := gin.New()
+	store := cookie.NewStore([]byte("test-session-secret"))
+	router.Use(sessions.Sessions("session", store))
+	router.GET("/login", func(c *gin.Context) {
+		completeLoginWithTwoFA(user, c)
+	})
+	router.GET("/state", func(c *gin.Context) {
+		session := sessions.Default(c)
+		c.JSON(http.StatusOK, gin.H{
+			"authenticated_user_id": session.Get("id"),
+			"pending_user_id":       session.Get("pending_user_id"),
+		})
+	})
+
+	loginRecorder := httptest.NewRecorder()
+	router.ServeHTTP(loginRecorder, httptest.NewRequest(http.MethodGet, "/login", nil))
+	require.Equal(t, http.StatusOK, loginRecorder.Code)
+
+	var loginResponse struct {
+		Success bool `json:"success"`
+		Data    struct {
+			RequireTwoFA bool `json:"require_2fa"`
+		} `json:"data"`
+	}
+	require.NoError(t, common.Unmarshal(loginRecorder.Body.Bytes(), &loginResponse))
+	assert.True(t, loginResponse.Success)
+	assert.True(t, loginResponse.Data.RequireTwoFA)
+
+	stateRequest := httptest.NewRequest(http.MethodGet, "/state", nil)
+	for _, responseCookie := range loginRecorder.Result().Cookies() {
+		stateRequest.AddCookie(responseCookie)
+	}
+	stateRecorder := httptest.NewRecorder()
+	router.ServeHTTP(stateRecorder, stateRequest)
+
+	var state struct {
+		AuthenticatedUserId *int `json:"authenticated_user_id"`
+		PendingUserId       int  `json:"pending_user_id"`
+	}
+	require.NoError(t, common.Unmarshal(stateRecorder.Body.Bytes(), &state))
+	assert.Nil(t, state.AuthenticatedUserId)
+	assert.Equal(t, user.Id, state.PendingUserId)
+}
+
+// TestCompleteLoginWithTwoFARejectsDisabledUser covers disabled accounts with
+// and without 2FA so external login methods cannot bypass the status check.
+func TestCompleteLoginWithTwoFARejectsDisabledUser(t *testing.T) {
+	require.NoError(t, appI18n.Init())
+
+	for _, twoFAEnabled := range []bool{false, true} {
+		t.Run(fmt.Sprintf("twofa_enabled_%t", twoFAEnabled), func(t *testing.T) {
+			db := setupModelListControllerTestDB(t)
+			require.NoError(t, db.AutoMigrate(&model.TwoFA{}))
+
+			user := &model.User{
+				Username: "disabled-external-login-user",
+				Role:     common.RoleCommonUser,
+				Status:   common.UserStatusDisabled,
+				Group:    "default",
+			}
+			require.NoError(t, db.Create(user).Error)
+			if twoFAEnabled {
+				require.NoError(t, db.Create(&model.TwoFA{
+					UserId:    user.Id,
+					Secret:    "test-secret",
+					IsEnabled: true,
+				}).Error)
+			}
+
+			router := gin.New()
+			store := cookie.NewStore([]byte("test-session-secret"))
+			router.Use(sessions.Sessions("session", store))
+			router.GET("/login", func(c *gin.Context) {
+				completeLoginWithTwoFA(user, c)
+			})
+			router.GET("/state", func(c *gin.Context) {
+				session := sessions.Default(c)
+				c.JSON(http.StatusOK, gin.H{
+					"authenticated_user_id": session.Get("id"),
+					"pending_user_id":       session.Get("pending_user_id"),
+				})
+			})
+
+			loginRequest := httptest.NewRequest(http.MethodGet, "/login", nil)
+			loginRequest.Header.Set("Accept-Language", appI18n.LangEn)
+			loginRecorder := httptest.NewRecorder()
+			router.ServeHTTP(loginRecorder, loginRequest)
+
+			var loginResponse struct {
+				Success bool   `json:"success"`
+				Message string `json:"message"`
+			}
+			require.NoError(t, common.Unmarshal(loginRecorder.Body.Bytes(), &loginResponse))
+			assert.False(t, loginResponse.Success)
+			assert.Equal(t, appI18n.Translate(appI18n.LangEn, appI18n.MsgOAuthUserBanned), loginResponse.Message)
+
+			stateRequest := httptest.NewRequest(http.MethodGet, "/state", nil)
+			for _, responseCookie := range loginRecorder.Result().Cookies() {
+				stateRequest.AddCookie(responseCookie)
+			}
+			stateRecorder := httptest.NewRecorder()
+			router.ServeHTTP(stateRecorder, stateRequest)
+
+			var state struct {
+				AuthenticatedUserId *int `json:"authenticated_user_id"`
+				PendingUserId       *int `json:"pending_user_id"`
+			}
+			require.NoError(t, common.Unmarshal(stateRecorder.Body.Bytes(), &state))
+			assert.Nil(t, state.AuthenticatedUserId)
+			assert.Nil(t, state.PendingUserId)
+		})
+	}
+}
+
+// TestSetupLoginRejectsUserDisabledWhileTwoFAIsPending protects the race where
+// an administrator disables an account between primary and second-factor login.
+func TestSetupLoginRejectsUserDisabledWhileTwoFAIsPending(t *testing.T) {
+	require.NoError(t, appI18n.Init())
+	db := setupModelListControllerTestDB(t)
+	require.NoError(t, db.AutoMigrate(&model.TwoFA{}))
+
+	user := &model.User{
+		Username: "disabled-during-twofa-user",
+		Role:     common.RoleCommonUser,
+		Status:   common.UserStatusEnabled,
+		Group:    "default",
+	}
+	require.NoError(t, db.Create(user).Error)
+	require.NoError(t, db.Create(&model.TwoFA{
+		UserId:    user.Id,
+		Secret:    "test-secret",
+		IsEnabled: true,
+	}).Error)
+
+	router := gin.New()
+	store := cookie.NewStore([]byte("test-session-secret"))
+	router.Use(sessions.Sessions("session", store))
+	router.GET("/login", func(c *gin.Context) {
+		completeLoginWithTwoFA(user, c)
+	})
+	router.GET("/finish", func(c *gin.Context) {
+		setupLogin(user, c)
+	})
+	router.GET("/state", func(c *gin.Context) {
+		session := sessions.Default(c)
+		c.JSON(http.StatusOK, gin.H{"authenticated_user_id": session.Get("id")})
+	})
+
+	loginRecorder := httptest.NewRecorder()
+	router.ServeHTTP(loginRecorder, httptest.NewRequest(http.MethodGet, "/login", nil))
+	require.Equal(t, http.StatusOK, loginRecorder.Code)
+
+	user.Status = common.UserStatusDisabled
+	require.NoError(t, db.Model(user).Update("status", user.Status).Error)
+	finishRequest := httptest.NewRequest(http.MethodGet, "/finish", nil)
+	finishRequest.Header.Set("Accept-Language", appI18n.LangEn)
+	for _, responseCookie := range loginRecorder.Result().Cookies() {
+		finishRequest.AddCookie(responseCookie)
+	}
+	finishRecorder := httptest.NewRecorder()
+	router.ServeHTTP(finishRecorder, finishRequest)
+
+	var finishResponse struct {
+		Success bool   `json:"success"`
+		Message string `json:"message"`
+	}
+	require.NoError(t, common.Unmarshal(finishRecorder.Body.Bytes(), &finishResponse))
+	assert.False(t, finishResponse.Success)
+	assert.Equal(t, appI18n.Translate(appI18n.LangEn, appI18n.MsgOAuthUserBanned), finishResponse.Message)
+
+	stateRequest := httptest.NewRequest(http.MethodGet, "/state", nil)
+	for _, responseCookie := range loginRecorder.Result().Cookies() {
+		stateRequest.AddCookie(responseCookie)
+	}
+	for _, responseCookie := range finishRecorder.Result().Cookies() {
+		stateRequest.AddCookie(responseCookie)
+	}
+	stateRecorder := httptest.NewRecorder()
+	router.ServeHTTP(stateRecorder, stateRequest)
+
+	var state struct {
+		AuthenticatedUserId *int `json:"authenticated_user_id"`
+	}
+	require.NoError(t, common.Unmarshal(stateRecorder.Body.Bytes(), &state))
+	assert.Nil(t, state.AuthenticatedUserId)
+}

controller/oauth.go

@@ -130,8 +130,8 @@ func HandleOAuth(c *gin.Context) {
 		return
 	}
 
-	// 9. Setup login
-	setupLogin(user, c)
+	// 9. Apply the same 2FA gate used by password login.
+	completeLoginWithTwoFA(user, c)
 }
 
 // handleOAuthBind handles binding OAuth account to existing user