From e9c01836f41b61c8a083196dcb11af5f36c5217f Mon Sep 17 00:00:00 2001 From: Delicious233 <101502465+DeliciousBuding@users.noreply.github.com> Date: Sun, 18 Oct 2026 15:25:57 +0800 Subject: [PATCH 01/22] fix(security): charge request memory by actual reads - Stop pre-reserving declared Content-Length: idle sockets could pin the global budget without sending a body. - Add a 10 minute ReadTimeout for request-body reads. - Do not pre-read/cache large import uploads before authentication; charge compressed import bodies to the same request-memory budget. - Refresh the CSP inline-script hash and add a drift test. --- api/middleware.go | 8 ++++++ api/middleware_test.go | 22 +++++++++++++++ main.go | 7 +++-- security/csp_test.go | 47 +++++++++++++++++++++++++++++++++ security/decompress.go | 24 +++++++++++++++-- security/middleware.go | 13 ++++++--- security/request_memory_test.go | 34 ++++++++++++++++++++++++ 7 files changed, 148 insertions(+), 7 deletions(-) create mode 100644 security/csp_test.go diff --git a/api/middleware.go b/api/middleware.go index bd77aac97..2d49d6d76 100644 --- a/api/middleware.go +++ b/api/middleware.go @@ -14,6 +14,8 @@ import ( "time" "github.com/gin-gonic/gin" + + "github.com/codex2api/security" ) // Version represents an API version @@ -181,6 +183,12 @@ func randInt() int { // BodyCacheMiddleware caches the request body for multiple reads func BodyCacheMiddleware() gin.HandlerFunc { return func(c *gin.Context) { + // 大批量导入端点由 handler 自己流式解析 multipart（>32MB 落临时文件）。 + // 这里不能预先 io.ReadAll：鉴权还没发生，200MB 的未认证请求会直接常驻内存。 + if security.IsLargeImportUploadPath(c.Request.URL.Path) { + c.Next() + return + } if body := GetRawBody(c); body != nil { c.Request.Body = io.NopCloser(bytes.NewReader(body)) c.Next() diff --git a/api/middleware_test.go b/api/middleware_test.go index 53bcb44a1..95f42d2d6 100644 --- a/api/middleware_test.go +++ b/api/middleware_test.go @@ -39,3 +39,25 @@ func TestCORSMiddlewareAllowsCodex2APIAffinityHeaderPreflight(t *testing.T) { t.Fatalf("Access-Control-Allow-Headers = %q, missing X-Codex2API-Affinity-Key", recorder.Header().Get("Access-Control-Allow-Headers")) } } + +func TestBodyCacheSkipsLargeImportUpload(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.Use(BodyCacheMiddleware()) + rawBodyCached := false + router.POST("/api/admin/accounts/import", func(c *gin.Context) { + _, rawBodyCached = c.Get("raw_body") + c.Status(http.StatusOK) + }) + + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodPost, "/api/admin/accounts/import", strings.NewReader("not read before auth")) + router.ServeHTTP(recorder, request) + + if recorder.Code != http.StatusOK { + t.Fatalf("status = %d", recorder.Code) + } + if rawBodyCached { + t.Fatal("large import body must not be cached/read before authentication") + } +} diff --git a/main.go b/main.go index 998042a78..ccd24d0ce 100644 --- a/main.go +++ b/main.go @@ -195,7 +195,7 @@ func main() { } } else { log.Printf("已加载持久化业务设置: ProxyURL=%s, MaxConcurrency=%d, GlobalRPM=%d, PgMaxConns=%d, RedisPoolSize=%d", - settings.ProxyURL, settings.MaxConcurrency, settings.GlobalRPM, settings.PgMaxConns, settings.RedisPoolSize) + security.MaskURLCredentials(settings.ProxyURL), settings.MaxConcurrency, settings.GlobalRPM, settings.PgMaxConns, settings.RedisPoolSize) } modelCooldownCtx, modelCooldownCancel := context.WithTimeout(context.Background(), 3*time.Second) modelCooldownSettings, modelCooldownErr := db.GetModelCooldownSettings(modelCooldownCtx) @@ -655,9 +655,12 @@ func main() { srv := &http.Server{ Addr: addr, Handler: r, - // ReadHeaderTimeout 防 Slowloris 慢速攻击；IdleTimeout 回收空闲 keep-alive 连接。 + // ReadHeaderTimeout 防 Slowloris 慢速攻击；ReadTimeout 限制整个请求体读取， + // 防止“声明大 Content-Length 但不发 body”长期占用连接/内存预算； + // IdleTimeout 回收空闲 keep-alive 连接。 // 注意：WriteTimeout 必须保持 0 —— LLM 流式响应可持续数分钟，任何固定写超时都会中途切断长回答。 ReadHeaderTimeout: 30 * time.Second, + ReadTimeout: 10 * time.Minute, IdleTimeout: 120 * time.Second, } go func() { diff --git a/security/csp_test.go b/security/csp_test.go new file mode 100644 index 000000000..6677c45a8 --- /dev/null +++ b/security/csp_test.go @@ -0,0 +1,47 @@ +package security + +import ( + "bytes" + "crypto/sha256" + "encoding/base64" + "os" + "regexp" + "strings" + "testing" +) + +// TestContentSecurityPolicyCoversInlineScripts guards against the CSP hash +// drifting from frontend/index.html (the theme bootstrap script is inline). +func TestContentSecurityPolicyCoversInlineScripts(t *testing.T) { + html, err := os.ReadFile("../frontend/index.html") + if err != nil { + t.Fatalf("read frontend/index.html: %v", err) + } + // The Linux build and Docker context use LF; normalize CRLF checkouts so the + // test computes the same hash as the served asset. + html = bytes.ReplaceAll(html, []byte("\r\n"), []byte("\n")) + re := regexp.MustCompile(`(?is)]*)>(.*?)`) + matches := re.FindAllSubmatch(html, -1) + if len(matches) == 0 { + t.Fatal("no inline scripts found in frontend/index.html") + } + inline := 0 + for i, match := range matches { + if strings.Contains(strings.ToLower(string(match[1])), "src=") { + continue + } + script := strings.TrimSpace(string(match[2])) + if script == "" { + continue + } + inline++ + sum := sha256.Sum256(match[2]) + hash := "sha256-" + base64.StdEncoding.EncodeToString(sum[:]) + if !strings.Contains(contentSecurityPolicy, hash) { + t.Errorf("inline script %d hash %s missing from CSP", i, hash) + } + } + if inline == 0 { + t.Fatal("no inline script content found in frontend/index.html") + } +} diff --git a/security/decompress.go b/security/decompress.go index 6f457e910..2cf71de18 100644 --- a/security/decompress.go +++ b/security/decompress.go @@ -125,13 +125,33 @@ func RequestBodyDecompressor(maxSize int64) gin.HandlerFunc { if cached, exists := c.Get("raw_body"); exists { body, _ = cached.([]byte) } + mem := requestMemoryFromContext(c) if body == nil { if c.Request.Body == nil || c.Request.Body == http.NoBody { c.Next() return } - data, err := io.ReadAll(io.LimitReader(c.Request.Body, maxSize+1)) + var data []byte + var err error + if IsLargeImportUploadPath(c.Request.URL.Path) { + // 导入端点在鉴权前不缓存 body；解压时也必须计入请求内存预算， + // 不能让未认证请求用压缩体绕过 BodyCache 的修复。 + reservation, admitted := TryAcquireRequestMemory(0) + if !admitted { + rejectRequestMemory(c) + return + } + defer reservation.Release() + mem = reservation + data, err = readRequestMemoryBounded(c.Request.Body, maxSize, reservation, 0) + } else { + data, err = io.ReadAll(io.LimitReader(c.Request.Body, maxSize+1)) + } if err != nil { + if errors.Is(err, ErrRequestMemoryBudget) { + rejectRequestMemory(c) + return + } c.JSON(http.StatusBadRequest, gin.H{ "error": gin.H{ "message": "读取请求体失败", @@ -152,7 +172,7 @@ func RequestBodyDecompressor(maxSize int64) gin.HandlerFunc { // 多个编码按施加顺序列出，解码逆序进行。 for i := len(encodings) - 1; i >= 0; i-- { var err error - decoded, err = decodeContentEncodingWithMemory(decoded, encodings[i], maxSize, requestMemoryFromContext(c)) + decoded, err = decodeContentEncodingWithMemory(decoded, encodings[i], maxSize, mem) if err != nil { if errors.Is(err, ErrRequestMemoryBudget) { rejectRequestMemory(c) diff --git a/security/middleware.go b/security/middleware.go index 25737838f..c7d81a7a0 100644 --- a/security/middleware.go +++ b/security/middleware.go @@ -15,6 +15,10 @@ import ( "github.com/gin-gonic/gin" ) +// contentSecurityPolicy is kept in a constant so a unit test can verify that +// every inline script in frontend/index.html is covered by its hash. +const contentSecurityPolicy = "default-src 'self'; script-src 'self' 'sha256-8jI3h1oN0D7cJ3eKd/Mjk+hKunZV0rZrmg9g4u/5+Nk='; style-src 'self' 'unsafe-inline'; img-src 'self' data: blob: https: http:; media-src 'self' data: blob: https: http:; font-src 'self'; connect-src 'self' https://api.github.com;" + // Security headers middleware func SecurityHeadersMiddleware() gin.HandlerFunc { return func(c *gin.Context) { @@ -26,8 +30,8 @@ func SecurityHeadersMiddleware() gin.HandlerFunc { // Content Security Policy // script-src 内联哈希对应 frontend/index.html 中的主题恢复脚本（在 React 挂载前同步 localStorage 主题到 ，避免刷新闪烁）。 - // 如修改该脚本，请同步更新此处哈希（浏览器控制台报错会给出新哈希）。 - c.Header("Content-Security-Policy", "default-src 'self'; script-src 'self' 'sha256-WBkfMfl2yQ2pGMyfjGOlZU/dKJRErT3nin/t7BPNjn0='; style-src 'self' 'unsafe-inline'; img-src 'self' data: blob: https: http:; media-src 'self' data: blob: https: http:; font-src 'self'; connect-src 'self' https://api.github.com;") + // 如修改该脚本，请同步更新 contentSecurityPolicy（security/csp_test.go 会检查）。 + c.Header("Content-Security-Policy", contentSecurityPolicy) // Strict Transport Security (HTTPS only) // c.Header("Strict-Transport-Security", "max-age=31536000; includeSubDomains; preload") @@ -80,7 +84,10 @@ func requestSizeLimiterWithMemory(maxSize int64, budget *requestMemoryBudget) gi }}) return } - prepaid := max(c.Request.ContentLength, 0) + // Charge the budget as bytes actually arrive. Reserving the declared + // Content-Length up front lets a client declare 47MB, send nothing, and + // pin the global budget; an empty body must reserve zero. + prepaid := int64(0) reservation, admitted := budget.acquire(prepaid) if !admitted { rejectRequestMemory(c) diff --git a/security/request_memory_test.go b/security/request_memory_test.go index fa0fe3cfe..e5a4e51c7 100644 --- a/security/request_memory_test.go +++ b/security/request_memory_test.go @@ -3,11 +3,13 @@ package security import ( "bytes" "compress/gzip" + "io" "net/http" "net/http/httptest" "strings" "sync" "testing" + "time" "github.com/gin-gonic/gin" ) @@ -85,6 +87,38 @@ func TestRequestMemorySingleRequestSizeStillReturns413(t *testing.T) { } } +func TestRequestMemoryDoesNotPrepayDeclaredLength(t *testing.T) { + budget := &requestMemoryBudget{RequestMemorySnapshot: RequestMemorySnapshot{LimitBytes: 64}} + router := gin.New() + router.Use(requestSizeLimiterWithMemory(128, budget)) + router.POST("/", func(c *gin.Context) { c.Status(200) }) + + pr, pw := io.Pipe() + req := httptest.NewRequest("POST", "/", pr) + req.ContentLength = 48 // declared, but send nothing yet + done := make(chan struct{}) + go func() { + defer close(done) + router.ServeHTTP(httptest.NewRecorder(), req) + }() + + time.Sleep(100 * time.Millisecond) + if budget.UsedBytes != 0 { + t.Fatalf("declared Content-Length was prepaid: used=%d", budget.UsedBytes) + } + + _, _ = pw.Write([]byte("x")) + _ = pw.Close() + select { + case <-done: + case <-time.After(3 * time.Second): + t.Fatal("request did not finish") + } + if budget.UsedBytes != 0 || budget.ActiveReservations != 0 { + t.Fatalf("reservation leaked: %+v", budget.RequestMemorySnapshot) + } +} + func TestRequestMemoryResizeAndConcurrentRelease(t *testing.T) { budget := &requestMemoryBudget{RequestMemorySnapshot: RequestMemorySnapshot{LimitBytes: 100}} r, ok := budget.acquire(80) From 3e604646ad3dc6e9eaff8637eeafa8897d431199 Mon Sep 17 00:00:00 2001 From: Delicious233 <101502465+DeliciousBuding@users.noreply.github.com> Date: Sun, 18 Oct 2026 15:26:08 +0800 Subject: [PATCH 02/22] fix(proxy): bound parse-cache memory and oversized input arrays - Request parse cache now has a 64MB byte budget, a 1MB per-entry cap and a 10 minute TTL. Previously 256 distinct 10-48MB bodies could stay resident indefinitely. - Reject Responses input and Chat messages arrays above 20k items before unmarshalling into maps; a 24MB body of tiny objects could otherwise expand into hundreds of MB. --- proxy/handler.go | 30 +++++++++ proxy/request_limits_test.go | 28 ++++++++ proxy/translator.go | 115 ++++++++++++++++++++++++++------- proxy/translator_cache_test.go | 55 ++++++++++++++++ 4 files changed, 205 insertions(+), 23 deletions(-) create mode 100644 proxy/request_limits_test.go create mode 100644 proxy/translator_cache_test.go diff --git a/proxy/handler.go b/proxy/handler.go index a0c7acb1a..9af204eb8 100644 --- a/proxy/handler.go +++ b/proxy/handler.go @@ -1786,6 +1786,30 @@ func readRawRequestBody(c *gin.Context) ([]byte, error) { return body, nil } +const ( + maxResponsesInputItems = 20000 + maxChatCompletionMessages = 20000 +) + +// rejectExcessiveJSONArray bounds item counts before the body is unmarshaled +// into maps. A 24MB body containing hundreds of thousands of tiny objects can +// expand into hundreds of MB during parsing; the request memory budget only +// counts raw bytes. +func rejectExcessiveJSONArray(c *gin.Context, rawBody []byte, path string, max int) bool { + if c == nil || len(rawBody) == 0 || max <= 0 { + return false + } + if count := gjson.GetBytes(rawBody, path+".#").Int(); count > int64(max) { + api.SendError(c, api.NewAPIError( + api.ErrCodeInvalidRequest, + fmt.Sprintf("%s contains too many items (max %d)", path, max), + api.ErrorTypeInvalidRequest, + )) + return true + } + return false +} + const ingressRequestBodyContextKey = "ingress_raw_body" func setIngressRequestBodyIfAbsent(c *gin.Context, body []byte) { @@ -3866,6 +3890,9 @@ func (h *Handler) Responses(c *gin.Context) { api.SendError(c, api.NewAPIError(api.ErrCodeInvalidRequest, "Failed to read request body", api.ErrorTypeInvalidRequest)) return } + if rejectExcessiveJSONArray(c, rawBody, "input", maxResponsesInputItems) { + return + } if contextErr := responsesHTTPContextEnvelopeError(rawBody); contextErr != nil { api.SendError(c, contextErr) return @@ -6815,6 +6842,9 @@ func (h *Handler) ChatCompletions(c *gin.Context) { api.SendError(c, api.NewAPIError(api.ErrCodeInvalidRequest, "Failed to read request body", api.ErrorTypeInvalidRequest)) return } + if rejectExcessiveJSONArray(c, rawBody, "messages", maxChatCompletionMessages) { + return + } h.capturePromptRequestIngress(c, rawBody) supportedModels := h.supportedModelIDs(c.Request.Context()) diff --git a/proxy/request_limits_test.go b/proxy/request_limits_test.go new file mode 100644 index 000000000..56e928407 --- /dev/null +++ b/proxy/request_limits_test.go @@ -0,0 +1,28 @@ +package proxy + +import ( + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" +) + +func TestRejectExcessiveJSONArray(t *testing.T) { + gin.SetMode(gin.TestMode) + body := []byte(`{"input":[{"a":1},{"a":2},{"a":3}]}`) + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + if !rejectExcessiveJSONArray(c, body, "input", 2) { + t.Fatal("expected rejection when item count exceeds limit") + } + if recorder.Code != 400 { + t.Fatalf("status = %d, want 400", recorder.Code) + } + + recorder = httptest.NewRecorder() + c, _ = gin.CreateTestContext(recorder) + if rejectExcessiveJSONArray(c, body, "input", 3) { + t.Fatal("unexpected rejection at the limit") + } +} diff --git a/proxy/translator.go b/proxy/translator.go index 9c23d16c5..ffb0f33b4 100644 --- a/proxy/translator.go +++ b/proxy/translator.go @@ -189,7 +189,12 @@ type openAIErrorResponse struct { // ==================== LRU 请求解析缓存 ==================== -const requestCacheSize = 256 +const ( + requestCacheSize = 256 + requestCacheMaxBytes = 64 << 20 + requestCacheMaxEntryBytes = 1 << 20 + requestCacheTTL = 10 * time.Minute +) // maxTools 上游 Codex API 允许的最大工具数量 const maxTools = 128 @@ -215,19 +220,28 @@ var responsesImageGenerationUnsupportedOptionFields = []string{ } type requestCacheEntry struct { - key [32]byte - req openAIRequest + key [32]byte + req openAIRequest + size int + expiresAt time.Time } type requestCache struct { - mu sync.Mutex - order *list.List - items map[[32]byte]*list.Element + mu sync.Mutex + order *list.List + items map[[32]byte]*list.Element + bytes int + maxBytes int + maxEntryBytes int + ttl time.Duration } var globalRequestCache = &requestCache{ - order: list.New(), - items: make(map[[32]byte]*list.Element, requestCacheSize), + order: list.New(), + items: make(map[[32]byte]*list.Element, requestCacheSize), + maxBytes: requestCacheMaxBytes, + maxEntryBytes: requestCacheMaxEntryBytes, + ttl: requestCacheTTL, } func firstNonEmptyAnyString(raw any) string { @@ -1696,6 +1710,46 @@ func truncateToolsPreservingImageGeneration(tools []any) []any { return truncated } +func (c *requestCache) ttlValue() time.Duration { + if c.ttl > 0 { + return c.ttl + } + return requestCacheTTL +} + +func (c *requestCache) maxBytesValue() int { + if c.maxBytes > 0 { + return c.maxBytes + } + return requestCacheMaxBytes +} + +func (c *requestCache) maxEntryBytesValue() int { + if c.maxEntryBytes > 0 { + return c.maxEntryBytes + } + return requestCacheMaxEntryBytes +} + +func (c *requestCache) removeElementLocked(elem *list.Element) { + entry := elem.Value.(*requestCacheEntry) + c.order.Remove(elem) + delete(c.items, entry.key) + c.bytes -= entry.size + if c.bytes < 0 { + c.bytes = 0 + } +} + +func (c *requestCache) pruneExpiredLocked(now time.Time) { + for _, elem := range c.items { + entry := elem.Value.(*requestCacheEntry) + if !entry.expiresAt.IsZero() && now.After(entry.expiresAt) { + c.removeElementLocked(elem) + } + } +} + func (c *requestCache) get(key [32]byte) (openAIRequest, bool) { c.mu.Lock() defer c.mu.Unlock() @@ -1703,29 +1757,44 @@ func (c *requestCache) get(key [32]byte) (openAIRequest, bool) { if !ok { return openAIRequest{}, false } + entry := elem.Value.(*requestCacheEntry) + if !entry.expiresAt.IsZero() && time.Now().After(entry.expiresAt) { + c.removeElementLocked(elem) + return openAIRequest{}, false + } c.order.MoveToFront(elem) - return elem.Value.(*requestCacheEntry).req, true + return entry.req, true } -func (c *requestCache) put(key [32]byte, req openAIRequest) { +func (c *requestCache) put(key [32]byte, req openAIRequest, size int) { + if size <= 0 || size > c.maxEntryBytesValue() { + return + } c.mu.Lock() defer c.mu.Unlock() + now := time.Now() + c.pruneExpiredLocked(now) + expiresAt := now.Add(c.ttlValue()) if elem, ok := c.items[key]; ok { + entry := elem.Value.(*requestCacheEntry) + c.bytes -= entry.size + entry.req = req + entry.size = size + entry.expiresAt = expiresAt + c.bytes += size c.order.MoveToFront(elem) - elem.Value.(*requestCacheEntry).req = req - return - } - elem := c.order.PushFront(&requestCacheEntry{key: key, req: req}) - c.items[key] = elem - if c.order.Len() <= requestCacheSize { - return + } else { + elem := c.order.PushFront(&requestCacheEntry{key: key, req: req, size: size, expiresAt: expiresAt}) + c.items[key] = elem + c.bytes += size } - tail := c.order.Back() - if tail == nil { - return + for c.order.Len() > requestCacheSize || c.bytes > c.maxBytesValue() { + tail := c.order.Back() + if tail == nil { + break + } + c.removeElementLocked(tail) } - c.order.Remove(tail) - delete(c.items, tail.Value.(*requestCacheEntry).key) } // cachedOrParse 从缓存获取或解析请求，返回结构体（Unmarshal 至多一次） @@ -1739,7 +1808,7 @@ func cachedOrParse(rawJSON []byte) openAIRequest { } var req openAIRequest _ = json.Unmarshal(rawJSON, &req) - globalRequestCache.put(key, req) + globalRequestCache.put(key, req, len(rawJSON)) return req } diff --git a/proxy/translator_cache_test.go b/proxy/translator_cache_test.go new file mode 100644 index 000000000..478da96d8 --- /dev/null +++ b/proxy/translator_cache_test.go @@ -0,0 +1,55 @@ +package proxy + +import ( + "container/list" + "testing" + "time" +) + +func newTestRequestCache(maxBytes, maxEntryBytes int, ttl time.Duration) *requestCache { + return &requestCache{ + order: list.New(), + items: make(map[[32]byte]*list.Element), + maxBytes: maxBytes, + maxEntryBytes: maxEntryBytes, + ttl: ttl, + } +} + +func TestRequestCacheEnforcesByteBudget(t *testing.T) { + c := newTestRequestCache(10, 10, time.Minute) + c.put([32]byte{1}, openAIRequest{}, 6) + c.put([32]byte{2}, openAIRequest{}, 6) + if len(c.items) != 1 { + t.Fatalf("entries = %d, want 1 after byte-budget eviction", len(c.items)) + } + if c.bytes > 10 { + t.Fatalf("bytes = %d, want <= 10", c.bytes) + } + if _, ok := c.get([32]byte{1}); ok { + t.Fatal("oldest entry should have been evicted") + } + if _, ok := c.get([32]byte{2}); !ok { + t.Fatal("newest entry should remain") + } +} + +func TestRequestCacheRejectsOversizedEntry(t *testing.T) { + c := newTestRequestCache(100, 8, time.Minute) + c.put([32]byte{1}, openAIRequest{}, 9) + if len(c.items) != 0 { + t.Fatal("entry larger than maxEntryBytes must not be cached") + } +} + +func TestRequestCacheExpiresEntries(t *testing.T) { + c := newTestRequestCache(100, 100, time.Nanosecond) + c.put([32]byte{1}, openAIRequest{}, 1) + time.Sleep(time.Millisecond) + if _, ok := c.get([32]byte{1}); ok { + t.Fatal("expired entry must not be returned") + } + if len(c.items) != 0 { + t.Fatal("expired entry should be removed on access") + } +} From 0482f9e1d1285bcc03220620d2de2055768549e3 Mon Sep 17 00:00:00 2001 From: Delicious233 <101502465+DeliciousBuding@users.noreply.github.com> Date: Sun, 18 Oct 2026 15:26:13 +0800 Subject: [PATCH 03/22] fix(images): cap external input images and trim the 202 receipt - Enforce the existing 32MB queue input budget on the legacy external job path too. 16 x 20MB remote images previously expanded to ~426MB of base64 in memory, a 447MB DB row and a 447MB 202 response. - Do not echo params_json (including all input image data URLs) in the 202 creation receipt; callers can read the job detail endpoint. --- admin/image_studio_external.go | 15 ++++++++++++++- admin/image_studio_external_budget_test.go | 20 ++++++++++++++++++++ admin/image_studio_test.go | 18 ++++++++++++++++-- 3 files changed, 50 insertions(+), 3 deletions(-) create mode 100644 admin/image_studio_external_budget_test.go diff --git a/admin/image_studio_external.go b/admin/image_studio_external.go index d99c017bb..e80b95c9f 100644 --- a/admin/image_studio_external.go +++ b/admin/image_studio_external.go @@ -168,7 +168,11 @@ func (h *Handler) CreateExternalImageJob(c *gin.Context) { } h.runImageGenerationJob(jobID, req, apiKey, opts) }() - c.JSON(http.StatusAccepted, externalImageJobResponse{Job: job}) + // 202 只返回小回执；params_json 可能包含全部输入图（数十 MB）， + // 不能随响应回显给客户端，否则单请求即可放大内存/带宽。 + responseJob := *job + responseJob.ParamsJSON = "" + c.JSON(http.StatusAccepted, externalImageJobResponse{Job: &responseJob}) } func (h *Handler) GetExternalImageJob(c *gin.Context) { @@ -258,15 +262,24 @@ func normalizeExternalImageJobFields(req *imageGenerationJobPayload) (bool, erro } func normalizeExternalImageJobInputs(ctx context.Context, req *imageGenerationJobPayload) error { + return normalizeExternalImageJobInputsWithLimit(ctx, req, queueInputBudget) +} + +func normalizeExternalImageJobInputsWithLimit(ctx context.Context, req *imageGenerationJobPayload, limit int64) error { if req == nil { return fmt.Errorf("body is required") } images := make([]string, 0, len(req.InputImages)) + total := 0 for _, imageURL := range req.InputImages { normalizedImage, err := normalizeExternalInputImage(ctx, imageURL) if err != nil { return err } + total += len(normalizedImage) + if int64(total) > limit { + return fmt.Errorf("input_images total size exceeds %d bytes", limit) + } images = append(images, normalizedImage) } req.InputImages = images diff --git a/admin/image_studio_external_budget_test.go b/admin/image_studio_external_budget_test.go new file mode 100644 index 000000000..7edcf0558 --- /dev/null +++ b/admin/image_studio_external_budget_test.go @@ -0,0 +1,20 @@ +package admin + +import ( + "bytes" + "context" + "encoding/base64" + "testing" +) + +func TestNormalizeExternalImageJobInputsTotalBudget(t *testing.T) { + img := "data:image/png;base64," + base64.StdEncoding.EncodeToString(bytes.Repeat([]byte("a"), 40)) + req := &imageGenerationJobPayload{InputImages: []string{img, img}} + if err := normalizeExternalImageJobInputsWithLimit(context.Background(), req, 150); err == nil { + t.Fatal("expected total input budget error") + } + req = &imageGenerationJobPayload{InputImages: []string{img, img}} + if err := normalizeExternalImageJobInputsWithLimit(context.Background(), req, 300); err != nil { + t.Fatalf("unexpected error under budget: %v", err) + } +} diff --git a/admin/image_studio_test.go b/admin/image_studio_test.go index 5b3c10142..43fdeb715 100644 --- a/admin/image_studio_test.go +++ b/admin/image_studio_test.go @@ -999,8 +999,15 @@ func TestExternalImageJobRouteNormalizesEditPayload(t *testing.T) { if created.Job == nil || created.Job.APIKeyID != keyID { t.Fatalf("created job = %#v, want api key id %d", created.Job, keyID) } + if created.Job.ParamsJSON != "" { + t.Fatal("202 receipt must not echo params_json/input images") + } + storedJob, err := db.GetImageGenerationJob(context.Background(), created.Job.ID) + if err != nil { + t.Fatalf("load stored job: %v", err) + } var params imageGenerationJobPayload - if err := json.Unmarshal([]byte(created.Job.ParamsJSON), ¶ms); err != nil { + if err := json.Unmarshal([]byte(storedJob.ParamsJSON), ¶ms); err != nil { t.Fatalf("decode params_json: %v", err) } if params.Prompt != "edit this" || params.Model != "gpt-image-2" || params.OutputFormat != "png" { @@ -1092,8 +1099,15 @@ func TestExternalImageJobRouteFetchesPublicInputImageAsDataURL(t *testing.T) { if created.Job == nil { t.Fatalf("created job is nil") } + if created.Job.ParamsJSON != "" { + t.Fatal("202 receipt must not echo params_json/input images") + } + storedJob, err := db.GetImageGenerationJob(context.Background(), created.Job.ID) + if err != nil { + t.Fatalf("load stored job: %v", err) + } var params imageGenerationJobPayload - if err := json.Unmarshal([]byte(created.Job.ParamsJSON), ¶ms); err != nil { + if err := json.Unmarshal([]byte(storedJob.ParamsJSON), ¶ms); err != nil { t.Fatalf("decode params_json: %v", err) } want := "data:image/png;base64," + base64.StdEncoding.EncodeToString(tiny) From 1e84e020d2670537afb428e5a4aedb9c8a7308e6 Mon Sep 17 00:00:00 2001 From: Delicious233 <101502465+DeliciousBuding@users.noreply.github.com> Date: Sun, 18 Oct 2026 15:26:18 +0800 Subject: [PATCH 04/22] fix(proxy): cap concurrent downstream WebSocket connections Responses/Realtime WebSocket handlers had no per-key or global connection limit and no idle deadline after the first frame. A single API key could keep hundreds of silent connections open, each holding a goroutine and fd. Reserve a slot before the upgrade (64 per key, 4096 globally) and release it when the handler exits. --- proxy/realtime_ws.go | 5 ++ proxy/responses_ws.go | 5 ++ proxy/ws_connection_limit.go | 91 +++++++++++++++++++++++++++++++ proxy/ws_connection_limit_test.go | 28 ++++++++++ 4 files changed, 129 insertions(+) create mode 100644 proxy/ws_connection_limit.go create mode 100644 proxy/ws_connection_limit_test.go diff --git a/proxy/realtime_ws.go b/proxy/realtime_ws.go index d9553bc6b..673ff1317 100644 --- a/proxy/realtime_ws.go +++ b/proxy/realtime_ws.go @@ -47,6 +47,11 @@ func (h *Handler) RealtimeWebSocket(c *gin.Context) { ), http.StatusUpgradeRequired) return } + releaseWS, ok := h.acquireDownstreamWebSocket(c) + if !ok { + return + } + defer releaseWS() if h != nil && h.store != nil { cfg := h.promptFilterConfigForRequest(c) if h.rejectRequiredNewAPIIdentity(c, cfg.Advanced.NewAPI, nil) { diff --git a/proxy/responses_ws.go b/proxy/responses_ws.go index 66314bdf0..e1be46e83 100644 --- a/proxy/responses_ws.go +++ b/proxy/responses_ws.go @@ -252,6 +252,11 @@ func (h *Handler) ResponsesWebSocket(c *gin.Context) { ), http.StatusUpgradeRequired) return } + releaseWS, ok := h.acquireDownstreamWebSocket(c) + if !ok { + return + } + defer releaseWS() if h != nil && h.store != nil { cfg := h.promptFilterConfigForRequest(c) if h.rejectRequiredNewAPIIdentity(c, cfg.Advanced.NewAPI, nil) { diff --git a/proxy/ws_connection_limit.go b/proxy/ws_connection_limit.go new file mode 100644 index 000000000..c15002ba8 --- /dev/null +++ b/proxy/ws_connection_limit.go @@ -0,0 +1,91 @@ +package proxy + +import ( + "net/http" + "sync" + + "github.com/gin-gonic/gin" + + "github.com/codex2api/api" +) + +const ( + defaultWSMaxConnectionsPerKey = 64 + defaultWSMaxConnections = 4096 +) + +// wsConnectionLimiter bounds downstream WebSocket connections per API key and +// globally. Connections are cheap to open but expensive to keep: each one holds +// a goroutine, file descriptor and read buffer for as long as the client stays +// silent. +type wsConnectionLimiter struct { + mu sync.Mutex + maxTotal int + maxPerKey int + total int + perKey map[int64]int +} + +func newWSConnectionLimiter(maxTotal, maxPerKey int) *wsConnectionLimiter { + if maxTotal <= 0 { + maxTotal = defaultWSMaxConnections + } + if maxPerKey <= 0 { + maxPerKey = defaultWSMaxConnectionsPerKey + } + return &wsConnectionLimiter{ + maxTotal: maxTotal, + maxPerKey: maxPerKey, + perKey: make(map[int64]int), + } +} + +var globalWSConnectionLimiter = newWSConnectionLimiter(defaultWSMaxConnections, defaultWSMaxConnectionsPerKey) + +// acquire registers one connection. The returned release is idempotent. +func (l *wsConnectionLimiter) acquire(keyID int64) (func(), bool) { + l.mu.Lock() + defer l.mu.Unlock() + if l.total >= l.maxTotal { + return nil, false + } + if keyID > 0 && l.perKey[keyID] >= l.maxPerKey { + return nil, false + } + l.total++ + if keyID > 0 { + l.perKey[keyID]++ + } + released := false + return func() { + l.mu.Lock() + defer l.mu.Unlock() + if released { + return + } + released = true + l.total-- + if keyID > 0 { + l.perKey[keyID]-- + if l.perKey[keyID] <= 0 { + delete(l.perKey, keyID) + } + } + }, true +} + +// acquireDownstreamWebSocket reserves a slot before the HTTP connection is +// hijacked. On rejection the caller must not upgrade. +func (h *Handler) acquireDownstreamWebSocket(c *gin.Context) (func(), bool) { + release, ok := globalWSConnectionLimiter.acquire(requestAPIKeyID(c)) + if !ok { + c.Header("Retry-After", "5") + api.SendErrorWithStatus(c, api.NewAPIError( + api.ErrCodeServiceUnavailable, + "too many concurrent WebSocket connections", + api.ErrorTypeServer, + ), http.StatusServiceUnavailable) + return nil, false + } + return release, true +} diff --git a/proxy/ws_connection_limit_test.go b/proxy/ws_connection_limit_test.go new file mode 100644 index 000000000..141614a60 --- /dev/null +++ b/proxy/ws_connection_limit_test.go @@ -0,0 +1,28 @@ +package proxy + +import "testing" + +func TestWSConnectionLimiterPerKeyAndGlobal(t *testing.T) { + l := newWSConnectionLimiter(3, 2) + r1, ok := l.acquire(1) + if !ok { + t.Fatal("first per-key connection rejected") + } + if _, ok := l.acquire(1); !ok { + t.Fatal("second per-key connection rejected") + } + if _, ok := l.acquire(1); ok { + t.Fatal("third per-key connection should be rejected") + } + if _, ok := l.acquire(2); !ok { + t.Fatal("other key should still be admitted") + } + if _, ok := l.acquire(3); ok { + t.Fatal("global limit should reject") + } + r1() + r1() // idempotent + if _, ok := l.acquire(1); !ok { + t.Fatal("released slot should be reusable") + } +} From 587bcf7f132a44f768d9a3b808dccb8a191a1d89 Mon Sep 17 00:00:00 2001 From: Delicious233 <101502465+DeliciousBuding@users.noreply.github.com> Date: Sun, 18 Oct 2026 15:26:23 +0800 Subject: [PATCH 05/22] fix(auth): keep self-service OAuth callbacks in pending review The public /auth/callback looked up the session by state only and always called upsertOAuthIdentityAccount, which inserts an enabled account. A self-service session could therefore skip /submit-code and put an unreviewed account straight into the scheduling pool. Branch on sess.SelfService and use upsertSelfServiceAccount (enabled=false + self-service tag), matching the normal portal submission path. --- admin/oauth.go | 45 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 45 insertions(+) diff --git a/admin/oauth.go b/admin/oauth.go index 4b06e054a..9fd35d950 100644 --- a/admin/oauth.go +++ b/admin/oauth.go @@ -22,6 +22,7 @@ import ( "github.com/codex2api/auth" "github.com/codex2api/database" "github.com/codex2api/proxy" + "github.com/codex2api/security" "github.com/gin-gonic/gin" ) @@ -777,6 +778,50 @@ func (h *Handler) OAuthCallback(c *gin.Context) { expiresIn: tokenResp.ExpiresIn, }) + // 公开自助门户发起的会话必须走「待审核」入库，不能因为 state 被带到公开 + // /auth/callback 就绕过 SubmitAccountPortalCode 的 enabled=false 分支。 + if sess.SelfService { + ctx, cancel := context.WithTimeout(c.Request.Context(), 30*time.Second) + defer cancel() + name := seed.email + if name == "" { + name = "self-service-account" + } + id, err := h.upsertSelfServiceAccount(ctx, name, proxyURL, seed, sess.ContactEmail) + if err == errDuplicateOAuthIdentity { + globalOAuthStore.delete(sessionID) + sess.ExchangeResult = &oauthExchangeResult{ + Success: true, + Message: "该账号已在系统中，无需重复提交", + } + c.String(http.StatusOK, oauthCallbackPage("提交成功", "该账号已在系统中，无需重复提交。", true)) + return + } + if err != nil { + sess.ExchangeResult = &oauthExchangeResult{ + Success: false, + Error: "账号写入数据库失败: " + err.Error(), + } + c.String(http.StatusOK, oauthCallbackPage("授权失败", "写入数据库失败: "+err.Error(), false)) + return + } + globalOAuthStore.delete(sessionID) + email := seed.email + if accountInfo != nil && accountInfo.Email != "" { + email = accountInfo.Email + } + security.SecurityAuditLog("ACCOUNT_PORTAL_SUBMITTED", "ip="+security.SanitizeLog(c.ClientIP())+" contact="+security.SanitizeLog(sess.ContactEmail)+" account="+security.SanitizeLog(email)) + sess.ExchangeResult = &oauthExchangeResult{ + Success: true, + Message: "账号已提交，等待管理员审核", + ID: id, + Email: email, + } + log.Printf("自助门户回调新增待审核账号: id=%d email=%s", id, email) + c.String(http.StatusOK, oauthCallbackPage("提交成功", "已提交，等待管理员审核后生效。感谢你的贡献！", true)) + return + } + // 自动添加账号 name := "" if seed.email != "" { From a7fda911762ae7cdb7ddece6ab659716aa14ae7f Mon Sep 17 00:00:00 2001 From: Delicious233 <101502465+DeliciousBuding@users.noreply.github.com> Date: Sun, 18 Oct 2026 15:27:06 +0800 Subject: [PATCH 06/22] fix(security): redact proxy credentials and TLS-wrap https proxies - Mask URL userinfo before logging persisted settings, proxy-pool events and transport fallback errors. Proxy URLs commonly contain user:pass. - The custom http-connect dialers treated https:// proxies as plaintext TCP and sent Proxy-Authorization in the clear. Establish TLS to the proxy (system roots, TLS 1.2+) before sending CONNECT; cover both the general uTLS transport and the auth client. --- admin/handler.go | 4 +-- auth/proxy_pool.go | 4 ++- auth/proxy_pool_integration.go | 11 ++++--- auth/utls_client.go | 28 ++++++++++++++-- auth/utls_proxy_tls_test.go | 60 ++++++++++++++++++++++++++++++++++ proxy/executor.go | 3 +- proxy/utls_proxy_tls_test.go | 60 ++++++++++++++++++++++++++++++++++ proxy/utls_transport.go | 31 +++++++++++++++--- 8 files changed, 185 insertions(+), 16 deletions(-) create mode 100644 auth/utls_proxy_tls_test.go create mode 100644 proxy/utls_proxy_tls_test.go diff --git a/admin/handler.go b/admin/handler.go index e6c5020a0..6455adc49 100644 --- a/admin/handler.go +++ b/admin/handler.go @@ -10955,7 +10955,7 @@ func (h *Handler) UpdateSettings(c *gin.Context) { if req.ProxyURL != nil { h.store.SetProxyURL(*req.ProxyURL) - log.Printf("设置已更新: proxy_url = %s", *req.ProxyURL) + log.Printf("设置已更新: proxy_url = %s", security.MaskURLCredentials(*req.ProxyURL)) } if req.PgMaxConns != nil { @@ -11140,7 +11140,7 @@ func (h *Handler) UpdateSettings(c *gin.Context) { v := strings.TrimSpace(*req.GithubProxyURL) h.store.SetGithubProxyURL(v) runtimeCfg.GithubProxyURL = v - log.Printf("设置已更新: github_proxy_url = %s", v) + log.Printf("设置已更新: github_proxy_url = %s", security.MaskURLCredentials(v)) } // Codex 过载熔断（配置只存 RuntimeSettings，热更新生效） diff --git a/auth/proxy_pool.go b/auth/proxy_pool.go index 7b594a0aa..e410d964b 100644 --- a/auth/proxy_pool.go +++ b/auth/proxy_pool.go @@ -11,6 +11,8 @@ import ( "sync" "sync/atomic" "time" + + "github.com/codex2api/security" ) // ProxyHealthStatus 代理健康状态 @@ -917,7 +919,7 @@ func (p *ProxyPool) RecoverIsolatedProxies() { // 在锁外执行回调，避免死锁 for _, entry := range recoveredEntries { - log.Printf("[ProxyPool] 代理已恢复: %s", entry.URL) + log.Printf("[ProxyPool] 代理已恢复: %s", security.MaskURLCredentials(entry.URL)) if p.onRecovery != nil { p.onRecovery(entry) } diff --git a/auth/proxy_pool_integration.go b/auth/proxy_pool_integration.go index 8524c2888..4888ec4bf 100644 --- a/auth/proxy_pool_integration.go +++ b/auth/proxy_pool_integration.go @@ -8,6 +8,7 @@ import ( "time" "github.com/codex2api/database" + "github.com/codex2api/security" ) // EnhancedProxyPool 增强的代理池管理器 @@ -54,16 +55,16 @@ func (e *EnhancedProxyPool) Init(ctx context.Context) error { // 设置回调 e.pool.SetOnIsolation(func(entry *ProxyEntry) { - log.Printf("[ProxyPool] 代理被隔离: %s (连续失败 %d 次)", entry.URL, entry.ConsecutiveFailures) + log.Printf("[ProxyPool] 代理被隔离: %s (连续失败 %d 次)", security.MaskURLCredentials(entry.URL), entry.ConsecutiveFailures) }) e.pool.SetOnRecovery(func(entry *ProxyEntry) { - log.Printf("[ProxyPool] 代理已恢复: %s", entry.URL) + log.Printf("[ProxyPool] 代理已恢复: %s", security.MaskURLCredentials(entry.URL)) }) e.pool.SetOnHealthCheck(func(result *HealthCheckResult) { if !result.Healthy && result.Error != nil { - log.Printf("[ProxyPool] 代理健康检查失败: %s, error: %v", result.URL, result.Error) + log.Printf("[ProxyPool] 代理健康检查失败: %s, error: %v", security.MaskURLCredentials(result.URL), result.Error) } }) @@ -169,7 +170,7 @@ func (e *EnhancedProxyPool) ReloadFromDB(ctx context.Context) error { } } e.pool.AddProxy(p.URL, weight) - log.Printf("[ProxyPool] 新增代理: %s", p.URL) + log.Printf("[ProxyPool] 新增代理: %s", security.MaskURLCredentials(p.URL)) } } @@ -181,7 +182,7 @@ func (e *EnhancedProxyPool) ReloadFromDB(ctx context.Context) error { for _, url := range currentProxies { if !newMap[url] { e.pool.RemoveProxy(url) - log.Printf("[ProxyPool] 移除代理: %s", url) + log.Printf("[ProxyPool] 移除代理: %s", security.MaskURLCredentials(url)) } } diff --git a/auth/utls_client.go b/auth/utls_client.go index 0c60aaf3f..aef048b74 100644 --- a/auth/utls_client.go +++ b/auth/utls_client.go @@ -3,6 +3,7 @@ package auth import ( "bufio" "context" + "crypto/tls" "encoding/base64" "fmt" "net" @@ -193,8 +194,10 @@ func buildAuthProxyDialer(proxyURL string) (xproxy.Dialer, error) { } type authHTTPConnectDialer struct { - proxyAddr string - authHeader string + proxyAddr string + proxyServerName string + useTLS bool + authHeader string } func (d *authHTTPConnectDialer) Dial(network, addr string) (net.Conn, error) { @@ -203,6 +206,21 @@ func (d *authHTTPConnectDialer) Dial(network, addr string) (net.Conn, error) { return nil, fmt.Errorf("连接代理服务器失败: %w", err) } + // https:// 代理先与代理完成 TLS，避免 Basic 凭据明文传输。 + if d.useTLS { + tlsConn := tls.Client(conn, &tls.Config{ + ServerName: d.proxyServerName, + MinVersion: tls.VersionTLS12, + }) + _ = tlsConn.SetDeadline(time.Now().Add(10 * time.Second)) + if err := tlsConn.Handshake(); err != nil { + conn.Close() + return nil, fmt.Errorf("代理 TLS 握手失败: %w", err) + } + _ = tlsConn.SetDeadline(time.Time{}) + conn = tlsConn + } + connectReq := fmt.Sprintf("CONNECT %s HTTP/1.1\r\nHost: %s\r\n", addr, addr) if d.authHeader != "" { connectReq += fmt.Sprintf("Proxy-Authorization: %s\r\n", d.authHeader) @@ -252,7 +270,11 @@ func buildAuthHTTPProxyDialer(u *url.URL) (xproxy.Dialer, error) { } } - d := &authHTTPConnectDialer{proxyAddr: addr} + d := &authHTTPConnectDialer{ + proxyAddr: addr, + proxyServerName: u.Hostname(), + useTLS: strings.EqualFold(u.Scheme, "https"), + } if u.User != nil { username := u.User.Username() diff --git a/auth/utls_proxy_tls_test.go b/auth/utls_proxy_tls_test.go new file mode 100644 index 000000000..e353a5ae3 --- /dev/null +++ b/auth/utls_proxy_tls_test.go @@ -0,0 +1,60 @@ +package auth + +import ( + "bytes" + "net" + "net/url" + "testing" + "time" +) + +func TestAuthHTTPProxyDialerUsesTLSForHTTPSProxy(t *testing.T) { + tests := []struct { + name string + scheme string + wantPrefix []byte + }{ + {name: "http proxy sends CONNECT in plaintext", scheme: "http", wantPrefix: []byte("CONNECT ")}, + {name: "https proxy starts with TLS ClientHello", scheme: "https", wantPrefix: []byte{0x16, 0x03}}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer ln.Close() + captured := make(chan []byte, 1) + go func() { + conn, err := ln.Accept() + if err != nil { + return + } + defer conn.Close() + _ = conn.SetReadDeadline(time.Now().Add(2 * time.Second)) + buf := make([]byte, 512) + n, _ := conn.Read(buf) + captured <- append([]byte(nil), buf[:n]...) + }() + + u, err := url.Parse(tc.scheme + "://user:pass@" + ln.Addr().String()) + if err != nil { + t.Fatal(err) + } + dialer, err := buildAuthHTTPProxyDialer(u) + if err != nil { + t.Fatal(err) + } + go func() { _, _ = dialer.Dial("tcp", "example.com:443") }() + + select { + case data := <-captured: + if !bytes.HasPrefix(data, tc.wantPrefix) { + t.Fatalf("first bytes = %q, want prefix %q", data, tc.wantPrefix) + } + case <-time.After(3 * time.Second): + t.Fatal("proxy dialer sent no bytes") + } + }) + } +} diff --git a/proxy/executor.go b/proxy/executor.go index 430492425..e2673e482 100644 --- a/proxy/executor.go +++ b/proxy/executor.go @@ -20,6 +20,7 @@ import ( "github.com/codex2api/auth" "github.com/codex2api/proxy/internal/jsonbody" + "github.com/codex2api/security" "github.com/google/uuid" "github.com/tidwall/gjson" "github.com/tidwall/sjson" @@ -260,7 +261,7 @@ func newCodexStandardTransport(proxyURL string) http.RoundTripper { baseDialer := &net.Dialer{Timeout: 10 * time.Second, KeepAlive: 30 * time.Second} transport.DialContext = baseDialer.DialContext if err := auth.ConfigureTransportProxy(transport, proxyURL, baseDialer); err != nil { - log.Printf("[CodexTransport] 代理配置失败，回退直连: proxy=%s err=%v", proxyURL, err) + log.Printf("[CodexTransport] 代理配置失败，回退直连: proxy=%s err=%v", security.MaskURLCredentials(proxyURL), err) transport.Proxy = nil transport.DialContext = baseDialer.DialContext } diff --git a/proxy/utls_proxy_tls_test.go b/proxy/utls_proxy_tls_test.go new file mode 100644 index 000000000..856873dd7 --- /dev/null +++ b/proxy/utls_proxy_tls_test.go @@ -0,0 +1,60 @@ +package proxy + +import ( + "bytes" + "net" + "net/url" + "testing" + "time" +) + +func TestHTTPProxyDialerUsesTLSForHTTPSProxy(t *testing.T) { + tests := []struct { + name string + scheme string + wantPrefix []byte + }{ + {name: "http proxy sends CONNECT in plaintext", scheme: "http", wantPrefix: []byte("CONNECT ")}, + {name: "https proxy starts with TLS ClientHello", scheme: "https", wantPrefix: []byte{0x16, 0x03}}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer ln.Close() + captured := make(chan []byte, 1) + go func() { + conn, err := ln.Accept() + if err != nil { + return + } + defer conn.Close() + _ = conn.SetReadDeadline(time.Now().Add(2 * time.Second)) + buf := make([]byte, 512) + n, _ := conn.Read(buf) + captured <- append([]byte(nil), buf[:n]...) + }() + + u, err := url.Parse(tc.scheme + "://user:pass@" + ln.Addr().String()) + if err != nil { + t.Fatal(err) + } + dialer, err := buildHTTPProxyDialer(u) + if err != nil { + t.Fatal(err) + } + go func() { _, _ = dialer.Dial("tcp", "example.com:443") }() + + select { + case data := <-captured: + if !bytes.HasPrefix(data, tc.wantPrefix) { + t.Fatalf("first bytes = %q, want prefix %q", data, tc.wantPrefix) + } + case <-time.After(3 * time.Second): + t.Fatal("proxy dialer sent no bytes") + } + }) + } +} diff --git a/proxy/utls_transport.go b/proxy/utls_transport.go index e70c41101..ec28944cc 100644 --- a/proxy/utls_transport.go +++ b/proxy/utls_transport.go @@ -3,6 +3,7 @@ package proxy import ( "bufio" "context" + "crypto/tls" "encoding/base64" "fmt" "log" @@ -104,7 +105,7 @@ func NewUTLSTransport(proxyURL string) http.RoundTripper { if proxyURL != "" { d, err := buildProxyDialer(proxyURL) if err != nil { - log.Printf("[UTLS] 代理配置失败，回退直连: proxy=%s err=%v", proxyURL, err) + log.Printf("[UTLS] 代理配置失败，回退直连: proxy=%s err=%v", security.MaskURLCredentials(proxyURL), err) dialer = xproxy.Direct } else { dialer = d @@ -145,8 +146,10 @@ func buildProxyDialer(proxyURL string) (xproxy.Dialer, error) { // httpConnectDialer 通过 HTTP CONNECT 方法建立隧道的拨号器 type httpConnectDialer struct { - proxyAddr string // 代理服务器地址（host:port） - authHeader string // Proxy-Authorization 头（可选） + proxyAddr string // 代理服务器地址（host:port） + proxyServerName string // https 代理 TLS SNI/证书校验名 + useTLS bool // https:// 代理需要先与代理建立 TLS + authHeader string // Proxy-Authorization 头（可选） } // Dial 通过 HTTP CONNECT 隧道连接到目标地址 @@ -157,6 +160,22 @@ func (d *httpConnectDialer) Dial(network, addr string) (net.Conn, error) { return nil, fmt.Errorf("连接代理服务器失败: %w", err) } + // https:// 代理必须先与代理完成 TLS 握手，否则 Proxy-Authorization + // 会以明文经过网络；证书校验保持开启。 + if d.useTLS { + tlsConn := tls.Client(conn, &tls.Config{ + ServerName: d.proxyServerName, + MinVersion: tls.VersionTLS12, + }) + _ = tlsConn.SetDeadline(time.Now().Add(10 * time.Second)) + if err := tlsConn.Handshake(); err != nil { + conn.Close() + return nil, fmt.Errorf("代理 TLS 握手失败: %w", err) + } + _ = tlsConn.SetDeadline(time.Time{}) + conn = tlsConn + } + // 2. 发送 CONNECT 请求建立隧道 connectReq := fmt.Sprintf("CONNECT %s HTTP/1.1\r\nHost: %s\r\n", addr, addr) if d.authHeader != "" { @@ -211,7 +230,11 @@ func buildHTTPProxyDialer(u *url.URL) (xproxy.Dialer, error) { } } - d := &httpConnectDialer{proxyAddr: addr} + d := &httpConnectDialer{ + proxyAddr: addr, + proxyServerName: u.Hostname(), + useTLS: strings.EqualFold(u.Scheme, "https"), + } // 处理代理认证 if u.User != nil { From 946d3575f5789f17dfab759467241257cd31f520 Mon Sep 17 00:00:00 2001 From: Delicious233 <101502465+DeliciousBuding@users.noreply.github.com> Date: Sun, 18 Oct 2026 15:27:11 +0800 Subject: [PATCH 07/22] chore(deploy): fix log workdir and validate build versions - Set WORKDIR /app in the runtime image so the default LOG_DIR=logs resolves to the compose-mounted /app/logs instead of /logs (audit logs were lost on container recreation). - Validate BUILD_VERSION against ^[0-9A-Za-z._-]{1,64}$ before it is interpolated into Dockerfile RUN shell commands. --- .github/workflows/docker-image.yml | 9 +++++++++ Dockerfile | 4 ++++ 2 files changed, 13 insertions(+) diff --git a/.github/workflows/docker-image.yml b/.github/workflows/docker-image.yml index 9219d629d..0d228b60d 100644 --- a/.github/workflows/docker-image.yml +++ b/.github/workflows/docker-image.yml @@ -30,6 +30,15 @@ jobs: - name: Checkout uses: actions/checkout@v5 + - name: Validate build version + env: + BUILD_VERSION: ${{ inputs.build_version || github.ref_name }} + run: | + if ! printf '%s' "$BUILD_VERSION" | grep -Eq '^[0-9A-Za-z._-]{1,64}$'; then + echo "Invalid BUILD_VERSION: must match ^[0-9A-Za-z._-]{1,64}\$" >&2 + exit 1 + fi + - name: Set up QEMU uses: docker/setup-qemu-action@v3 diff --git a/Dockerfile b/Dockerfile index bfe0fb9b2..bcee6122f 100644 --- a/Dockerfile +++ b/Dockerfile @@ -48,6 +48,10 @@ RUN apk --no-cache add ca-certificates tzdata COPY --from=go-builder /codex2api /usr/local/bin/codex2api +# 默认 LOG_DIR=logs 是相对路径；固定工作目录才能落到 compose 挂载的 /app/logs， +# 否则审计日志会写到容器根目录的 /logs，容器重建后丢失。 +WORKDIR /app + EXPOSE 8080 ENTRYPOINT ["/usr/local/bin/codex2api"] From 9554e2565c96dd67d8ffacf5866d05b7c5af3bca Mon Sep 17 00:00:00 2001 From: Delicious233 <101502465+DeliciousBuding@users.noreply.github.com> Date: Sun, 18 Oct 2026 15:40:22 +0800 Subject: [PATCH 08/22] test(security): use locked budget snapshot in prepay test The new test read budget counters while the request goroutine was still running; the race detector caught it. Read through a mutex-protected snapshot instead. --- security/request_memory.go | 9 +++++++++ security/request_memory_test.go | 8 ++++---- 2 files changed, 13 insertions(+), 4 deletions(-) diff --git a/security/request_memory.go b/security/request_memory.go index ff547bc63..fc0005818 100644 --- a/security/request_memory.go +++ b/security/request_memory.go @@ -50,6 +50,15 @@ func GetRequestMemorySnapshot() RequestMemorySnapshot { return processRequestMemory.RequestMemorySnapshot } +func (b *requestMemoryBudget) snapshot() RequestMemorySnapshot { + if b == nil { + return RequestMemorySnapshot{} + } + b.mu.Lock() + defer b.mu.Unlock() + return b.RequestMemorySnapshot +} + // RequestMemoryReservation may be shared by queue message copies. Release is // idempotent; a released reservation cannot be grown again. type RequestMemoryReservation struct { diff --git a/security/request_memory_test.go b/security/request_memory_test.go index e5a4e51c7..db32c5208 100644 --- a/security/request_memory_test.go +++ b/security/request_memory_test.go @@ -103,8 +103,8 @@ func TestRequestMemoryDoesNotPrepayDeclaredLength(t *testing.T) { }() time.Sleep(100 * time.Millisecond) - if budget.UsedBytes != 0 { - t.Fatalf("declared Content-Length was prepaid: used=%d", budget.UsedBytes) + if used := budget.snapshot().UsedBytes; used != 0 { + t.Fatalf("declared Content-Length was prepaid: used=%d", used) } _, _ = pw.Write([]byte("x")) @@ -114,8 +114,8 @@ func TestRequestMemoryDoesNotPrepayDeclaredLength(t *testing.T) { case <-time.After(3 * time.Second): t.Fatal("request did not finish") } - if budget.UsedBytes != 0 || budget.ActiveReservations != 0 { - t.Fatalf("reservation leaked: %+v", budget.RequestMemorySnapshot) + if snap := budget.snapshot(); snap.UsedBytes != 0 || snap.ActiveReservations != 0 { + t.Fatalf("reservation leaked: %+v", snap) } } From f984972b321076072b8b52181a125297e0c4f58d Mon Sep 17 00:00:00 2001 From: Delicious233 <101502465+DeliciousBuding@users.noreply.github.com> Date: Sun, 18 Oct 2026 15:50:56 +0800 Subject: [PATCH 09/22] fix: address security review feedback - Mask both username and password (and malformed userinfo through the last '@') before logging proxy URLs; do not log wrapped proxy errors, which may contain the raw URL. - Reject BUILD_VERSION values containing newlines before they reach the metadata action. - Reject duplicate top-level input/messages fields before counting array items; gjson sees the first array while encoding/json decodes the last. - Charge parsed request-cache entries at 3x the raw body size to account for retained slices/RawMessage copies. - Apply a per-IP limit to anonymous WebSocket connections instead of relying only on the global cap. - Insert self-service accounts disabled from the first write, removing the enabled-then-disabled window if a follow-up update fails. --- .github/workflows/docker-image.yml | 8 ++++-- admin/self_service.go | 7 ++--- database/postgres.go | 19 ++++++++++++ proxy/executor.go | 2 +- proxy/handler.go | 6 ++++ proxy/request_limits.go | 46 ++++++++++++++++++++++++++++++ proxy/request_limits_test.go | 16 +++++++++++ proxy/translator.go | 5 +++- proxy/utls_transport.go | 2 +- proxy/ws_connection_limit.go | 37 ++++++++++++++++++------ proxy/ws_connection_limit_test.go | 27 ++++++++++++++---- security/mask_url_test.go | 24 ++++++++++++++++ security/validator.go | 12 ++++---- security/validator_test.go | 2 +- 14 files changed, 180 insertions(+), 33 deletions(-) create mode 100644 proxy/request_limits.go create mode 100644 security/mask_url_test.go diff --git a/.github/workflows/docker-image.yml b/.github/workflows/docker-image.yml index 0d228b60d..b8584c791 100644 --- a/.github/workflows/docker-image.yml +++ b/.github/workflows/docker-image.yml @@ -34,8 +34,12 @@ jobs: env: BUILD_VERSION: ${{ inputs.build_version || github.ref_name }} run: | - if ! printf '%s' "$BUILD_VERSION" | grep -Eq '^[0-9A-Za-z._-]{1,64}$'; then - echo "Invalid BUILD_VERSION: must match ^[0-9A-Za-z._-]{1,64}\$" >&2 + if [[ "$BUILD_VERSION" == *$'\n'* || "$BUILD_VERSION" == *$'\r'* ]]; then + echo "Invalid BUILD_VERSION: newlines are not allowed" >&2 + exit 1 + fi + if ! [[ "$BUILD_VERSION" =~ ^[0-9A-Za-z._-]{1,64}$ ]]; then + echo "Invalid BUILD_VERSION: must match ^[0-9A-Za-z._-]{1,64}$" >&2 exit 1 fi diff --git a/admin/self_service.go b/admin/self_service.go index b7207db4d..1bfa05621 100644 --- a/admin/self_service.go +++ b/admin/self_service.go @@ -253,14 +253,11 @@ func (h *Handler) upsertSelfServiceAccount(ctx context.Context, name, proxyURL s } } - id, err := h.db.InsertAccountWithCredentials(ctx, name, h.newCodexAccountCredentials(seed), proxyURL) + id, err := h.db.InsertDisabledAccountWithCredentials(ctx, name, h.newCodexAccountCredentials(seed), proxyURL) if err != nil { return 0, err } - // 待审核：禁用 + 备注 + 打标；不调用 Store.AddAccount，故不进调度池。 - if err := h.db.SetAccountEnabled(ctx, id, false); err != nil { - return 0, err - } + // 待审核：插入时即 disabled；备注/标签失败时账号仍保持禁用，不会进调度池。 note := "自助提交联系人: " + contactEmail if err := h.db.UpdateAccountNote(ctx, id, note); err != nil { return 0, err diff --git a/database/postgres.go b/database/postgres.go index 65339143a..9beaa3a9e 100644 --- a/database/postgres.go +++ b/database/postgres.go @@ -8782,6 +8782,25 @@ func (db *DB) InsertAccountWithCredentials(ctx context.Context, name string, cre ) } +// InsertDisabledAccountWithCredentials creates an account in the pending state +// from the first write. Self-service submissions must not briefly exist as an +// enabled account if a follow-up update fails. +func (db *DB) InsertDisabledAccountWithCredentials(ctx context.Context, name string, credentials map[string]interface{}, proxyURL string) (int64, error) { + if credentials == nil { + credentials = map[string]interface{}{} + } + credJSON, err := json.Marshal(encryptSensitiveCredentials(credentials)) + if err != nil { + return 0, err + } + return db.insertAccountRowWithFamily(ctx, + `INSERT INTO accounts (name, credentials, proxy_url, enabled) VALUES ($1, $2, $3, FALSE) RETURNING id`, + `INSERT INTO accounts (name, credentials, proxy_url, enabled) VALUES ($1, $2, $3, FALSE)`, + credentials, + name, credJSON, proxyURL, + ) +} + func (db *DB) InsertOpenAIResponsesAccount(ctx context.Context, name string, credentials map[string]interface{}, proxyURL string) (int64, error) { if credentials == nil { credentials = map[string]interface{}{} diff --git a/proxy/executor.go b/proxy/executor.go index e2673e482..475806efc 100644 --- a/proxy/executor.go +++ b/proxy/executor.go @@ -261,7 +261,7 @@ func newCodexStandardTransport(proxyURL string) http.RoundTripper { baseDialer := &net.Dialer{Timeout: 10 * time.Second, KeepAlive: 30 * time.Second} transport.DialContext = baseDialer.DialContext if err := auth.ConfigureTransportProxy(transport, proxyURL, baseDialer); err != nil { - log.Printf("[CodexTransport] 代理配置失败，回退直连: proxy=%s err=%v", security.MaskURLCredentials(proxyURL), err) + log.Printf("[CodexTransport] 代理配置失败，回退直连: proxy=%s err_type=%T", security.MaskURLCredentials(proxyURL), err) transport.Proxy = nil transport.DialContext = baseDialer.DialContext } diff --git a/proxy/handler.go b/proxy/handler.go index 9af204eb8..676c42419 100644 --- a/proxy/handler.go +++ b/proxy/handler.go @@ -3890,6 +3890,9 @@ func (h *Handler) Responses(c *gin.Context) { api.SendError(c, api.NewAPIError(api.ErrCodeInvalidRequest, "Failed to read request body", api.ErrorTypeInvalidRequest)) return } + if rejectDuplicateJSONField(c, rawBody, "input") { + return + } if rejectExcessiveJSONArray(c, rawBody, "input", maxResponsesInputItems) { return } @@ -6842,6 +6845,9 @@ func (h *Handler) ChatCompletions(c *gin.Context) { api.SendError(c, api.NewAPIError(api.ErrCodeInvalidRequest, "Failed to read request body", api.ErrorTypeInvalidRequest)) return } + if rejectDuplicateJSONField(c, rawBody, "messages") { + return + } if rejectExcessiveJSONArray(c, rawBody, "messages", maxChatCompletionMessages) { return } diff --git a/proxy/request_limits.go b/proxy/request_limits.go new file mode 100644 index 000000000..2ff655b43 --- /dev/null +++ b/proxy/request_limits.go @@ -0,0 +1,46 @@ +package proxy + +import ( + "fmt" + + "github.com/gin-gonic/gin" + "github.com/tidwall/gjson" + + "github.com/codex2api/api" + "github.com/codex2api/proxy/internal/jsonbody" +) + +// rejectDuplicateJSONField rejects duplicate top-level array fields before the +// item-count guard. gjson selects the first duplicate, while encoding/json +// decodes the last one, so a small first array could hide a massive second one. +func rejectDuplicateJSONField(c *gin.Context, rawBody []byte, field string) bool { + if c == nil || len(rawBody) == 0 || !gjson.ValidBytes(rawBody) { + return false + } + root := jsonbody.View(rawBody) + if !root.IsObject() { + return false + } + seen := false + duplicate := false + root.ForEach(func(key, _ gjson.Result) bool { + if key.String() != field { + return true + } + if seen { + duplicate = true + return false + } + seen = true + return true + }) + if !duplicate { + return false + } + api.SendError(c, api.NewAPIError( + api.ErrCodeInvalidRequest, + fmt.Sprintf("duplicate top-level field %q is not allowed", field), + api.ErrorTypeInvalidRequest, + )) + return true +} diff --git a/proxy/request_limits_test.go b/proxy/request_limits_test.go index 56e928407..902b9f467 100644 --- a/proxy/request_limits_test.go +++ b/proxy/request_limits_test.go @@ -7,6 +7,22 @@ import ( "github.com/gin-gonic/gin" ) +func TestRejectDuplicateJSONField(t *testing.T) { + gin.SetMode(gin.TestMode) + body := []byte(`{"messages":[],"messages":[{"a":1}]}`) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + if !rejectDuplicateJSONField(c, body, "messages") { + t.Fatal("expected duplicate top-level field rejection") + } + if recorder.Code != 400 { + t.Fatalf("status = %d, want 400", recorder.Code) + } + if rejectDuplicateJSONField(c, []byte(`{"messages":[]}`), "messages") { + t.Fatal("unexpected rejection for a single field") + } +} + func TestRejectExcessiveJSONArray(t *testing.T) { gin.SetMode(gin.TestMode) body := []byte(`{"input":[{"a":1},{"a":2},{"a":3}]}`) diff --git a/proxy/translator.go b/proxy/translator.go index ffb0f33b4..3cbea99bf 100644 --- a/proxy/translator.go +++ b/proxy/translator.go @@ -1808,7 +1808,10 @@ func cachedOrParse(rawJSON []byte) openAIRequest { } var req openAIRequest _ = json.Unmarshal(rawJSON, &req) - globalRequestCache.put(key, req, len(rawJSON)) + // The parsed request retains more than the raw JSON (message slices, + // RawMessage copies, tool structs). Charge a conservative multiplier so the + // byte budget bounds retained heap, not just wire bytes. + globalRequestCache.put(key, req, len(rawJSON)*3) return req } diff --git a/proxy/utls_transport.go b/proxy/utls_transport.go index ec28944cc..8a6cedcd3 100644 --- a/proxy/utls_transport.go +++ b/proxy/utls_transport.go @@ -105,7 +105,7 @@ func NewUTLSTransport(proxyURL string) http.RoundTripper { if proxyURL != "" { d, err := buildProxyDialer(proxyURL) if err != nil { - log.Printf("[UTLS] 代理配置失败，回退直连: proxy=%s err=%v", security.MaskURLCredentials(proxyURL), err) + log.Printf("[UTLS] 代理配置失败，回退直连: proxy=%s err_type=%T", security.MaskURLCredentials(proxyURL), err) dialer = xproxy.Direct } else { dialer = d diff --git a/proxy/ws_connection_limit.go b/proxy/ws_connection_limit.go index c15002ba8..1778da01d 100644 --- a/proxy/ws_connection_limit.go +++ b/proxy/ws_connection_limit.go @@ -11,19 +11,22 @@ import ( const ( defaultWSMaxConnectionsPerKey = 64 + defaultWSMaxConnectionsPerIP = 16 defaultWSMaxConnections = 4096 ) -// wsConnectionLimiter bounds downstream WebSocket connections per API key and -// globally. Connections are cheap to open but expensive to keep: each one holds -// a goroutine, file descriptor and read buffer for as long as the client stays -// silent. +// wsConnectionLimiter bounds downstream WebSocket connections per API key, +// per client IP for anonymous mode, and globally. Connections are cheap to +// open but expensive to keep: each one holds a goroutine, file descriptor and +// read buffer for as long as the client stays silent. type wsConnectionLimiter struct { mu sync.Mutex maxTotal int maxPerKey int + maxPerIP int total int perKey map[int64]int + perIP map[string]int } func newWSConnectionLimiter(maxTotal, maxPerKey int) *wsConnectionLimiter { @@ -36,26 +39,36 @@ func newWSConnectionLimiter(maxTotal, maxPerKey int) *wsConnectionLimiter { return &wsConnectionLimiter{ maxTotal: maxTotal, maxPerKey: maxPerKey, + maxPerIP: defaultWSMaxConnectionsPerIP, perKey: make(map[int64]int), + perIP: make(map[string]int), } } var globalWSConnectionLimiter = newWSConnectionLimiter(defaultWSMaxConnections, defaultWSMaxConnectionsPerKey) // acquire registers one connection. The returned release is idempotent. -func (l *wsConnectionLimiter) acquire(keyID int64) (func(), bool) { +func (l *wsConnectionLimiter) acquire(keyID int64, clientIP string) (func(), bool) { l.mu.Lock() defer l.mu.Unlock() if l.total >= l.maxTotal { return nil, false } - if keyID > 0 && l.perKey[keyID] >= l.maxPerKey { + useKey := keyID > 0 + useIP := !useKey && clientIP != "" + if useKey && l.perKey[keyID] >= l.maxPerKey { + return nil, false + } + if useIP && l.perIP[clientIP] >= l.maxPerIP { return nil, false } l.total++ - if keyID > 0 { + if useKey { l.perKey[keyID]++ } + if useIP { + l.perIP[clientIP]++ + } released := false return func() { l.mu.Lock() @@ -65,19 +78,25 @@ func (l *wsConnectionLimiter) acquire(keyID int64) (func(), bool) { } released = true l.total-- - if keyID > 0 { + if useKey { l.perKey[keyID]-- if l.perKey[keyID] <= 0 { delete(l.perKey, keyID) } } + if useIP { + l.perIP[clientIP]-- + if l.perIP[clientIP] <= 0 { + delete(l.perIP, clientIP) + } + } }, true } // acquireDownstreamWebSocket reserves a slot before the HTTP connection is // hijacked. On rejection the caller must not upgrade. func (h *Handler) acquireDownstreamWebSocket(c *gin.Context) (func(), bool) { - release, ok := globalWSConnectionLimiter.acquire(requestAPIKeyID(c)) + release, ok := globalWSConnectionLimiter.acquire(requestAPIKeyID(c), c.ClientIP()) if !ok { c.Header("Retry-After", "5") api.SendErrorWithStatus(c, api.NewAPIError( diff --git a/proxy/ws_connection_limit_test.go b/proxy/ws_connection_limit_test.go index 141614a60..44c4a049b 100644 --- a/proxy/ws_connection_limit_test.go +++ b/proxy/ws_connection_limit_test.go @@ -4,25 +4,40 @@ import "testing" func TestWSConnectionLimiterPerKeyAndGlobal(t *testing.T) { l := newWSConnectionLimiter(3, 2) - r1, ok := l.acquire(1) + r1, ok := l.acquire(1, "") if !ok { t.Fatal("first per-key connection rejected") } - if _, ok := l.acquire(1); !ok { + if _, ok := l.acquire(1, ""); !ok { t.Fatal("second per-key connection rejected") } - if _, ok := l.acquire(1); ok { + if _, ok := l.acquire(1, ""); ok { t.Fatal("third per-key connection should be rejected") } - if _, ok := l.acquire(2); !ok { + if _, ok := l.acquire(2, ""); !ok { t.Fatal("other key should still be admitted") } - if _, ok := l.acquire(3); ok { + if _, ok := l.acquire(3, ""); ok { t.Fatal("global limit should reject") } r1() r1() // idempotent - if _, ok := l.acquire(1); !ok { + if _, ok := l.acquire(1, ""); !ok { t.Fatal("released slot should be reusable") } } + +func TestWSConnectionLimiterAnonymousPerIP(t *testing.T) { + l := newWSConnectionLimiter(100, 2) + for i := 0; i < defaultWSMaxConnectionsPerIP; i++ { + if _, ok := l.acquire(0, "203.0.113.7"); !ok { + t.Fatalf("connection %d rejected before per-IP limit", i) + } + } + if _, ok := l.acquire(0, "203.0.113.7"); ok { + t.Fatal("per-IP limit should reject anonymous connection") + } + if _, ok := l.acquire(0, "203.0.113.8"); !ok { + t.Fatal("different IP should be admitted") + } +} diff --git a/security/mask_url_test.go b/security/mask_url_test.go new file mode 100644 index 000000000..f50c6012d --- /dev/null +++ b/security/mask_url_test.go @@ -0,0 +1,24 @@ +package security + +import ( + "strings" + "testing" +) + +func TestMaskURLCredentials(t *testing.T) { + tests := []string{ + "http://user:pass@proxy.example:8080", + "http://user@proxy.example:8080", + "http://user:p@ss@proxy.example:8080", + "http://user:p@ss@%zz", + "prefix http://user:pass@proxy.example:8080 suffix", + } + for _, input := range tests { + got := MaskURLCredentials(input) + for _, secret := range []string{"user", "pass", "p@ss", "ss@"} { + if strings.Contains(got, secret) { + t.Errorf("MaskURLCredentials(%q) = %q, still contains %q", input, got, secret) + } + } + } +} diff --git a/security/validator.go b/security/validator.go index 3d30d84b0..7610f9dfa 100644 --- a/security/validator.go +++ b/security/validator.go @@ -196,7 +196,7 @@ func MaskSensitiveData(input string) string { return result } -var embeddedURLCredentials = regexp.MustCompile(`(?i)([a-z][a-z0-9+.-]*://)[^\s/@]+@`) +var embeddedURLCredentials = regexp.MustCompile(`(?i)([a-z][a-z0-9+.-]*://)[^/\s]*@`) // MaskURLCredentials masks userinfo in URLs, including URLs embedded in labels. func MaskURLCredentials(input string) string { @@ -204,12 +204,10 @@ func MaskURLCredentials(input string) string { if err != nil || parsed.Scheme == "" || parsed.Host == "" || parsed.User == nil { return embeddedURLCredentials.ReplaceAllString(input, "${1}redacted@") } - username := parsed.User.Username() - if _, hasPassword := parsed.User.Password(); hasPassword { - parsed.User = url.UserPassword(username, "redacted") - } else { - parsed.User = url.User(username) - } + // Redact both username and password: username-only proxy credentials and + // malformed userinfo (which may contain further '@' characters) must not + // leak either half. + parsed.User = url.User("redacted") return parsed.String() } diff --git a/security/validator_test.go b/security/validator_test.go index 0c1396ac5..00b702406 100644 --- a/security/validator_test.go +++ b/security/validator_test.go @@ -83,7 +83,7 @@ func TestMaskSensitiveData(t *testing.T) { {"sk-abcdef1234567890abcdef", "sk-****MASKED****"}, {"hello world", "hello world"}, {"token=550e8400-e29b-41d4-a716-446655440000", "token=****UUID-MASKED****"}, - {"http://user:secret@proxy.example:8080", "http://user:redacted@proxy.example:8080"}, + {"http://user:secret@proxy.example:8080", "http://redacted@proxy.example:8080"}, } for _, test := range tests { From 5ff88cd2c62e4a0bd5a1f09c52800002d1fbae2b Mon Sep 17 00:00:00 2001 From: Delicious233 <101502465+DeliciousBuding@users.noreply.github.com> Date: Sun, 18 Oct 2026 15:59:05 +0800 Subject: [PATCH 10/22] fix: address remaining review feedback - Charge parsed message/tool-call slice overhead in the request cache estimate; a 20k-message request is now treated as over the per-entry limit instead of retaining megabytes while reporting only wire bytes. - Roll back a partially initialized self-service account when the pending-review note or tag cannot be written, so a retry is not reported as already submitted. --- admin/self_service.go | 14 ++++++++++++++ admin/self_service_test.go | 22 ++++++++++++++++++++++ proxy/translator.go | 29 ++++++++++++++++++++++++++--- proxy/translator_cache_test.go | 17 +++++++++++++++++ 4 files changed, 79 insertions(+), 3 deletions(-) diff --git a/admin/self_service.go b/admin/self_service.go index 1bfa05621..174d9e332 100644 --- a/admin/self_service.go +++ b/admin/self_service.go @@ -258,11 +258,15 @@ func (h *Handler) upsertSelfServiceAccount(ctx context.Context, name, proxyURL s return 0, err } // 待审核：插入时即 disabled；备注/标签失败时账号仍保持禁用，不会进调度池。 + // 但仍需回滚半成品：否则同一身份重试会命中重复检查，提交者收到"已在系统中" + // 的成功语义，而账号永远缺少联系人/标签等待审核信息。 note := "自助提交联系人: " + contactEmail if err := h.db.UpdateAccountNote(ctx, id, note); err != nil { + h.rollbackSelfServiceAccountInsert(ctx, id) return 0, err } if err := h.db.SetAccountTags(ctx, id, []string{selfServiceTag}); err != nil { + h.rollbackSelfServiceAccountInsert(ctx, id) return 0, err } h.db.InsertAccountEventAsync(id, "added", selfServiceSource) @@ -272,3 +276,13 @@ func (h *Handler) upsertSelfServiceAccount(ctx context.Context, name, proxyURL s log.Printf("自助门户新增待审核账号 %d (%s, 联系人 %s)", id, name, contactEmail) return id, nil } + +// rollbackSelfServiceAccountInsert 软删本次插入但未完成待审核设置的账号， +// 使同一身份可以重新提交。软删失败时账号保持 disabled，仍不会进调度池。 +func (h *Handler) rollbackSelfServiceAccountInsert(ctx context.Context, id int64) { + if err := h.db.SoftDeleteAccount(ctx, id); err != nil { + log.Printf("回滚未完成的自助账号 %d 失败: %v", id, err) + return + } + h.invalidateAccountSnapshotCaches() +} diff --git a/admin/self_service_test.go b/admin/self_service_test.go index 506759571..9cc326654 100644 --- a/admin/self_service_test.go +++ b/admin/self_service_test.go @@ -180,6 +180,28 @@ func TestUpsertSelfServiceAccountPendingAndNotScheduled(t *testing.T) { } } +func TestSelfServiceRollbackAllowsRetry(t *testing.T) { + handler, db := newSelfServiceHandler(t, true) + ctx := context.Background() + seed := tokenCredentialSeed{ + refreshToken: "rt-rb", + accessToken: "at-rb", + email: "rollback@example.com", + accountID: "acct-rb", + workspaceID: "workspace-rb", + } + + id, err := db.InsertDisabledAccountWithCredentials(ctx, "rollback@example.com", handler.newCodexAccountCredentials(seed), "") + if err != nil { + t.Fatalf("InsertDisabledAccountWithCredentials: %v", err) + } + handler.rollbackSelfServiceAccountInsert(ctx, id) + + if _, err := handler.upsertSelfServiceAccount(ctx, "rollback@example.com", "", seed, "c@x.com"); err != nil { + t.Fatalf("回滚后同一身份应可重新提交, got %v", err) + } +} + func TestApproveSelfServiceAccountLoadsIntoPool(t *testing.T) { gin.SetMode(gin.TestMode) handler, db := newSelfServiceHandler(t, true) diff --git a/proxy/translator.go b/proxy/translator.go index 3cbea99bf..989cdff91 100644 --- a/proxy/translator.go +++ b/proxy/translator.go @@ -1797,6 +1797,28 @@ func (c *requestCache) put(key [32]byte, req openAIRequest, size int) { } } +// 解析后常驻堆的开销估算：除原始 JSON 外，还需计入切片本身。 +// 20000 条短消息的 JSON 只有约 60KB，但 []openAIMessage 的后备数组就占 ~1.5MB。 +const ( + parsedRequestSizeCeiling = requestCacheMaxEntryBytes + 1 + parsedMessageOverhead = 128 // openAIMessage 结构体 + 切片增长余量 + parsedToolCallOverhead = 128 // openAIToolCall 结构体 + 切片增长余量 + parsedRawJSONOverhead = 32 // json.RawMessage 头部 +) + +func estimateParsedRequestSize(rawJSON []byte, req *openAIRequest) int { + size := int64(len(rawJSON)) * 3 + size += int64(cap(req.Messages)) * parsedMessageOverhead + for i := range req.Messages { + size += int64(cap(req.Messages[i].ToolCalls)) * parsedToolCallOverhead + } + size += int64(cap(req.Tools)) * parsedRawJSONOverhead + if size > parsedRequestSizeCeiling { + return parsedRequestSizeCeiling + } + return int(size) +} + // cachedOrParse 从缓存获取或解析请求，返回结构体（Unmarshal 至多一次） func cachedOrParse(rawJSON []byte) openAIRequest { if len(rawJSON) == 0 { @@ -1809,9 +1831,10 @@ func cachedOrParse(rawJSON []byte) openAIRequest { var req openAIRequest _ = json.Unmarshal(rawJSON, &req) // The parsed request retains more than the raw JSON (message slices, - // RawMessage copies, tool structs). Charge a conservative multiplier so the - // byte budget bounds retained heap, not just wire bytes. - globalRequestCache.put(key, req, len(rawJSON)*3) + // RawMessage copies, tool structs). Charge a conservative multiplier plus + // per-slice overhead so the byte budget bounds retained heap, not just wire + // bytes. + globalRequestCache.put(key, req, estimateParsedRequestSize(rawJSON, &req)) return req } diff --git a/proxy/translator_cache_test.go b/proxy/translator_cache_test.go index 478da96d8..41dd1c2c1 100644 --- a/proxy/translator_cache_test.go +++ b/proxy/translator_cache_test.go @@ -53,3 +53,20 @@ func TestRequestCacheExpiresEntries(t *testing.T) { t.Fatal("expired entry should be removed on access") } } + +func TestEstimateParsedRequestSizeCountsMessageSlices(t *testing.T) { + raw := []byte(`{"messages":[{},{},{}]}`) // ~24 bytes + // 20000 条短消息的原始 JSON 很小，但切片常驻内存远超单条上限：必须按不可缓存处理。 + sparse := openAIRequest{Messages: make([]openAIMessage, 20000)} + if got := estimateParsedRequestSize(raw, &sparse); got != parsedRequestSizeCeiling { + t.Fatalf("20k-message estimate = %d, want clamped to %d", got, parsedRequestSizeCeiling) + } + small := openAIRequest{Messages: make([]openAIMessage, 4)} + want := len(raw)*3 + 4*parsedMessageOverhead + if got := estimateParsedRequestSize(raw, &small); got != want { + t.Fatalf("estimate = %d, want %d", got, want) + } + if got := estimateParsedRequestSize(raw, &openAIRequest{}); got != len(raw)*3 { + t.Fatalf("empty request estimate = %d, want %d", got, len(raw)*3) + } +} From 11a9a36d23e3a6b833a52eb9e9aa9109a2c9f672 Mon Sep 17 00:00:00 2001 From: Delicious233 <101502465+DeliciousBuding@users.noreply.github.com> Date: Sun, 18 Oct 2026 16:37:26 +0800 Subject: [PATCH 11/22] fix: close quota burst and upstream memory gaps - Add per-key in-process admission counters so concurrent bursts cannot overshoot RPM/RPD while the 60s aggregate cache is stale. Verified in WSL: 20 concurrent requests with rpm=2 produced 2 accepted jobs and 18 429s. - Throttle repeated admin auth failures per IP and hash both secrets before the constant-time compare; successful logins still work during a block. - Bound upstream error bodies, SSE lines and decompressed WebSocket messages; cap turn-state provenance entries and pagination offsets. - Reject non-image MIME types for inline serving, mark admin responses no-store, strip newlines in audit logs, extend secret masking, cap strict image dimensions, harden .dockerignore and fix audit-gate on Windows. - Use a background cleanup context for self-service rollback. --- .dockerignore | 8 ++ admin/accounts_paged.go | 3 + admin/auth_throttle.go | 103 +++++++++++++++++ admin/auth_throttle_test.go | 96 ++++++++++++++++ admin/handler.go | 46 +++++++- admin/image_mime_test.go | 26 +++++ admin/image_studio.go | 39 +++++-- admin/key_usage_public.go | 3 + admin/pagination_test.go | 28 +++++ admin/self_service.go | 9 +- admin/self_service_test.go | 2 +- frontend/scripts/audit-gate.mjs | 5 +- internal/imageupscale/imageupscale.go | 11 ++ internal/imageupscale/imageupscale_test.go | 17 +++ proxy/antigravity_stream_errors.go | 2 +- proxy/apikey_admission.go | 128 +++++++++++++++++++++ proxy/apikey_admission_test.go | 79 +++++++++++++ proxy/apikey_limits.go | 70 ++++++----- proxy/codex_turn_state.go | 24 +++- proxy/codex_turn_state_cap_test.go | 40 +++++++ proxy/executor.go | 18 +-- proxy/executor_body_limit_test.go | 48 ++++++++ proxy/grok_namespace_tools.go | 2 +- proxy/grok_protocol.go | 2 +- proxy/handler.go | 5 + proxy/handler_gemini.go | 4 +- proxy/openai_responses_ws.go | 16 ++- proxy/openai_responses_ws_limit_test.go | 64 +++++++++++ proxy/responses_ws.go | 2 +- proxy/sse_line.go | 27 +++++ proxy/sse_line_test.go | 26 +++++ security/middleware.go | 18 +-- security/middleware_test.go | 12 ++ security/validator.go | 7 +- security/validator_test.go | 5 + 35 files changed, 924 insertions(+), 71 deletions(-) create mode 100644 admin/auth_throttle.go create mode 100644 admin/auth_throttle_test.go create mode 100644 admin/image_mime_test.go create mode 100644 admin/pagination_test.go create mode 100644 proxy/apikey_admission.go create mode 100644 proxy/apikey_admission_test.go create mode 100644 proxy/codex_turn_state_cap_test.go create mode 100644 proxy/executor_body_limit_test.go create mode 100644 proxy/openai_responses_ws_limit_test.go create mode 100644 proxy/sse_line.go create mode 100644 proxy/sse_line_test.go diff --git a/.dockerignore b/.dockerignore index 084c85721..7aa702f26 100644 --- a/.dockerignore +++ b/.dockerignore @@ -19,6 +19,14 @@ Thumbs.db # Env .env .env.local +.env.* +!.env.example + +# Logs and local worktrees must not enter the build context. +logs/ +*.log +.worktrees/ +.local/ # Local credentials and account-test fixtures must never enter a build context. .private-account-fixtures/ diff --git a/admin/accounts_paged.go b/admin/accounts_paged.go index 5b6e91cf0..aed37e0c5 100644 --- a/admin/accounts_paged.go +++ b/admin/accounts_paged.go @@ -300,6 +300,9 @@ func parseAccountPageQuery(c *gin.Context) (accountPageQuery, error) { if err != nil || value < 1 { return query, fmt.Errorf("page must be a positive integer") } + if value > maxAdminPage { + return query, fmt.Errorf("page is too large") + } query.Page = value } if raw := strings.TrimSpace(c.Query("page_size")); raw != "" { diff --git a/admin/auth_throttle.go b/admin/auth_throttle.go new file mode 100644 index 000000000..74b8dc828 --- /dev/null +++ b/admin/auth_throttle.go @@ -0,0 +1,103 @@ +package admin + +import ( + "sync" + "time" +) + +const ( + adminAuthFailureWindow = 10 * time.Minute + adminAuthBlockDuration = 10 * time.Minute + adminAuthMaxFailures = 5 + adminAuthMaxEntries = 4096 +) + +// adminAuthThrottle 记录按 IP 的管理端鉴权失败次数。 +// 只在失败时计数，成功请求不受影响；窗口内失败达到阈值后阻断一段时间。 +type adminAuthThrottle struct { + mu sync.Mutex + entries map[string]*adminAuthThrottleEntry +} + +type adminAuthThrottleEntry struct { + failures int + windowStart time.Time + blockedUntil time.Time + lastSeen time.Time +} + +func newAdminAuthThrottle() *adminAuthThrottle { + return &adminAuthThrottle{entries: make(map[string]*adminAuthThrottleEntry)} +} + +// retryAfter 返回该 IP 剩余的阻断时间；未阻断时返回 0。 +func (t *adminAuthThrottle) retryAfter(ip string, now time.Time) time.Duration { + if t == nil { + return 0 + } + t.mu.Lock() + defer t.mu.Unlock() + entry := t.entries[ip] + if entry == nil || !now.Before(entry.blockedUntil) { + return 0 + } + return entry.blockedUntil.Sub(now) +} + +// recordFailure 记录一次失败；达到阈值时开始阻断并返回阻断时长。 +func (t *adminAuthThrottle) recordFailure(ip string, now time.Time) time.Duration { + if t == nil { + return 0 + } + t.mu.Lock() + defer t.mu.Unlock() + entry := t.entries[ip] + if entry == nil || now.Sub(entry.windowStart) > adminAuthFailureWindow { + entry = &adminAuthThrottleEntry{windowStart: now} + t.entries[ip] = entry + } + entry.lastSeen = now + if now.Before(entry.blockedUntil) { + return entry.blockedUntil.Sub(now) + } + entry.failures++ + if entry.failures < adminAuthMaxFailures { + return 0 + } + entry.failures = 0 + entry.windowStart = now + entry.blockedUntil = now.Add(adminAuthBlockDuration) + t.evictLocked(now) + return adminAuthBlockDuration +} + +func (t *adminAuthThrottle) reset(ip string) { + if t == nil { + return + } + t.mu.Lock() + delete(t.entries, ip) + t.mu.Unlock() +} + +// evictLocked 在条目过多时清理过期记录，避免伪造来源地址把内存撑大。 +func (t *adminAuthThrottle) evictLocked(now time.Time) { + if len(t.entries) <= adminAuthMaxEntries { + return + } + for ip, entry := range t.entries { + if !now.Before(entry.blockedUntil) && now.Sub(entry.lastSeen) > adminAuthFailureWindow { + delete(t.entries, ip) + } + } + for len(t.entries) > adminAuthMaxEntries { + var oldestIP string + var oldest time.Time + for ip, entry := range t.entries { + if oldestIP == "" || entry.lastSeen.Before(oldest) { + oldestIP, oldest = ip, entry.lastSeen + } + } + delete(t.entries, oldestIP) + } +} diff --git a/admin/auth_throttle_test.go b/admin/auth_throttle_test.go new file mode 100644 index 000000000..4a4e97e6e --- /dev/null +++ b/admin/auth_throttle_test.go @@ -0,0 +1,96 @@ +package admin + +import ( + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/gin-gonic/gin" +) + +func TestAdminAuthThrottleBlocksAfterFailures(t *testing.T) { + throttle := newAdminAuthThrottle() + now := time.Unix(1_700_000_000, 0) + + for i := 0; i < adminAuthMaxFailures-1; i++ { + if blocked := throttle.recordFailure("203.0.113.9", now); blocked != 0 { + t.Fatalf("failure %d blocked early: %v", i+1, blocked) + } + } + if retry := throttle.retryAfter("203.0.113.9", now); retry != 0 { + t.Fatalf("retryAfter before threshold = %v, want 0", retry) + } + + if blocked := throttle.recordFailure("203.0.113.9", now); blocked != adminAuthBlockDuration { + t.Fatalf("threshold failure blocked for %v, want %v", blocked, adminAuthBlockDuration) + } + if retry := throttle.retryAfter("203.0.113.9", now.Add(time.Minute)); retry != adminAuthBlockDuration-time.Minute { + t.Fatalf("retryAfter = %v, want %v", retry, adminAuthBlockDuration-time.Minute) + } + if retry := throttle.retryAfter("203.0.113.10", now); retry != 0 { + t.Fatalf("other IP should not be blocked, got %v", retry) + } + if retry := throttle.retryAfter("203.0.113.9", now.Add(adminAuthBlockDuration)); retry != 0 { + t.Fatalf("block should expire, got %v", retry) + } +} + +func TestAdminAuthThrottleResetsOnSuccess(t *testing.T) { + throttle := newAdminAuthThrottle() + now := time.Unix(1_700_000_000, 0) + for i := 0; i < adminAuthMaxFailures-1; i++ { + throttle.recordFailure("203.0.113.11", now) + } + throttle.reset("203.0.113.11") + if blocked := throttle.recordFailure("203.0.113.11", now); blocked != 0 { + t.Fatalf("reset should clear failure count, blocked for %v", blocked) + } +} + +func TestAdminAuthThrottleExpiresOldFailures(t *testing.T) { + throttle := newAdminAuthThrottle() + now := time.Unix(1_700_000_000, 0) + for i := 0; i < adminAuthMaxFailures-1; i++ { + throttle.recordFailure("203.0.113.12", now) + } + later := now.Add(adminAuthFailureWindow + time.Second) + if blocked := throttle.recordFailure("203.0.113.12", later); blocked != 0 { + t.Fatalf("failures outside the window should not accumulate, blocked for %v", blocked) + } +} + +func TestAdminAuthMiddlewareThrottlesFailedAttempts(t *testing.T) { + gin.SetMode(gin.TestMode) + handler, _ := newSelfServiceHandler(t, false) + router := gin.New() + handler.RegisterRoutes(router) + + request := func(key string) *httptest.ResponseRecorder { + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/api/admin/stats", nil) + req.Header.Set("X-Admin-Key", key) + router.ServeHTTP(rec, req) + return rec + } + + for i := 0; i < adminAuthMaxFailures-1; i++ { + if rec := request("wrong-key"); rec.Code != http.StatusUnauthorized { + t.Fatalf("failure %d returned %d, want 401", i+1, rec.Code) + } + } + if rec := request("wrong-key"); rec.Code != http.StatusTooManyRequests { + t.Fatalf("threshold failure returned %d, want 429", rec.Code) + } + if rec := request("wrong-key"); rec.Code != http.StatusTooManyRequests { + t.Fatalf("failure during block returned %d, want 429", rec.Code) + } else if rec.Header().Get("Retry-After") == "" { + t.Fatal("429 response missing Retry-After") + } + if rec := request("admin-secret"); rec.Code != http.StatusOK { + t.Fatalf("correct key during block returned %d, want 200", rec.Code) + } + if rec := request("wrong-key"); rec.Code != http.StatusUnauthorized { + t.Fatalf("failure after successful login returned %d, want 401", rec.Code) + } +} diff --git a/admin/handler.go b/admin/handler.go index 6455adc49..28a6ebd7d 100644 --- a/admin/handler.go +++ b/admin/handler.go @@ -110,6 +110,7 @@ type Handler struct { cacheDriver string cacheLabel string adminSecretEnv string + adminAuthThrottle *adminAuthThrottle imageProxy *proxy.Handler antigravitySyncAccount func(context.Context, int64) antigravityRefreshItem antigravityCapabilityProbe antigravityCapabilityExecutor @@ -428,6 +429,9 @@ const ( importDBWaitBackoff = 10 * time.Second ) +// maxAdminPage 限制分页页码，避免 (page-1)*page_size 溢出为负 OFFSET。 +const maxAdminPage = 100000 + type importRuntimeLoadSnapshot struct { RPM int64 Active int64 @@ -1042,6 +1046,7 @@ func NewHandler(store *auth.Store, db *database.DB, tc cache.TokenCache, rl *pro cacheDriver: tc.Driver(), cacheLabel: tc.Label(), adminSecretEnv: adminSecretEnv, + adminAuthThrottle: newAdminAuthThrottle(), imageProxy: proxy.NewHandler(store, db, nil, nil), chartCacheData: make(map[string]*chartCacheEntry), accountListCache: make(map[string]*accountListSnapshot), @@ -1436,19 +1441,43 @@ func (h *Handler) adminAuthMiddleware() gin.HandlerFunc { } } + clientIP := c.ClientIP() + // 清理输入 adminKey = security.SanitizeInput(adminKey) // 使用安全比较防止时序攻击 if !security.SecureCompare(adminKey, adminSecret) { // 记录安全审计日志 - security.SecurityAuditLog("ADMIN_AUTH_FAILED", fmt.Sprintf("path=%s ip=%s source=%s", c.Request.URL.Path, c.ClientIP(), source)) + security.SecurityAuditLog("ADMIN_AUTH_FAILED", fmt.Sprintf("path=%s ip=%s source=%s", c.Request.URL.Path, security.SanitizeLog(clientIP), source)) + // 阻断只针对失败尝试；正确密钥仍可通过，避免攻击者用共享出口 IP 锁死管理员。 + if retry := h.adminAuthThrottle.retryAfter(clientIP, time.Now()); retry > 0 { + c.Header("Retry-After", strconv.Itoa(int((retry+time.Second-1)/time.Second))) + c.JSON(http.StatusTooManyRequests, gin.H{ + "error": "认证失败次数过多，请稍后再试", + }) + c.Abort() + return + } + if blocked := h.adminAuthThrottle.recordFailure(clientIP, time.Now()); blocked > 0 { + c.Header("Retry-After", strconv.Itoa(int((blocked+time.Second-1)/time.Second))) + c.JSON(http.StatusTooManyRequests, gin.H{ + "error": "认证失败次数过多，请稍后再试", + }) + c.Abort() + return + } c.JSON(http.StatusUnauthorized, gin.H{ "error": "管理密钥无效或缺失", }) c.Abort() return } + h.adminAuthThrottle.reset(clientIP) + + // 管理端响应包含 raw_key、auth.json、settings 等秘密，禁止中间缓存保存。 + c.Header("Cache-Control", "no-store") + c.Header("Pragma", "no-cache") // 成功认证，记录审计日志 if security.IsSensitiveEndpoint(c.Request.URL.Path) { @@ -8618,10 +8647,12 @@ type createKeyReq struct { } // generateKey 生成随机 API Key -func generateKey() string { +func generateKey() (string, error) { b := make([]byte, 24) - rand.Read(b) - return "sk-" + hex.EncodeToString(b) + if _, err := rand.Read(b); err != nil { + return "", fmt.Errorf("generate api key: %w", err) + } + return "sk-" + hex.EncodeToString(b), nil } // CreateAPIKey 创建新 API 密钥（增强版，带输入验证） @@ -8678,7 +8709,12 @@ func (h *Handler) CreateAPIKey(c *gin.Context) { key := req.Key if key == "" { - key = generateKey() + var err error + key, err = generateKey() + if err != nil { + writeInternalError(c, err) + return + } } else { // 验证用户提供的key格式 key = security.SanitizeInput(key) diff --git a/admin/image_mime_test.go b/admin/image_mime_test.go new file mode 100644 index 000000000..0a62e5193 --- /dev/null +++ b/admin/image_mime_test.go @@ -0,0 +1,26 @@ +package admin + +import "testing" + +func TestSafeInlineImageMimeType(t *testing.T) { + tests := []struct { + input string + wantType string + wantInline bool + }{ + {"image/png", "image/png", true}, + {"IMAGE/JPEG; charset=binary", "image/jpeg", true}, + {"image/webp", "image/webp", true}, + {"image/gif", "image/gif", true}, + {"text/html", "application/octet-stream", false}, + {"image/svg+xml", "application/octet-stream", false}, + {"", "application/octet-stream", false}, + } + for _, test := range tests { + gotType, gotInline := safeInlineImageMimeType(test.input) + if gotType != test.wantType || gotInline != test.wantInline { + t.Errorf("safeInlineImageMimeType(%q) = (%q, %v), want (%q, %v)", + test.input, gotType, gotInline, test.wantType, test.wantInline) + } + } +} diff --git a/admin/image_studio.go b/admin/image_studio.go index 5500cc30d..e11aabaeb 100644 --- a/admin/image_studio.go +++ b/admin/image_studio.go @@ -707,12 +707,23 @@ func (h *Handler) serveImageAssetFile(c *gin.Context, asset *database.ImageAsset if opts.download { disposition = "attachment" } + contentType, inlineSafe := safeInlineImageMimeType(asset.MimeType) + if !inlineSafe { + // 上游或客户端回显的 MIME 不可信；非白名单类型强制下载，避免 text/html、 + // SVG 等在同源页面里执行脚本。 + disposition = "attachment" + } filename := sanitizeDownloadFilename(asset.Filename) if opts.thumbKB > 0 && !opts.download { - data, contentType, err := cachedImageThumbnail(c.Request.Context(), asset, backend, opts.thumbKB) + data, thumbType, err := cachedImageThumbnail(c.Request.Context(), asset, backend, opts.thumbKB) if err == nil { - c.Header("Content-Disposition", fmt.Sprintf(`inline; filename="%s"`, thumbnailFilename(filename))) - c.Data(http.StatusOK, contentType, data) + thumbDisposition := "inline" + thumbContentType, thumbSafe := safeInlineImageMimeType(thumbType) + if !thumbSafe { + thumbDisposition = "attachment" + } + c.Header("Content-Disposition", fmt.Sprintf(`%s; filename="%s"`, thumbDisposition, thumbnailFilename(filename))) + c.Data(http.StatusOK, thumbContentType, data) return } // Preserve the original endpoint's full-image fallback for formats that @@ -722,9 +733,7 @@ func (h *Handler) serveImageAssetFile(c *gin.Context, asset *database.ImageAsset } } - if strings.TrimSpace(asset.MimeType) != "" { - c.Header("Content-Type", asset.MimeType) - } + c.Header("Content-Type", contentType) c.Header("Content-Disposition", fmt.Sprintf(`%s; filename="%s"`, disposition, filename)) rc, size, err := backend.Open(c.Request.Context(), asset.StoragePath) @@ -733,10 +742,6 @@ func (h *Handler) serveImageAssetFile(c *gin.Context, asset *database.ImageAsset return } defer rc.Close() - contentType := strings.TrimSpace(asset.MimeType) - if contentType == "" { - contentType = "application/octet-stream" - } if size < 0 { // S3 偶尔不带 Content-Length，退化为 chunked。 c.Status(http.StatusOK) @@ -747,6 +752,17 @@ func (h *Handler) serveImageAssetFile(c *gin.Context, asset *database.ImageAsset c.DataFromReader(http.StatusOK, size, contentType, rc, nil) } +// safeInlineImageMimeType 只放行浏览器按图片渲染的常见类型。 +// 其他值（包括 text/html、image/svg+xml 和空值）返回 application/octet-stream。 +func safeInlineImageMimeType(value string) (string, bool) { + mediaType := strings.ToLower(strings.TrimSpace(strings.SplitN(value, ";", 2)[0])) + switch mediaType { + case "image/png", "image/jpeg", "image/webp", "image/gif": + return mediaType, true + } + return "application/octet-stream", false +} + func imageAssetPathAllowed(storagePath string) bool { storagePath = strings.TrimSpace(storagePath) if imagestore.IsS3Ref(storagePath) { @@ -1714,6 +1730,9 @@ func paginationParams(c *gin.Context, defaultPageSize int) (int, int) { page = parsed } } + if page > maxAdminPage { + page = maxAdminPage + } pageSize := defaultPageSize if raw := c.Query("page_size"); raw != "" { if parsed, err := strconv.Atoi(raw); err == nil && parsed > 0 { diff --git a/admin/key_usage_public.go b/admin/key_usage_public.go index d2403ded9..4a3b3a0eb 100644 --- a/admin/key_usage_public.go +++ b/admin/key_usage_public.go @@ -166,6 +166,9 @@ func parsePublicAPIKeyUsageLogPagination(pageRaw, pageSizeRaw string) (int, int, if err != nil || parsed < 1 { return 0, 0, errors.New("page 参数无效，需要正整数") } + if parsed > maxAdminPage { + return 0, 0, errors.New("page 参数过大") + } page = parsed } diff --git a/admin/pagination_test.go b/admin/pagination_test.go new file mode 100644 index 000000000..1c09ca3cd --- /dev/null +++ b/admin/pagination_test.go @@ -0,0 +1,28 @@ +package admin + +import ( + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" +) + +func TestPaginationParamsCapsPage(t *testing.T) { + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest("GET", "/?page=999999999&page_size=100", nil) + page, pageSize := paginationParams(c, 20) + if page != maxAdminPage || pageSize != 100 { + t.Fatalf("paginationParams = (%d, %d), want (%d, 100)", page, pageSize, maxAdminPage) + } +} + +func TestPublicUsagePaginationRejectsOversizedPage(t *testing.T) { + if _, _, err := parsePublicAPIKeyUsageLogPagination("999999999", "100"); err == nil { + t.Fatal("oversized page should be rejected") + } + page, size, err := parsePublicAPIKeyUsageLogPagination("2", "100") + if err != nil || page != 2 || size != 100 { + t.Fatalf("normal page = (%d, %d, %v)", page, size, err) + } +} diff --git a/admin/self_service.go b/admin/self_service.go index 174d9e332..0a4b731e9 100644 --- a/admin/self_service.go +++ b/admin/self_service.go @@ -262,11 +262,11 @@ func (h *Handler) upsertSelfServiceAccount(ctx context.Context, name, proxyURL s // 的成功语义，而账号永远缺少联系人/标签等待审核信息。 note := "自助提交联系人: " + contactEmail if err := h.db.UpdateAccountNote(ctx, id, note); err != nil { - h.rollbackSelfServiceAccountInsert(ctx, id) + h.rollbackSelfServiceAccountInsert(id) return 0, err } if err := h.db.SetAccountTags(ctx, id, []string{selfServiceTag}); err != nil { - h.rollbackSelfServiceAccountInsert(ctx, id) + h.rollbackSelfServiceAccountInsert(id) return 0, err } h.db.InsertAccountEventAsync(id, "added", selfServiceSource) @@ -279,7 +279,10 @@ func (h *Handler) upsertSelfServiceAccount(ctx context.Context, name, proxyURL s // rollbackSelfServiceAccountInsert 软删本次插入但未完成待审核设置的账号， // 使同一身份可以重新提交。软删失败时账号保持 disabled，仍不会进调度池。 -func (h *Handler) rollbackSelfServiceAccountInsert(ctx context.Context, id int64) { +// 使用独立的有界上下文，避免请求超时后回滚也失败。 +func (h *Handler) rollbackSelfServiceAccountInsert(id int64) { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() if err := h.db.SoftDeleteAccount(ctx, id); err != nil { log.Printf("回滚未完成的自助账号 %d 失败: %v", id, err) return diff --git a/admin/self_service_test.go b/admin/self_service_test.go index 9cc326654..c59831de5 100644 --- a/admin/self_service_test.go +++ b/admin/self_service_test.go @@ -195,7 +195,7 @@ func TestSelfServiceRollbackAllowsRetry(t *testing.T) { if err != nil { t.Fatalf("InsertDisabledAccountWithCredentials: %v", err) } - handler.rollbackSelfServiceAccountInsert(ctx, id) + handler.rollbackSelfServiceAccountInsert(id) if _, err := handler.upsertSelfServiceAccount(ctx, "rollback@example.com", "", seed, "c@x.com"); err != nil { t.Fatalf("回滚后同一身份应可重新提交, got %v", err) diff --git a/frontend/scripts/audit-gate.mjs b/frontend/scripts/audit-gate.mjs index 0c9c13393..9ef2f4444 100644 --- a/frontend/scripts/audit-gate.mjs +++ b/frontend/scripts/audit-gate.mjs @@ -20,9 +20,12 @@ const ALLOWLIST = new Map([ const BLOCKING_SEVERITIES = new Set(["high", "critical"]); -const result = spawnSync("npm", ["audit", "--omit=dev", "--json"], { +const npmCommand = process.platform === "win32" ? "npm.cmd" : "npm"; +const result = spawnSync(npmCommand, ["audit", "--omit=dev", "--json"], { encoding: "utf8", maxBuffer: 32 * 1024 * 1024, + // Windows 上 npm 是 .cmd，需要经 cmd.exe 启动。 + shell: process.platform === "win32", }); // npm audit 检出漏洞时退出码非 0,这属正常输出而非执行失败,只有拿不到 stdout 才是真失败。 diff --git a/internal/imageupscale/imageupscale.go b/internal/imageupscale/imageupscale.go index a45440968..f8b780b86 100644 --- a/internal/imageupscale/imageupscale.go +++ b/internal/imageupscale/imageupscale.go @@ -22,6 +22,13 @@ import ( const maxUpscalerResponseBytes = 64 << 20 +// 目标尺寸上限：单边 16384，总像素 4000 万（RGBA 约 160MB）。 +// 超过上限的 strict_size 不会被当作精确目标，避免 NewRGBA 直接分配 GB 级内存。 +const ( + maxTargetDimension = 16384 + maxTargetPixels = 40_000_000 +) + // Result describes one applied upscale. Data is the re-encoded image. type Result struct { Data []byte @@ -118,6 +125,10 @@ func parseSize(size string) (int, int, bool) { if err != nil || height <= 0 { return 0, 0, false } + if width > maxTargetDimension || height > maxTargetDimension || + int64(width)*int64(height) > maxTargetPixels { + return 0, 0, false + } return width, height, true } diff --git a/internal/imageupscale/imageupscale_test.go b/internal/imageupscale/imageupscale_test.go index b1f4aa3d2..3db08ef29 100644 --- a/internal/imageupscale/imageupscale_test.go +++ b/internal/imageupscale/imageupscale_test.go @@ -24,6 +24,23 @@ func strictTestPNG(t *testing.T, width, height int) []byte { return buf.Bytes() } +func TestParseSizeRejectsOversizedTargets(t *testing.T) { + tests := []string{ + "20000x20000", + "16385x100", + "100x16385", + "10000x5000", // 50M pixels + } + for _, size := range tests { + if width, height, ok := ParseSize(size); ok { + t.Errorf("ParseSize(%q) = %dx%d, want rejected", size, width, height) + } + } + if width, height, ok := ParseSize("16384x2048"); !ok || width != 16384 || height != 2048 { + t.Errorf("ParseSize(16384x2048) = %dx%d, %v; want 16384x2048", width, height, ok) + } +} + func TestBytesWithFitHonorsStrictSizeWithoutUpscaleTier(t *testing.T) { t.Setenv("IMAGE_UPSCALER_ENDPOINT", "") data, contentType, method, err := BytesWithFit(context.Background(), strictTestPNG(t, 4, 4), "", "12x6", "pad", true) diff --git a/proxy/antigravity_stream_errors.go b/proxy/antigravity_stream_errors.go index 616870fb5..8e4603ff8 100644 --- a/proxy/antigravity_stream_errors.go +++ b/proxy/antigravity_stream_errors.go @@ -89,7 +89,7 @@ func (r *antigravitySSEReader) next() ([]byte, error) { case r.err != nil: err = r.err default: - line, err = r.reader.ReadBytes('\n') + line, err = readBoundedSSELine(r.reader, maxSSELineBytes) } if len(line) > 0 { line = bytes.TrimRight(line, "\r\n") diff --git a/proxy/apikey_admission.go b/proxy/apikey_admission.go new file mode 100644 index 000000000..e61e15d97 --- /dev/null +++ b/proxy/apikey_admission.go @@ -0,0 +1,128 @@ +package proxy + +import ( + "sync" + "time" +) + +const ( + apiKeyAdmissionRPMSlots = 60 + apiKeyAdmissionMaxKeys = 8192 + apiKeyAdmissionIdleTTL = 10 * time.Minute +) + +// apiKeyAdmissionLimiter 在进程内按实际请求递增 RPM/RPD 计数。 +// +// 数据库聚合 + 60s 缓存适合展示和跨进程兜底，但不能约束同一缓存窗口内的并发 +// 突发：所有请求都读到同一个旧值。这里在准入点用原子计数补上这一段，缓存只负责 +// 历史用量。多实例部署时每个副本各有一份计数，跨副本仍需共享计数（见 issue）。 +type apiKeyAdmissionLimiter struct { + mu sync.Mutex + entries map[int64]*apiKeyAdmissionEntry +} + +type apiKeyAdmissionEntry struct { + rpmBuckets [apiKeyAdmissionRPMSlots]int64 + rpmSecond int64 + rpdStart time.Time + rpdCount int64 + lastSeen time.Time +} + +func newAPIKeyAdmissionLimiter() *apiKeyAdmissionLimiter { + return &apiKeyAdmissionLimiter{entries: make(map[int64]*apiKeyAdmissionEntry)} +} + +func (h *Handler) admissionLimiter() *apiKeyAdmissionLimiter { + if h == nil { + return nil + } + h.admissionOnce.Do(func() { + h.admission = newAPIKeyAdmissionLimiter() + }) + return h.admission +} + +// allow 在 RPM/RPD 限额内为 requests 个上游请求预留计数。 +// limit <= 0 表示该维度不限制。 +func (l *apiKeyAdmissionLimiter) allow(keyID int64, rpm, rpd, requests int, now time.Time) bool { + if l == nil || keyID <= 0 || requests <= 0 { + return true + } + if rpm <= 0 && rpd <= 0 { + return true + } + l.mu.Lock() + defer l.mu.Unlock() + + entry := l.entries[keyID] + if entry == nil { + entry = &apiKeyAdmissionEntry{} + l.entries[keyID] = entry + l.evictLocked(now) + } + entry.lastSeen = now + + if rpm > 0 && entry.rpmUsed(now)+int64(requests) > int64(rpm) { + return false + } + if rpd > 0 { + if entry.rpdStart.IsZero() || now.Sub(entry.rpdStart) >= apiKeyRPDWindow { + entry.rpdStart = now + entry.rpdCount = 0 + } + if entry.rpdCount+int64(requests) > int64(rpd) { + return false + } + } + + // 两个维度都通过后再一起记账，避免 RPD 拒绝时仍消耗 RPM 配额。 + if rpm > 0 { + entry.rpmBuckets[now.Unix()%apiKeyAdmissionRPMSlots] += int64(requests) + } + if rpd > 0 { + entry.rpdCount += int64(requests) + } + return true +} + +// rpmUsed 推进秒级桶并返回最近 60 秒的计数，不写入本次请求。 +func (e *apiKeyAdmissionEntry) rpmUsed(now time.Time) int64 { + second := now.Unix() + switch { + case e.rpmSecond == 0 || second < e.rpmSecond || second-e.rpmSecond >= apiKeyAdmissionRPMSlots: + e.rpmBuckets = [apiKeyAdmissionRPMSlots]int64{} + e.rpmSecond = second + default: + for s := e.rpmSecond + 1; s <= second; s++ { + e.rpmBuckets[s%apiKeyAdmissionRPMSlots] = 0 + } + e.rpmSecond = second + } + var used int64 + for _, count := range e.rpmBuckets { + used += count + } + return used +} + +func (l *apiKeyAdmissionLimiter) evictLocked(now time.Time) { + if len(l.entries) <= apiKeyAdmissionMaxKeys { + return + } + for key, entry := range l.entries { + if now.Sub(entry.lastSeen) > apiKeyAdmissionIdleTTL && now.Sub(entry.rpdStart) >= apiKeyRPDWindow { + delete(l.entries, key) + } + } + for len(l.entries) > apiKeyAdmissionMaxKeys { + var oldestKey int64 + var oldest time.Time + for key, entry := range l.entries { + if oldestKey == 0 || entry.lastSeen.Before(oldest) { + oldestKey, oldest = key, entry.lastSeen + } + } + delete(l.entries, oldestKey) + } +} diff --git a/proxy/apikey_admission_test.go b/proxy/apikey_admission_test.go new file mode 100644 index 000000000..b68cedd5f --- /dev/null +++ b/proxy/apikey_admission_test.go @@ -0,0 +1,79 @@ +package proxy + +import ( + "sync" + "sync/atomic" + "testing" + "time" +) + +func TestAPIKeyAdmissionLimiterRejectsBurst(t *testing.T) { + limiter := newAPIKeyAdmissionLimiter() + now := time.Unix(1_700_000_000, 0) + + for i := 0; i < 2; i++ { + if !limiter.allow(1, 2, 0, 1, now) { + t.Fatalf("request %d within RPM rejected", i+1) + } + } + if limiter.allow(1, 2, 0, 1, now) { + t.Fatal("third request in the same second should exceed RPM=2") + } + if !limiter.allow(2, 2, 0, 1, now) { + t.Fatal("other API key should not be affected") + } + if !limiter.allow(1, 2, 0, 1, now.Add(time.Minute)) { + t.Fatal("request should be admitted after the RPM window rolls") + } +} + +func TestAPIKeyAdmissionLimiterConcurrentBurst(t *testing.T) { + limiter := newAPIKeyAdmissionLimiter() + now := time.Unix(1_700_000_000, 0) + var allowed atomic.Int64 + var wg sync.WaitGroup + for i := 0; i < 20; i++ { + wg.Add(1) + go func() { + defer wg.Done() + if limiter.allow(1, 5, 0, 1, now) { + allowed.Add(1) + } + }() + } + wg.Wait() + if got := allowed.Load(); got != 5 { + t.Fatalf("allowed = %d, want exactly 5 (RPM limit)", got) + } +} + +func TestAPIKeyAdmissionLimiterReservesBatchRequests(t *testing.T) { + limiter := newAPIKeyAdmissionLimiter() + now := time.Unix(1_700_000_000, 0) + if !limiter.allow(1, 10, 0, 4, now) { + t.Fatal("batch of 4 within RPM=10 rejected") + } + if limiter.allow(1, 10, 0, 7, now) { + t.Fatal("batch of 7 would exceed the remaining RPM budget") + } + if !limiter.allow(1, 10, 0, 6, now) { + t.Fatal("batch of 6 should fit exactly") + } +} + +func TestAPIKeyAdmissionLimiterEnforcesRPD(t *testing.T) { + limiter := newAPIKeyAdmissionLimiter() + now := time.Unix(1_700_000_000, 0) + if !limiter.allow(1, 0, 2, 1, now) { + t.Fatal("first RPD request rejected") + } + if !limiter.allow(1, 0, 2, 1, now.Add(time.Hour)) { + t.Fatal("second RPD request rejected") + } + if limiter.allow(1, 0, 2, 1, now.Add(2*time.Hour)) { + t.Fatal("third RPD request should exceed RPD=2") + } + if !limiter.allow(1, 0, 2, 1, now.Add(apiKeyRPDWindow)) { + t.Fatal("RPD window should reset after 24h") + } +} diff --git a/proxy/apikey_limits.go b/proxy/apikey_limits.go index 5b633c283..e2553df03 100644 --- a/proxy/apikey_limits.go +++ b/proxy/apikey_limits.go @@ -77,42 +77,40 @@ func (h *Handler) payloadRuleIdentity(c *gin.Context) *PayloadRuleIdentity { // EnforceAPIKeyLimits checks API key scoped model and rate/cost limits. func (h *Handler) EnforceAPIKeyLimits(c *gin.Context, model string) (int, string) { - return h.enforceAPIKeyLimits(c, model) + return h.enforceAPIKeyLimitsN(c, model, 1) } // EnforceAPIKeyLimitsForRequests performs the normal API-key checks and also // verifies that a batched request cannot cross the RPM/RPD boundary by issuing -// multiple upstream calls after a single admission check. Each upstream call -// still runs the normal handler checks and records its own usage entry. +// multiple upstream calls after a single admission check. The whole batch is +// reserved atomically so concurrent batches cannot overshoot the window. func (h *Handler) EnforceAPIKeyLimitsForRequests(c *gin.Context, model string, requests int) (int, string) { - if status, msg := h.enforceAPIKeyLimits(c, model); status != 0 { - return status, msg - } - if requests <= 1 { - return 0, "" + if requests < 1 { + requests = 1 } + // 先做批量边界检查，避免整批被拒绝时仍消耗一个 RPM 配额。 + // Each upstream call still runs the normal handler checks and records its own + // usage entry, so the single admission below only reserves the batch size. row := apiKeyRowFromContext(c) - if row == nil || row.ID <= 0 { - return 0, "" - } - - ctx, cancel := context.WithTimeout(c.Request.Context(), 2*time.Second) - defer cancel() - if row.Limits.RPM > 0 { - count, err := h.apiKeyWindowRequests(ctx, row.ID, "rpm", apiKeyRPMWindow) - if err == nil && apiKeyBatchWouldExceed(count, row.Limits.RPM, requests) { - return http.StatusTooManyRequests, - fmt.Sprintf("API key rate limit exceeded: batch of %d would exceed %d requests per minute", requests, row.Limits.RPM) + if row != nil && row.ID > 0 && requests > 1 { + ctx, cancel := context.WithTimeout(c.Request.Context(), 2*time.Second) + defer cancel() + if row.Limits.RPM > 0 { + count, err := h.apiKeyWindowRequests(ctx, row.ID, "rpm", apiKeyRPMWindow) + if err == nil && apiKeyBatchWouldExceed(count, row.Limits.RPM, requests) { + return http.StatusTooManyRequests, + fmt.Sprintf("API key rate limit exceeded: batch of %d would exceed %d requests per minute", requests, row.Limits.RPM) + } } - } - if row.Limits.RPD > 0 { - count, err := h.apiKeyWindowRequests(ctx, row.ID, "rpd", apiKeyRPDWindow) - if err == nil && apiKeyBatchWouldExceed(count, row.Limits.RPD, requests) { - return http.StatusTooManyRequests, - fmt.Sprintf("API key rate limit exceeded: batch of %d would exceed %d requests per day", requests, row.Limits.RPD) + if row.Limits.RPD > 0 { + count, err := h.apiKeyWindowRequests(ctx, row.ID, "rpd", apiKeyRPDWindow) + if err == nil && apiKeyBatchWouldExceed(count, row.Limits.RPD, requests) { + return http.StatusTooManyRequests, + fmt.Sprintf("API key rate limit exceeded: batch of %d would exceed %d requests per day", requests, row.Limits.RPD) + } } } - return 0, "" + return h.enforceAPIKeyLimitsN(c, model, requests) } func apiKeyBatchWouldExceed(current int64, limit, requests int) bool { @@ -129,6 +127,15 @@ func apiKeyBatchWouldExceed(current int64, limit, requests int) bool { // - http.StatusForbidden (403): 模型不在白名单 / 在黑名单 // - http.StatusTooManyRequests (429): rpm/rpd/cost/token 任一窗口超额 func (h *Handler) enforceAPIKeyLimits(c *gin.Context, model string) (int, string) { + return h.enforceAPIKeyLimitsN(c, model, 1) +} + +// enforceAPIKeyLimitsN 是准入检查的实际实现。requests 表示本次准入会消耗的 +// 上游请求数（批量生图一次可产生多个请求），用于原子预留 RPM/RPD 计数。 +func (h *Handler) enforceAPIKeyLimitsN(c *gin.Context, model string, requests int) (int, string) { + if requests < 1 { + requests = 1 + } var daybreakStatus int var daybreakMessage string model, daybreakStatus, daybreakMessage = h.checkDaybreakRequest(c, model) @@ -264,6 +271,17 @@ func (h *Handler) enforceAPIKeyLimits(c *gin.Context, model string) (int, string } } + // 7. 进程内准入计数。数据库聚合 + 60s 缓存读到的是同一份旧值，无法约束 + // 缓存窗口内的并发突发；这里按实际请求数原子递增后再放行。 + if !h.admissionLimiter().allow(row.ID, limits.RPM, limits.RPD, requests, time.Now()) { + if limits.RPM > 0 { + return http.StatusTooManyRequests, + fmt.Sprintf("API key rate limit exceeded: %d requests per minute", limits.RPM) + } + return http.StatusTooManyRequests, + fmt.Sprintf("API key rate limit exceeded: %d requests per day", limits.RPD) + } + return 0, "" } diff --git a/proxy/codex_turn_state.go b/proxy/codex_turn_state.go index 99238bf7e..f489b8d28 100644 --- a/proxy/codex_turn_state.go +++ b/proxy/codex_turn_state.go @@ -21,7 +21,12 @@ import ( // // 键使用 affinityKey(下游会话标识 + API Key,与账号粘性绑定同源),保证 // 同一段对话的记录/守卫两侧落在同一键上。无会话标识时不做跟踪(保持透传)。 -const codexTurnStateProvenanceTTL = time.Hour +const ( + codexTurnStateProvenanceTTL = time.Hour + // 溯源表只用于剥离跨账号回带，条目可以安全丢弃；设硬上限避免高基数 + // 会话键在 TTL 内无限增长。 + codexTurnStateMaxOrigins = 65536 +) type codexTurnStateOrigin struct { accountID int64 @@ -183,11 +188,28 @@ func sweepCodexTurnStateOrigins() { return } now := time.Now() + live := 0 codexTurnStateOrigins.Range(func(key, value any) bool { origin, ok := value.(codexTurnStateOrigin) if !ok || (!origin.expiresAt.IsZero() && now.After(origin.expiresAt)) { codexTurnStateOrigins.Delete(key) + return true + } + live++ + return true + }) + if live <= codexTurnStateMaxOrigins { + return + } + // 超过上限时丢弃任意条目。溯源缺失只会让守卫放行一次回带，不会影响 + // 转发正确性；这比让内存随会话键数量持续增长更安全。 + excess := live - codexTurnStateMaxOrigins + codexTurnStateOrigins.Range(func(key, _ any) bool { + if excess <= 0 { + return false } + codexTurnStateOrigins.Delete(key) + excess-- return true }) } diff --git a/proxy/codex_turn_state_cap_test.go b/proxy/codex_turn_state_cap_test.go new file mode 100644 index 000000000..5c9dc3f99 --- /dev/null +++ b/proxy/codex_turn_state_cap_test.go @@ -0,0 +1,40 @@ +package proxy + +import ( + "fmt" + "strings" + "testing" + + "github.com/codex2api/auth" +) + +func TestCodexTurnStateOriginsCapped(t *testing.T) { + const prefix = "cap-test-" + remove := func() { + codexTurnStateOrigins.Range(func(key, _ any) bool { + if name, ok := key.(string); ok && strings.HasPrefix(name, prefix) { + codexTurnStateOrigins.Delete(key) + } + return true + }) + } + remove() + t.Cleanup(remove) + + account := &auth.Account{DBID: 1} + for i := 0; i < codexTurnStateMaxOrigins+5000; i++ { + noteCodexTurnStateProvenance(fmt.Sprintf("%s%d", prefix, i), account) + } + + live := 0 + codexTurnStateOrigins.Range(func(key, _ any) bool { + if name, ok := key.(string); ok && strings.HasPrefix(name, prefix) { + live++ + } + return true + }) + // 最后一次清扫后最多再写入 255 条。 + if live > codexTurnStateMaxOrigins+255 { + t.Fatalf("turn-state origins = %d, want <= %d", live, codexTurnStateMaxOrigins+255) + } +} diff --git a/proxy/executor.go b/proxy/executor.go index 475806efc..ebbc10749 100644 --- a/proxy/executor.go +++ b/proxy/executor.go @@ -1030,23 +1030,27 @@ func isCodexAccessRestrictedResponse(resp *http.Response) bool { if resp == nil || resp.StatusCode != http.StatusForbidden || resp.Body == nil { return false } - body, err := io.ReadAll(resp.Body) - _ = resp.Body.Close() - resp.Body = io.NopCloser(bytes.NewReader(body)) - resp.ContentLength = int64(len(body)) + // 只读取用于识别的前缀，再把前缀和剩余流拼回去，避免把巨大的错误体 + // 整体读进内存（上游可能由第三方中转控制）。 + const peekLimit = 64 << 10 + peek, err := io.ReadAll(io.LimitReader(resp.Body, peekLimit)) + resp.Body = struct { + io.Reader + io.Closer + }{io.MultiReader(bytes.NewReader(peek), resp.Body), resp.Body} if err != nil { return false } - if strings.EqualFold(strings.TrimSpace(gjson.GetBytes(body, "error.code").String()), "codex_access_restricted") { + if strings.EqualFold(strings.TrimSpace(gjson.GetBytes(peek, "error.code").String()), "codex_access_restricted") { return true } // 部分中转（sub2api 的 codex_cli_only、zzzcoding 等第三方网关）用 // {"error":{"message":"This account only allows Codex official clients", // "type":"forbidden_error"}} 表达同样的官方客户端限制。 - if !strings.EqualFold(strings.TrimSpace(gjson.GetBytes(body, "error.type").String()), "forbidden_error") { + if !strings.EqualFold(strings.TrimSpace(gjson.GetBytes(peek, "error.type").String()), "forbidden_error") { return false } - msg := strings.ToLower(strings.TrimSpace(gjson.GetBytes(body, "error.message").String())) + msg := strings.ToLower(strings.TrimSpace(gjson.GetBytes(peek, "error.message").String())) return strings.Contains(msg, "official clients") || strings.Contains(msg, "codex client") || strings.Contains(msg, "codex 客户端") diff --git a/proxy/executor_body_limit_test.go b/proxy/executor_body_limit_test.go new file mode 100644 index 000000000..d4cb5ce73 --- /dev/null +++ b/proxy/executor_body_limit_test.go @@ -0,0 +1,48 @@ +package proxy + +import ( + "bytes" + "io" + "net/http" + "strings" + "testing" +) + +func TestCodexAccessRestrictedPeekPreservesBody(t *testing.T) { + body := []byte(`{"error":{"code":"codex_access_restricted","message":"restricted"}}`) + resp := &http.Response{ + StatusCode: http.StatusForbidden, + Body: io.NopCloser(bytes.NewReader(body)), + } + if !isCodexAccessRestrictedResponse(resp) { + t.Fatal("expected codex_access_restricted to be detected") + } + got, err := io.ReadAll(resp.Body) + if err != nil { + t.Fatalf("read restored body: %v", err) + } + if !bytes.Equal(got, body) { + t.Fatalf("restored body = %q, want %q", got, body) + } +} + +func TestCodexAccessRestrictedPeekDoesNotBufferWholeBody(t *testing.T) { + prefix := `{"error":{"code":"other","padding":"` + suffix := `"}}` + padding := strings.Repeat("x", 2<<20) + body := []byte(prefix + padding + suffix) + resp := &http.Response{ + StatusCode: http.StatusForbidden, + Body: io.NopCloser(bytes.NewReader(body)), + } + if isCodexAccessRestrictedResponse(resp) { + t.Fatal("non-restricted error must not be detected") + } + got, err := io.ReadAll(resp.Body) + if err != nil { + t.Fatalf("read restored body: %v", err) + } + if !bytes.Equal(got, body) { + t.Fatalf("restored body length = %d, want %d", len(got), len(body)) + } +} diff --git a/proxy/grok_namespace_tools.go b/proxy/grok_namespace_tools.go index 09100479e..21a5124f4 100644 --- a/proxy/grok_namespace_tools.go +++ b/proxy/grok_namespace_tools.go @@ -1069,7 +1069,7 @@ func newGrokNamespaceReverser(body io.ReadCloser, streaming bool, aliases map[st reverser := newGrokStreamReverser(aliases) reader := bufio.NewReader(body) for { - line, err := reader.ReadBytes('\n') + line, err := readBoundedSSELine(reader, maxSSELineBytes) if len(line) > 0 { rewritten := reverser.rewriteLine(line) if len(rewritten) > 0 { diff --git a/proxy/grok_protocol.go b/proxy/grok_protocol.go index 28eff12fa..6be4a4357 100644 --- a/proxy/grok_protocol.go +++ b/proxy/grok_protocol.go @@ -1084,7 +1084,7 @@ func (r *chatToResponsesReader) emitChatTerminal(status string) { func readSSEDataLine(reader *bufio.Reader) ([]byte, error) { var data []byte for { - line, err := reader.ReadBytes('\n') + line, err := readBoundedSSELine(reader, maxSSELineBytes) if len(line) > 0 { line = bytes.TrimRight(line, "\r\n") if bytes.HasPrefix(line, []byte("data:")) { diff --git a/proxy/handler.go b/proxy/handler.go index 676c42419..ddf0a413b 100644 --- a/proxy/handler.go +++ b/proxy/handler.go @@ -67,6 +67,8 @@ type Handler struct { apiKeyGateMu sync.Mutex promptRiskMu sync.Mutex apiKeyGate *apiKeyConcurrencyLimiter + admissionOnce sync.Once + admission *apiKeyAdmissionLimiter scopeUsageMu sync.Mutex scopeUsage *apiKeyScopeUsageTracker scopeDeltaInit sync.Once @@ -89,6 +91,9 @@ const ( apiKeyCountCacheNamespace = "api-key-count" apiKeyCacheTTL = 5 * time.Minute apiKeyCountCacheTTL = 30 * time.Second + // 上游错误体只用于识别错误类型和回传给客户端；限制上限，避免恶意或 + // 失控上游用超大 body 把网关内存打满。 + maxUpstreamErrorBodyBytes = 1 << 20 ) type apiKeyRuntimeRecord struct { diff --git a/proxy/handler_gemini.go b/proxy/handler_gemini.go index 9f5b4acb5..e2594e707 100644 --- a/proxy/handler_gemini.go +++ b/proxy/handler_gemini.go @@ -269,7 +269,7 @@ func (h *Handler) handleGeminiCountTokens(c *gin.Context, model string, rawBody } if resp.StatusCode != http.StatusOK { - errBody, _ := io.ReadAll(resp.Body) + errBody, _ := io.ReadAll(io.LimitReader(resp.Body, maxUpstreamErrorBodyBytes)) _ = resp.Body.Close() lastStatusCode = resp.StatusCode lastBody = errBody @@ -541,7 +541,7 @@ func (h *Handler) handleGeminiGenerateContent(c *gin.Context, model string, rawB } if resp.StatusCode != http.StatusOK { - errBody, _ := io.ReadAll(resp.Body) + errBody, _ := io.ReadAll(io.LimitReader(resp.Body, maxUpstreamErrorBodyBytes)) rememberContinuousRetryHTTPFailure(c.Request.Context(), resp, errBody) _ = resp.Body.Close() upstreamCancel() diff --git a/proxy/openai_responses_ws.go b/proxy/openai_responses_ws.go index 87979bd4d..e117b1161 100644 --- a/proxy/openai_responses_ws.go +++ b/proxy/openai_responses_ws.go @@ -25,8 +25,13 @@ const ( openAIResponsesWebsocketBetaHeader = "responses_websockets=2026-02-06" openAIResponsesWSMaxConns = 8 openAIResponsesWSMaxAge = 30 * time.Minute + // gorilla 的 SetReadLimit 按压缩后的帧字节计数；解压后的消息大小必须 + // 单独限制，否则小体积 deflate 炸弹可以放大成 GB 级内存。 + openAIResponsesWSMaxMessageBytes = 32 << 20 ) +var errOpenAIResponsesWSMessageTooLarge = fmt.Errorf("upstream websocket message exceeds %d bytes", openAIResponsesWSMaxMessageBytes) + // openAIResponsesRelayUsesUpstreamWebsocket 判断这一次出站是否拨中转账号自己的 // Responses WebSocket。生图仍走 HTTP。全局 codex_force_websocket 不参与这个判断。 func openAIResponsesRelayUsesUpstreamWebsocket(account *auth.Account, body []byte) bool { @@ -321,12 +326,19 @@ func openAIResponsesWSStream(ctx context.Context, pool *openAIResponsesWSPool, c } func readOpenAIResponsesWS(conn *websocket.Conn, emit func([]byte) bool) error { - conn.SetReadLimit(32 << 20) + conn.SetReadLimit(openAIResponsesWSMaxMessageBytes) for { - _, data, err := conn.ReadMessage() + _, reader, err := conn.NextReader() if err != nil { return err } + data, err := io.ReadAll(io.LimitReader(reader, openAIResponsesWSMaxMessageBytes+1)) + if err != nil { + return err + } + if len(data) > openAIResponsesWSMaxMessageBytes { + return errOpenAIResponsesWSMessageTooLarge + } if len(bytes.TrimSpace(data)) == 0 { continue } diff --git a/proxy/openai_responses_ws_limit_test.go b/proxy/openai_responses_ws_limit_test.go new file mode 100644 index 000000000..364e3853a --- /dev/null +++ b/proxy/openai_responses_ws_limit_test.go @@ -0,0 +1,64 @@ +package proxy + +import ( + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/gorilla/websocket" +) + +func dialCompressedWSTestServer(t *testing.T, payload string) (*websocket.Conn, func()) { + t.Helper() + upgrader := websocket.Upgrader{ + CheckOrigin: func(*http.Request) bool { return true }, + EnableCompression: true, + } + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + return + } + defer conn.Close() + _ = conn.WriteMessage(websocket.TextMessage, []byte(payload)) + })) + dialer := websocket.Dialer{EnableCompression: true} + conn, _, err := dialer.Dial("ws"+strings.TrimPrefix(server.URL, "http"), nil) + if err != nil { + server.Close() + t.Fatalf("dial: %v", err) + } + return conn, func() { + _ = conn.Close() + server.Close() + } +} + +func TestReadOpenAIResponsesWSRejectsDecompressedOversize(t *testing.T) { + payload := strings.Repeat("a", openAIResponsesWSMaxMessageBytes+1024) + conn, cleanup := dialCompressedWSTestServer(t, payload) + defer cleanup() + + err := readOpenAIResponsesWS(conn, func([]byte) bool { return true }) + if err != errOpenAIResponsesWSMessageTooLarge { + t.Fatalf("err = %v, want errOpenAIResponsesWSMessageTooLarge", err) + } +} + +func TestReadOpenAIResponsesWSReplaysNormalMessage(t *testing.T) { + conn, cleanup := dialCompressedWSTestServer(t, "hello") + defer cleanup() + + var got string + err := readOpenAIResponsesWS(conn, func(data []byte) bool { + got = string(data) + return false + }) + if err != nil { + t.Fatalf("read: %v", err) + } + if got != "hello" { + t.Fatalf("emit got %q, want hello", got) + } +} diff --git a/proxy/responses_ws.go b/proxy/responses_ws.go index e1be46e83..708a4f6ac 100644 --- a/proxy/responses_ws.go +++ b/proxy/responses_ws.go @@ -893,7 +893,7 @@ func (h *Handler) forwardResponsesWebSocketTurn(c *gin.Context, conn *websocket. if wsHTTPFallback.ForceHTTP() && !useWebsocket { wsHTTPFallback.LogHTTPAttemptCompletion("/v1/responses", account.ID(), attempt+1, durationMs, 0, resp.StatusCode) } - errBody, _ := io.ReadAll(resp.Body) + errBody, _ := io.ReadAll(io.LimitReader(resp.Body, maxUpstreamErrorBodyBytes)) rememberContinuousRetryHTTPFailure(c.Request.Context(), resp, errBody) resp.Body.Close() if c.Request.Context().Err() != nil { diff --git a/proxy/sse_line.go b/proxy/sse_line.go new file mode 100644 index 000000000..b33afa143 --- /dev/null +++ b/proxy/sse_line.go @@ -0,0 +1,27 @@ +package proxy + +import ( + "bufio" + "errors" +) + +// 上游 SSE 单行上限。正常事件远小于此值；超过说明上游发送了异常数据， +// 继续按 ReadBytes 累积会直接放大到整个行长的内存占用。 +const maxSSELineBytes = 8 << 20 + +var errSSELineTooLong = errors.New("upstream SSE line exceeds limit") + +// readBoundedSSELine 读取一行，内存占用不超过 limit（不含 bufio 自身缓冲）。 +func readBoundedSSELine(reader *bufio.Reader, limit int) ([]byte, error) { + var line []byte + for { + fragment, err := reader.ReadSlice('\n') + if len(line)+len(fragment) > limit { + return nil, errSSELineTooLong + } + line = append(line, fragment...) + if err != bufio.ErrBufferFull { + return line, err + } + } +} diff --git a/proxy/sse_line_test.go b/proxy/sse_line_test.go new file mode 100644 index 000000000..d573c4c4b --- /dev/null +++ b/proxy/sse_line_test.go @@ -0,0 +1,26 @@ +package proxy + +import ( + "bufio" + "strings" + "testing" +) + +func TestReadBoundedSSELine(t *testing.T) { + long := strings.Repeat("a", 10000) + "\n" + reader := bufio.NewReaderSize(strings.NewReader(long), 16) + line, err := readBoundedSSELine(reader, maxSSELineBytes) + if err != nil || len(line) != len(long) { + t.Fatalf("len=%d err=%v, want %d bytes", len(line), err, len(long)) + } + + reader = bufio.NewReaderSize(strings.NewReader("123456789\n"), 4) + if _, err := readBoundedSSELine(reader, 8); err != errSSELineTooLong { + t.Fatalf("err = %v, want errSSELineTooLong", err) + } + + reader = bufio.NewReaderSize(strings.NewReader("12345678\n"), 4) + if line, err := readBoundedSSELine(reader, 9); err != nil || string(line) != "12345678\n" { + t.Fatalf("line=%q err=%v", line, err) + } +} diff --git a/security/middleware.go b/security/middleware.go index c7d81a7a0..65f3f2d03 100644 --- a/security/middleware.go +++ b/security/middleware.go @@ -3,6 +3,8 @@ package security import ( "bytes" "crypto/rand" + "crypto/sha256" + "crypto/subtle" "encoding/hex" "errors" "fmt" @@ -148,14 +150,11 @@ func GenerateSecureToken(length int) (string, error) { // SecureCompare compares two strings in constant time to prevent timing attacks func SecureCompare(a, b string) bool { - if len(a) != len(b) { - return false - } - var result byte - for i := 0; i < len(a); i++ { - result |= a[i] ^ b[i] - } - return result == 0 + // Hash both values first so the comparison takes the same path regardless + // of length; a direct byte loop leaks the expected secret's length. + aHash := sha256.Sum256([]byte(a)) + bHash := sha256.Sum256([]byte(b)) + return subtle.ConstantTimeCompare(aHash[:], bHash[:]) == 1 } // SanitizeLog sanitizes data for logging @@ -163,6 +162,9 @@ func SanitizeLog(data string) string { if data == "" { return data } + // Keep one log event per line: URL paths and query values can contain + // escaped newlines that would otherwise forge additional audit lines. + data = strings.NewReplacer("\r", " ", "\n", " ").Replace(data) // Mask sensitive patterns data = MaskSensitiveData(data) // Truncate very long strings diff --git a/security/middleware_test.go b/security/middleware_test.go index 507a476c5..a9a3ccd83 100644 --- a/security/middleware_test.go +++ b/security/middleware_test.go @@ -231,3 +231,15 @@ func TestGenerateSecureToken(t *testing.T) { t.Errorf("GenerateSecureToken(16) length = %d, expected 32", len(token)) } } + +func TestSanitizeLogKeepsOneEventPerLine(t *testing.T) { + got := SanitizeLog("path=/v1/x%0aADMIN_ACCESS ip=203.0.113.9\r\nsource=env") + for _, bad := range []string{"\n", "\r"} { + if strings.Contains(got, bad) { + t.Fatalf("SanitizeLog left %q in %q", bad, got) + } + } + if !strings.Contains(got, "ADMIN_ACCESS") { + t.Fatalf("SanitizeLog dropped payload text: %q", got) + } +} diff --git a/security/validator.go b/security/validator.go index 7610f9dfa..b43181cc1 100644 --- a/security/validator.go +++ b/security/validator.go @@ -71,11 +71,16 @@ var ( regexp.MustCompile(`(?i)(refresh[_-]?token["']?\s*[:=]\s*["']?)([^"'\s&]+)`), regexp.MustCompile(`(?i)(access[_-]?token["']?\s*[:=]\s*["']?)([^"'\s&]+)`), regexp.MustCompile(`(?i)(id[_-]?token["']?\s*[:=]\s*["']?)([^"'\s&]+)`), + regexp.MustCompile(`(?i)(session[_-]?token["']?\s*[:=]\s*["']?)([^"'\s&]+)`), regexp.MustCompile(`(?i)(bearer\s+)([^\s&]+)`), regexp.MustCompile(`(?i)(["']?api[_-]?key["']?\s*[:=]\s*["']?)([^"'\s&]+)`), + regexp.MustCompile(`(?i)(["']?client[_-]?secret["']?\s*[:=]\s*["']?)([^"'\s&]+)`), regexp.MustCompile(`(?i)(["']?secret["']?\s*[:=]\s*["']?)([^"'\s&]+)`), regexp.MustCompile(`(?i)(["']?password["']?\s*[:=]\s*["']?)([^"'\s&]+)`), - regexp.MustCompile(`(?i)(sk-)([a-zA-Z0-9]{20,})`), // sk- 前缀单独作为组1 + regexp.MustCompile(`(?i)(["']?agent[_-]?private[_-]?key["']?\s*[:=]\s*["']?)([^"'\s&]+)`), + regexp.MustCompile(`(?i)(sk-)([a-zA-Z0-9]{20,})`), // sk- 前缀单独作为组1 + regexp.MustCompile(`\beyJ[A-Za-z0-9_-]{8,}\.[A-Za-z0-9_-]{8,}\.[A-Za-z0-9_-]{8,}\b`), // 裸 JWT + regexp.MustCompile(`-----BEGIN [A-Z ]*PRIVATE KEY-----[\s\S]*?-----END [A-Z ]*PRIVATE KEY-----`), } // UUID pattern for token masking diff --git a/security/validator_test.go b/security/validator_test.go index 00b702406..f4ffa9602 100644 --- a/security/validator_test.go +++ b/security/validator_test.go @@ -84,6 +84,11 @@ func TestMaskSensitiveData(t *testing.T) { {"hello world", "hello world"}, {"token=550e8400-e29b-41d4-a716-446655440000", "token=****UUID-MASKED****"}, {"http://user:secret@proxy.example:8080", "http://redacted@proxy.example:8080"}, + {"session_token=abc123", "session_token=****MASKED****"}, + {"client_secret: s3cr3t", "client_secret: ****MASKED****"}, + {"agent_private_key=abc", "agent_private_key=****MASKED****"}, + {"jwt eyJhbGciOiJIUzI1NiJ9.eyJzdWIiOiIxMjM0NTY3ODkwIn0.dozjgNryP4J3jVmNHl0w5N_XgL0n3I9PlFUP0THsR8U", "jwt ****MASKED****"}, + {"-----BEGIN RSA PRIVATE KEY-----\nMIIEow\n-----END RSA PRIVATE KEY-----", "****MASKED****"}, } for _, test := range tests { From 27dab23dd0411f291586be36c2a610ec79f31f9d Mon Sep 17 00:00:00 2001 From: Delicious233 <101502465+DeliciousBuding@users.noreply.github.com> Date: Sun, 18 Oct 2026 16:52:19 +0800 Subject: [PATCH 12/22] fix: address admission review feedback - Stop reserving RPM/RPD twice for batched image jobs: the enqueue path only checks the batch boundary, and each upstream call reserves one slot in the standard Images handler. Re-verified in WSL: 20 concurrent synchronous requests with rpm=2 gave 2 admitted and 18 429s. - Evict admin throttle and admission entries before inserting a new key/IP, so the maps stay bounded and a newly inserted key is not immediately dropped. - Return whether RPM or RPD was exceeded in the 429 message. --- admin/auth_throttle.go | 10 +++-- admin/auth_throttle_test.go | 12 +++++ proxy/apikey_admission.go | 17 +++---- proxy/apikey_admission_test.go | 82 ++++++++++++++++++++++++++++------ proxy/apikey_limits.go | 35 +++++++-------- 5 files changed, 113 insertions(+), 43 deletions(-) diff --git a/admin/auth_throttle.go b/admin/auth_throttle.go index 74b8dc828..8901327be 100644 --- a/admin/auth_throttle.go +++ b/admin/auth_throttle.go @@ -53,6 +53,9 @@ func (t *adminAuthThrottle) recordFailure(ip string, now time.Time) time.Duratio defer t.mu.Unlock() entry := t.entries[ip] if entry == nil || now.Sub(entry.windowStart) > adminAuthFailureWindow { + // 新来源也要触发清理：攻击者可以给大量 IP 各留 1–4 次失败， + // 只在达到阻断阈值时清理会让表无限增长。 + t.evictLocked(now) entry = &adminAuthThrottleEntry{windowStart: now} t.entries[ip] = entry } @@ -67,7 +70,6 @@ func (t *adminAuthThrottle) recordFailure(ip string, now time.Time) time.Duratio entry.failures = 0 entry.windowStart = now entry.blockedUntil = now.Add(adminAuthBlockDuration) - t.evictLocked(now) return adminAuthBlockDuration } @@ -80,9 +82,9 @@ func (t *adminAuthThrottle) reset(ip string) { t.mu.Unlock() } -// evictLocked 在条目过多时清理过期记录，避免伪造来源地址把内存撑大。 +// evictLocked 在写入新条目之前清理，保证插入后不超过上限。 func (t *adminAuthThrottle) evictLocked(now time.Time) { - if len(t.entries) <= adminAuthMaxEntries { + if len(t.entries) < adminAuthMaxEntries { return } for ip, entry := range t.entries { @@ -90,7 +92,7 @@ func (t *adminAuthThrottle) evictLocked(now time.Time) { delete(t.entries, ip) } } - for len(t.entries) > adminAuthMaxEntries { + for len(t.entries) >= adminAuthMaxEntries { var oldestIP string var oldest time.Time for ip, entry := range t.entries { diff --git a/admin/auth_throttle_test.go b/admin/auth_throttle_test.go index 4a4e97e6e..0f9b666b6 100644 --- a/admin/auth_throttle_test.go +++ b/admin/auth_throttle_test.go @@ -1,6 +1,7 @@ package admin import ( + "fmt" "net/http" "net/http/httptest" "testing" @@ -94,3 +95,14 @@ func TestAdminAuthMiddlewareThrottlesFailedAttempts(t *testing.T) { t.Fatalf("failure after successful login returned %d, want 401", rec.Code) } } + +func TestAdminAuthThrottleBoundsTrackedIPs(t *testing.T) { + throttle := newAdminAuthThrottle() + now := time.Unix(1_700_000_000, 0) + for i := 0; i < adminAuthMaxEntries+200; i++ { + throttle.recordFailure(fmt.Sprintf("198.51.100.%d", i), now) + } + if got := len(throttle.entries); got > adminAuthMaxEntries { + t.Fatalf("tracked IPs = %d, want <= %d", got, adminAuthMaxEntries) + } +} diff --git a/proxy/apikey_admission.go b/proxy/apikey_admission.go index e61e15d97..cdc7fca9d 100644 --- a/proxy/apikey_admission.go +++ b/proxy/apikey_admission.go @@ -44,27 +44,28 @@ func (h *Handler) admissionLimiter() *apiKeyAdmissionLimiter { } // allow 在 RPM/RPD 限额内为 requests 个上游请求预留计数。 -// limit <= 0 表示该维度不限制。 -func (l *apiKeyAdmissionLimiter) allow(keyID int64, rpm, rpd, requests int, now time.Time) bool { +// limit <= 0 表示该维度不限制。拒绝时返回 "rpm" 或 "rpd"。 +func (l *apiKeyAdmissionLimiter) allow(keyID int64, rpm, rpd, requests int, now time.Time) (bool, string) { if l == nil || keyID <= 0 || requests <= 0 { - return true + return true, "" } if rpm <= 0 && rpd <= 0 { - return true + return true, "" } l.mu.Lock() defer l.mu.Unlock() entry := l.entries[keyID] if entry == nil { + // 先清理再插入，否则满表时新键会被自己的淘汰逻辑删掉。 + l.evictLocked(now) entry = &apiKeyAdmissionEntry{} l.entries[keyID] = entry - l.evictLocked(now) } entry.lastSeen = now if rpm > 0 && entry.rpmUsed(now)+int64(requests) > int64(rpm) { - return false + return false, "rpm" } if rpd > 0 { if entry.rpdStart.IsZero() || now.Sub(entry.rpdStart) >= apiKeyRPDWindow { @@ -72,7 +73,7 @@ func (l *apiKeyAdmissionLimiter) allow(keyID int64, rpm, rpd, requests int, now entry.rpdCount = 0 } if entry.rpdCount+int64(requests) > int64(rpd) { - return false + return false, "rpd" } } @@ -83,7 +84,7 @@ func (l *apiKeyAdmissionLimiter) allow(keyID int64, rpm, rpd, requests int, now if rpd > 0 { entry.rpdCount += int64(requests) } - return true + return true, "" } // rpmUsed 推进秒级桶并返回最近 60 秒的计数，不写入本次请求。 diff --git a/proxy/apikey_admission_test.go b/proxy/apikey_admission_test.go index b68cedd5f..bb51271f2 100644 --- a/proxy/apikey_admission_test.go +++ b/proxy/apikey_admission_test.go @@ -1,28 +1,37 @@ package proxy import ( + "net/http" + "path/filepath" "sync" "sync/atomic" "testing" "time" + + "github.com/codex2api/database" ) +func admissionAllow(l *apiKeyAdmissionLimiter, keyID int64, rpm, rpd, requests int, now time.Time) bool { + ok, _ := l.allow(keyID, rpm, rpd, requests, now) + return ok +} + func TestAPIKeyAdmissionLimiterRejectsBurst(t *testing.T) { limiter := newAPIKeyAdmissionLimiter() now := time.Unix(1_700_000_000, 0) for i := 0; i < 2; i++ { - if !limiter.allow(1, 2, 0, 1, now) { + if !admissionAllow(limiter, 1, 2, 0, 1, now) { t.Fatalf("request %d within RPM rejected", i+1) } } - if limiter.allow(1, 2, 0, 1, now) { + if admissionAllow(limiter, 1, 2, 0, 1, now) { t.Fatal("third request in the same second should exceed RPM=2") } - if !limiter.allow(2, 2, 0, 1, now) { + if !admissionAllow(limiter, 2, 2, 0, 1, now) { t.Fatal("other API key should not be affected") } - if !limiter.allow(1, 2, 0, 1, now.Add(time.Minute)) { + if !admissionAllow(limiter, 1, 2, 0, 1, now.Add(time.Minute)) { t.Fatal("request should be admitted after the RPM window rolls") } } @@ -36,7 +45,7 @@ func TestAPIKeyAdmissionLimiterConcurrentBurst(t *testing.T) { wg.Add(1) go func() { defer wg.Done() - if limiter.allow(1, 5, 0, 1, now) { + if admissionAllow(limiter, 1, 5, 0, 1, now) { allowed.Add(1) } }() @@ -50,13 +59,13 @@ func TestAPIKeyAdmissionLimiterConcurrentBurst(t *testing.T) { func TestAPIKeyAdmissionLimiterReservesBatchRequests(t *testing.T) { limiter := newAPIKeyAdmissionLimiter() now := time.Unix(1_700_000_000, 0) - if !limiter.allow(1, 10, 0, 4, now) { + if !admissionAllow(limiter, 1, 10, 0, 4, now) { t.Fatal("batch of 4 within RPM=10 rejected") } - if limiter.allow(1, 10, 0, 7, now) { + if admissionAllow(limiter, 1, 10, 0, 7, now) { t.Fatal("batch of 7 would exceed the remaining RPM budget") } - if !limiter.allow(1, 10, 0, 6, now) { + if !admissionAllow(limiter, 1, 10, 0, 6, now) { t.Fatal("batch of 6 should fit exactly") } } @@ -64,16 +73,63 @@ func TestAPIKeyAdmissionLimiterReservesBatchRequests(t *testing.T) { func TestAPIKeyAdmissionLimiterEnforcesRPD(t *testing.T) { limiter := newAPIKeyAdmissionLimiter() now := time.Unix(1_700_000_000, 0) - if !limiter.allow(1, 0, 2, 1, now) { + if !admissionAllow(limiter, 1, 0, 2, 1, now) { t.Fatal("first RPD request rejected") } - if !limiter.allow(1, 0, 2, 1, now.Add(time.Hour)) { + if !admissionAllow(limiter, 1, 0, 2, 1, now.Add(time.Hour)) { t.Fatal("second RPD request rejected") } - if limiter.allow(1, 0, 2, 1, now.Add(2*time.Hour)) { - t.Fatal("third RPD request should exceed RPD=2") + if ok, reason := limiter.allow(1, 0, 2, 1, now.Add(2*time.Hour)); ok || reason != "rpd" { + t.Fatalf("third RPD request = (%v, %q), want (false, rpd)", ok, reason) } - if !limiter.allow(1, 0, 2, 1, now.Add(apiKeyRPDWindow)) { + if !admissionAllow(limiter, 1, 0, 2, 1, now.Add(apiKeyRPDWindow)) { t.Fatal("RPD window should reset after 24h") } } + +func TestAPIKeyAdmissionLimiterKeepsNewKeyWhenFull(t *testing.T) { + limiter := newAPIKeyAdmissionLimiter() + now := time.Unix(1_700_000_000, 0) + for i := 1; i <= apiKeyAdmissionMaxKeys; i++ { + if !admissionAllow(limiter, int64(i), 1, 0, 1, now) { + t.Fatalf("key %d rejected while filling the map", i) + } + } + newKey := int64(apiKeyAdmissionMaxKeys + 1) + if !admissionAllow(limiter, newKey, 1, 0, 1, now) { + t.Fatal("new key rejected when the map is full") + } + if _, ok := limiter.entries[newKey]; !ok { + t.Fatal("new key was evicted immediately after insertion") + } + if got := len(limiter.entries); got > apiKeyAdmissionMaxKeys+1 { + t.Fatalf("entries = %d, want <= %d", got, apiKeyAdmissionMaxKeys+1) + } +} + +func TestBatchAdmissionDoesNotReserveSlotsTwice(t *testing.T) { + db, err := database.New("sqlite", filepath.Join(t.TempDir(), "admission.db")) + if err != nil { + t.Fatalf("database.New: %v", err) + } + t.Cleanup(func() { _ = db.Close() }) + h := &Handler{db: db} + limits := database.APIKeyLimits{RPM: 3} + c := newImageLimitCtx(t, "/v1/images/generations", `{"prompt":"cat"}`, limits) + + // The enqueue check must not consume RPM; each upstream call reserves one. + if status, msg := h.EnforceAPIKeyLimitsForRequests(c, "gpt-image-2", 3); status != 0 { + t.Fatalf("batch admission = %d (%s), want pass", status, msg) + } + if _, ok := h.admissionLimiter().entries[1]; ok { + t.Fatal("batch admission reserved RPM slots before execution") + } + for i := 0; i < 3; i++ { + if status, msg := h.EnforceAPIKeyLimits(c, "gpt-image-2"); status != 0 { + t.Fatalf("per-call admission %d = %d (%s), want pass", i+1, status, msg) + } + } + if status, _ := h.EnforceAPIKeyLimits(c, "gpt-image-2"); status != http.StatusTooManyRequests { + t.Fatalf("fourth per-call admission = %d, want 429", status) + } +} diff --git a/proxy/apikey_limits.go b/proxy/apikey_limits.go index e2553df03..a497e9f31 100644 --- a/proxy/apikey_limits.go +++ b/proxy/apikey_limits.go @@ -77,20 +77,20 @@ func (h *Handler) payloadRuleIdentity(c *gin.Context) *PayloadRuleIdentity { // EnforceAPIKeyLimits checks API key scoped model and rate/cost limits. func (h *Handler) EnforceAPIKeyLimits(c *gin.Context, model string) (int, string) { - return h.enforceAPIKeyLimitsN(c, model, 1) + return h.enforceAPIKeyLimitsChecks(c, model, true) } // EnforceAPIKeyLimitsForRequests performs the normal API-key checks and also // verifies that a batched request cannot cross the RPM/RPD boundary by issuing -// multiple upstream calls after a single admission check. The whole batch is -// reserved atomically so concurrent batches cannot overshoot the window. +// multiple upstream calls after a single admission check. +// +// It deliberately does not reserve the batch here: every generated image is an +// upstream request that re-enters the standard Images handler and reserves one +// RPM/RPD slot there. Reserving the batch here as well would charge 2N. func (h *Handler) EnforceAPIKeyLimitsForRequests(c *gin.Context, model string, requests int) (int, string) { if requests < 1 { requests = 1 } - // 先做批量边界检查，避免整批被拒绝时仍消耗一个 RPM 配额。 - // Each upstream call still runs the normal handler checks and records its own - // usage entry, so the single admission below only reserves the batch size. row := apiKeyRowFromContext(c) if row != nil && row.ID > 0 && requests > 1 { ctx, cancel := context.WithTimeout(c.Request.Context(), 2*time.Second) @@ -110,7 +110,7 @@ func (h *Handler) EnforceAPIKeyLimitsForRequests(c *gin.Context, model string, r } } } - return h.enforceAPIKeyLimitsN(c, model, requests) + return h.enforceAPIKeyLimitsChecks(c, model, false) } func apiKeyBatchWouldExceed(current int64, limit, requests int) bool { @@ -127,15 +127,12 @@ func apiKeyBatchWouldExceed(current int64, limit, requests int) bool { // - http.StatusForbidden (403): 模型不在白名单 / 在黑名单 // - http.StatusTooManyRequests (429): rpm/rpd/cost/token 任一窗口超额 func (h *Handler) enforceAPIKeyLimits(c *gin.Context, model string) (int, string) { - return h.enforceAPIKeyLimitsN(c, model, 1) + return h.enforceAPIKeyLimitsChecks(c, model, true) } -// enforceAPIKeyLimitsN 是准入检查的实际实现。requests 表示本次准入会消耗的 -// 上游请求数（批量生图一次可产生多个请求），用于原子预留 RPM/RPD 计数。 -func (h *Handler) enforceAPIKeyLimitsN(c *gin.Context, model string, requests int) (int, string) { - if requests < 1 { - requests = 1 - } +// enforceAPIKeyLimitsChecks 执行全部限额检查；reserve 为 false 时跳过进程内 +// RPM/RPD 计数，供批量入队复用（实际执行时由每个上游调用另行预留）。 +func (h *Handler) enforceAPIKeyLimitsChecks(c *gin.Context, model string, reserve bool) (int, string) { var daybreakStatus int var daybreakMessage string model, daybreakStatus, daybreakMessage = h.checkDaybreakRequest(c, model) @@ -273,13 +270,15 @@ func (h *Handler) enforceAPIKeyLimitsN(c *gin.Context, model string, requests in // 7. 进程内准入计数。数据库聚合 + 60s 缓存读到的是同一份旧值，无法约束 // 缓存窗口内的并发突发；这里按实际请求数原子递增后再放行。 - if !h.admissionLimiter().allow(row.ID, limits.RPM, limits.RPD, requests, time.Now()) { - if limits.RPM > 0 { + if reserve { + if allowed, reason := h.admissionLimiter().allow(row.ID, limits.RPM, limits.RPD, 1, time.Now()); !allowed { + if reason == "rpd" { + return http.StatusTooManyRequests, + fmt.Sprintf("API key rate limit exceeded: %d requests per day", limits.RPD) + } return http.StatusTooManyRequests, fmt.Sprintf("API key rate limit exceeded: %d requests per minute", limits.RPM) } - return http.StatusTooManyRequests, - fmt.Sprintf("API key rate limit exceeded: %d requests per day", limits.RPD) } return 0, "" From 3d0a1365e03542a498960e7b3be3c109444c7038 Mon Sep 17 00:00:00 2001 From: Delicious233 <101502465+DeliciousBuding@users.noreply.github.com> Date: Sun, 18 Oct 2026 16:58:36 +0800 Subject: [PATCH 13/22] test(admin): assert no-store on authenticated responses --- admin/auth_throttle_test.go | 2 ++ 1 file changed, 2 insertions(+) diff --git a/admin/auth_throttle_test.go b/admin/auth_throttle_test.go index 0f9b666b6..1b05626dd 100644 --- a/admin/auth_throttle_test.go +++ b/admin/auth_throttle_test.go @@ -90,6 +90,8 @@ func TestAdminAuthMiddlewareThrottlesFailedAttempts(t *testing.T) { } if rec := request("admin-secret"); rec.Code != http.StatusOK { t.Fatalf("correct key during block returned %d, want 200", rec.Code) + } else if got := rec.Header().Get("Cache-Control"); got != "no-store" { + t.Fatalf("admin response Cache-Control = %q, want no-store", got) } if rec := request("wrong-key"); rec.Code != http.StatusUnauthorized { t.Fatalf("failure after successful login returned %d, want 401", rec.Code) From 92da4c219d60adb832f6fe0027567f959817f32c Mon Sep 17 00:00:00 2001 From: Delicious233 <101502465+DeliciousBuding@users.noreply.github.com> Date: Sun, 18 Oct 2026 17:08:05 +0800 Subject: [PATCH 14/22] test(ui): wait for concurrency badge before measuring layout --- frontend/tests/intelligenceCheck.browser.test.mjs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/frontend/tests/intelligenceCheck.browser.test.mjs b/frontend/tests/intelligenceCheck.browser.test.mjs index 4c6c099b3..37ba30858 100644 --- a/frontend/tests/intelligenceCheck.browser.test.mjs +++ b/frontend/tests/intelligenceCheck.browser.test.mjs @@ -187,7 +187,9 @@ test('save errors keep drafts and failed requests restore the manual action',asy async function assertConcurrencyInline(scope) { const concurrency = scope.locator('span[title*="当前正在处理 2 个请求"]').first() const detect = scope.getByRole('button', {name: '立即检测', exact: true}) + await Promise.all([concurrency.waitFor(), detect.waitFor()]) const [counter, action] = await Promise.all([concurrency.boundingBox(), detect.boundingBox()]) + assert.ok(counter && action, 'concurrency and action must both be visible') assert.ok(Math.abs(counter.y + counter.height/2 - action.y - action.height/2) < 2, 'concurrency stays on the action row') } async function assertActionsFollowStatus(scope) { From 8a21315bc4ae08c8741e9d7c93a5d7486faadb6a Mon Sep 17 00:00:00 2001 From: Delicious233 <101502465+DeliciousBuding@users.noreply.github.com> Date: Sun, 18 Oct 2026 17:28:05 +0800 Subject: [PATCH 15/22] fix(oauth): escape callback page and reject replayed sessions - HTML-escape the callback title and message; both can carry database errors or token claim values. - Compare states in constant time across all sessions instead of returning on the first string match. - Reject a callback whose session already completed successfully, while keeping the session available for the existing polling endpoint. - Upgrade github.com/klauspost/compress to 1.18.7 (govulncheck reported GO-2026-5841 in 1.17.6; the affected path was not called). --- admin/oauth.go | 28 +++++++++++++++++++++----- admin/oauth_test.go | 48 +++++++++++++++++++++++++++++++++++++++++++++ go.mod | 2 +- go.sum | 4 ++-- 4 files changed, 74 insertions(+), 8 deletions(-) diff --git a/admin/oauth.go b/admin/oauth.go index 9fd35d950..0b772ca2f 100644 --- a/admin/oauth.go +++ b/admin/oauth.go @@ -10,6 +10,7 @@ import ( "encoding/json" "errors" "fmt" + "html" "io" "log" "net/http" @@ -98,16 +99,26 @@ func (s *oauthSessionStore) delete(id string) { s.mu.Unlock() } -// findByState 通过 state 查找 session（回调端点使用，返回 sessionID + session） +// findByState 通过 state 查找 session（回调端点使用，返回 sessionID + session）。 +// 遍历全部条目并使用常量时间比较，避免按匹配位置提前返回。 func (s *oauthSessionStore) findByState(state string) (string, *oauthSession, bool) { s.mu.Lock() defer s.mu.Unlock() + var foundID string + var foundSess *oauthSession for id, sess := range s.sessions { - if sess.State == state && time.Since(sess.CreatedAt) <= oauthSessionTTL { - return id, sess, true + if !security.SecureCompare(sess.State, state) { + continue } + if time.Since(sess.CreatedAt) > oauthSessionTTL { + continue + } + foundID, foundSess = id, sess + } + if foundSess == nil { + return "", nil, false } - return "", nil, false + return foundID, foundSess, true } func (s *oauthSessionStore) cleanupLoop() { @@ -741,6 +752,10 @@ func (h *Handler) OAuthCallback(c *gin.Context) { c.String(http.StatusBadRequest, oauthCallbackPage("授权失败", "OAuth 会话不存在或已过期，请重新发起授权", false)) return } + if sess.ExchangeResult != nil && sess.ExchangeResult.Success { + c.String(http.StatusBadRequest, oauthCallbackPage("授权失败", "该授权回调已处理，请重新发起授权", false)) + return + } // 记录回调信息 sess.CallbackCode = code @@ -918,7 +933,8 @@ func (h *Handler) PollOAuthCallback(c *gin.Context) { }) } -// oauthCallbackPage 生成简单的 HTML 回调结果页面 +// oauthCallbackPage 生成简单的 HTML 回调结果页面。 +// title/message 可能来自数据库错误或 token claims，必须转义后再插入。 func oauthCallbackPage(title, message string, success bool) string { color := "#e53e3e" icon := "❌" @@ -926,6 +942,8 @@ func oauthCallbackPage(title, message string, success bool) string { color = "#38a169" icon = "✔" } + title = html.EscapeString(title) + message = html.EscapeString(message) return fmt.Sprintf(` %s