From f736b231a3e7ce8f2b8cc991c95f385e9ca0a5c2 Mon Sep 17 00:00:00 2001 From: Lukasz Raczylo Date: Sat, 22 Aug 2026 13:43:37 +0100 Subject: [PATCH] fix(routes): admit before decode, bound in-flight bodies, narrow unified cap MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Close a HIGH-severity OOM path on the unified and media JSON routes (security review finding 1). Three parts: - checkAndCount now runs via a new admitRequest helper BEFORE the request body is read at all, on both routes_unified.go's runUnified and every routes_media.go handler. The rate-limit scopes need only the already-resolved user/group, so a caller already over their requestsPerMinute is refused without paying for the read+decode of a body that was always going to be discarded. Measured: an encoding/json decode into map[string]any costs ~11x the wire body in live heap, so the prior ordering let an already-limited caller burn ~110MiB per request before ever seeing the 429. - A new bounded, buffered-channel semaphore (Gateway.bodyAdmission, mirroring limits.go's spawnTokens shape) caps how many unified/media requests may concurrently hold a decoded body in memory. Acquired non-blockingly at each route's entry via acquireBodyAdmission; on exhaustion, returns 503 with Retry-After rather than blocking. Sized by defaultBodyAdmissionCap (GOMAXPROCS * 8, clamped to [16, 128]); Config.MaxInFlightBodyRequests, when set, always overrides the self-tuned default. - The unified chat/embeddings routes now decode against a new, narrower maxUnifiedRequestBytes (4MiB) instead of the shared maxRequestBytes (10MiB), which the media JSON routes, audio/transcriptions' multipart cap, and the MCP federation JSON-RPC routes keep unchanged. 1MiB (the audit's own first suggestion) was evaluated and rejected: this package documents and exercises inline base64 image_url content parts in chat messages (translate_anthropic.go, translate_gemini.go), and a single moderately-sized photo commonly exceeds 1MiB once base64-encoded. 4MiB preserves that documented single-inline-image case while still cutting worst-case per-request decode amplification from ~110MiB to ~44MiB. routes_media.go's resolveMediaRequest is split: model resolution stays (renamed resolveMediaModel), the limit check moves into the shared admitRequest. Tests: admission-before-decode is asserted directly (a 429'd request's body is never Read), semaphore exhaustion is asserted deterministically by filling the channel directly, the self-tune/override contract and its clamp bounds are pinned, and a zero-config run of ordinary traffic is asserted to see no new 503s. A dedicated test proves the chosen 4MiB cap still accepts a realistic ~1.5MiB inline image — the case 1MiB would have broken. --- llmgateway.go | 78 +++++++++- routes_media.go | 67 ++++++--- routes_unified.go | 185 +++++++++++++++++++----- routes_unified_test.go | 315 +++++++++++++++++++++++++++++++++++++++++ 4 files changed, 593 insertions(+), 52 deletions(-) diff --git a/llmgateway.go b/llmgateway.go index 6991b7a..056b9d4 100644 --- a/llmgateway.go +++ b/llmgateway.go @@ -9,6 +9,7 @@ import ( "fmt" "net/http" "os" + "runtime" ) // Config is the plugin's dynamic configuration, populated by Traefik from @@ -61,6 +62,22 @@ type Config struct { // above: its own Enabled field is the on/off signal. Cache CacheConfig `json:"cache"` PassthroughUnknown bool `json:"passthroughUnknown,omitempty"` + // MaxInFlightBodyRequests caps how many unified and media JSON/ + // multipart requests (routes_unified.go's runUnified; routes_media.go's + // handleImagesGenerations/handleAudioSpeech/handleAudioTranscriptions) + // may concurrently hold a decoded request body in memory at once + // (security review finding 1b, 2026-08-22) — see + // acquireBodyAdmission's own doc comment (routes_unified.go) for the + // OOM this bounds. 0 (the default) self-tunes from GOMAXPROCS + // (defaultBodyAdmissionCap) rather than adding an operator knob for a + // value the runtime can already infer (house rule: prefer self-tuning + // over operator knobs where the right value can be inferred; add a + // flag only when it genuinely cannot be). A positive value pins an + // exact cap instead and always wins over the auto-tuned default + // (house rule: an explicit override always wins) — e.g. a deployment + // with a known, tighter memory budget than GOMAXPROCS alone would + // infer. + MaxInFlightBodyRequests int `json:"maxInFlightBodyRequests,omitempty"` } // ProviderConfig describes one upstream LLM provider. @@ -303,6 +320,45 @@ func CreateConfig() *Config { return &Config{} } +// minBodyAdmissionCap/maxBodyAdmissionCap clamp defaultBodyAdmissionCap's +// self-tuned result (security review finding 1b, 2026-08-22): the floor +// keeps a single-vCPU sidecar from self-tuning down to a cap so small it +// throttles ordinary traffic, and the ceiling keeps a very large host from +// self-tuning up to a cap so high the semaphore stops bounding anything — +// each in-flight request can hold tens of MiB (maxUnifiedRequestBytes's +// own doc comment, routes_unified.go), so even the ceiling still bounds +// worst-case memory to a fixed, known order of magnitude. +const ( + minBodyAdmissionCap = 16 + maxBodyAdmissionCap = 128 +) + +// bodyAdmissionCapPerCPU is defaultBodyAdmissionCap's GOMAXPROCS +// multiplier: generous headroom for a body-buffering request, which spends +// most of its lifetime waiting on network I/O (upstream latency) rather +// than burning CPU, so a cap tied 1:1 to CPU count would leave real +// concurrency on the table. +const bodyAdmissionCapPerCPU = 8 + +// defaultBodyAdmissionCap returns the self-tuned in-flight cap +// acquireBodyAdmission's semaphore is sized to when Config. +// MaxInFlightBodyRequests is left at 0 (security review finding 1b, +// 2026-08-22): GOMAXPROCS-derived rather than a single constant that would +// be wrong for both a 1-vCPU sidecar and a 32-vCPU ingress node alike +// (house rule: prefer self-tuning over operator knobs where the right +// value can be inferred from the environment), clamped to +// [minBodyAdmissionCap, maxBodyAdmissionCap]. +func defaultBodyAdmissionCap() int { + n := runtime.GOMAXPROCS(0) * bodyAdmissionCapPerCPU + if n < minBodyAdmissionCap { + return minBodyAdmissionCap + } + if n > maxBodyAdmissionCap { + return maxBodyAdmissionCap + } + return n +} + // Gateway is the Traefik middleware handler. type Gateway struct { next http.Handler @@ -329,7 +385,14 @@ type Gateway struct { // pooled connections through. nil when config.Redis is absent, same // as buildRedisClient's own nil-for-unconfigured contract. redisClient *respClient - name string + // bodyAdmission is the buffered-channel semaphore acquireBodyAdmission + // (routes_unified.go) claims from and releases: security review + // finding 1b, 2026-08-22. Sized once, here, by newGateway (see + // defaultBodyAdmissionCap/Config.MaxInFlightBodyRequests) — never + // resized afterward, matching a Go channel's own fixed-capacity + // contract. + bodyAdmission chan struct{} + name string } // New creates the middleware. NOTE: no tail call — Yaegi zeroes @@ -384,6 +447,19 @@ func newGateway(ctx context.Context, next http.Handler, config *Config, name str } g := &Gateway{next: next, name: name, cfg: config, auth: auth} + // bodyAdmissionCap: an explicit MaxInFlightBodyRequests always wins + // over the self-tuned default (house rule: explicit override always + // wins) — see Config.MaxInFlightBodyRequests' own doc comment. 0 (the + // field's zero value, and every config that predates this field) means + // "self-tune", never "cap at zero" — a channel with a zero buffer + // would reject every request outright, which is never the intent of + // an operator who simply never set this field. + bodyAdmissionCap := defaultBodyAdmissionCap() + if config.MaxInFlightBodyRequests > 0 { + bodyAdmissionCap = config.MaxInFlightBodyRequests + } + g.bodyAdmission = make(chan struct{}, bodyAdmissionCap) + if config.Users != nil && config.Users.File != "" { if err = attachUsersFile(auth, config.Users.File, g); err != nil { return nil, err diff --git a/routes_media.go b/routes_media.go index c05cb2a..633b2c2 100644 --- a/routes_media.go +++ b/routes_media.go @@ -55,24 +55,29 @@ func (g *Gateway) decodeMediaJSONRequest(sw *statusTrackingWriter, r *http.Reque return req, model, true } -// resolveMediaRequest runs the shared model-routing/authz/limits steps -// every media route needs before calling its adapter (spec §3's "auth → -// model routing (registry) → group authz → limits" — group authz is -// folded into modelRegistry.resolve itself, the same as -// routes_unified.go's runUnified). On success it returns the resolved -// adapter and the upstream model id to route to; ok reports whether the -// caller may proceed — a resolve or limit failure has already written its -// response to sw. -func (g *Gateway) resolveMediaRequest(sw *statusTrackingWriter, u *user, grp *group, requestedModel string) (adapter providerAdapter, upstreamModel string, ok bool) { +// resolveMediaModel runs the shared model-routing/authz step every media +// route needs before calling its adapter (spec §3's "auth → model routing +// (registry) → group authz" — group authz is folded into +// modelRegistry.resolve itself, the same as routes_unified.go's +// runUnified). On success it returns the resolved adapter and the +// upstream model id to route to; ok reports whether the caller may +// proceed — a resolve failure has already written its response to sw. +// +// This used to also run checkAndCount (a limit check), and was named +// resolveMediaRequest — security review finding 1a, 2026-08-22, split +// that out into admitRequest (routes_unified.go), called by every media +// handler BEFORE the request body is even read: checkAndCount needs only +// u and grp, both already known ahead of decode, so there was never a +// reason to make a rate-limited caller pay for reading and decoding a +// body that was always going to be discarded. See admitRequest's own doc +// comment for the full rationale, shared verbatim with runUnified's +// identical fix. +func (g *Gateway) resolveMediaModel(sw *statusTrackingWriter, grp *group, requestedModel string) (adapter providerAdapter, upstreamModel string, ok bool) { adapter, upstreamModel, _, err := g.registry.resolve(requestedModel, grp) if err != nil { writeModelResolveError(sw, err) return nil, "", false } - if violation := g.limiter.checkAndCount(withTotalScope(buildLimitScopes(u, grp))); violation != nil { - writeLimitViolation(sw, violation) - return nil, "", false - } return adapter, upstreamModel, true } @@ -81,16 +86,24 @@ func (g *Gateway) resolveMediaRequest(sw *statusTrackingWriter, u *user, grp *gr // Imagen (translate_gemini_images.go); anthropic always answers 501 // (provider_anthropic.go's imagesGeneration). Images are never cached // (spec §2) and never cost-accounted (spec §3) — only the request -// counters resolveMediaRequest's checkAndCount call already incremented -// move. +// counters admitRequest's checkAndCount call already incremented move. func (g *Gateway) handleImagesGenerations(w http.ResponseWriter, r *http.Request, u *user, grp *group) { sw := &statusTrackingWriter{ResponseWriter: w} + release, ok := g.acquireBodyAdmission(sw) + defer release() + if !ok { + return + } + if _, ok = g.admitRequest(sw, u, grp); !ok { + return + } + req, requestedModel, ok := g.decodeMediaJSONRequest(sw, r) if !ok { return } - adapter, upstreamModel, ok := g.resolveMediaRequest(sw, u, grp, requestedModel) + adapter, upstreamModel, ok := g.resolveMediaModel(sw, grp, requestedModel) if !ok { return } @@ -119,11 +132,20 @@ func (g *Gateway) handleImagesGenerations(w http.ResponseWriter, r *http.Request func (g *Gateway) handleAudioSpeech(w http.ResponseWriter, r *http.Request, u *user, grp *group) { sw := &statusTrackingWriter{ResponseWriter: w} + release, ok := g.acquireBodyAdmission(sw) + defer release() + if !ok { + return + } + if _, ok = g.admitRequest(sw, u, grp); !ok { + return + } + req, requestedModel, ok := g.decodeMediaJSONRequest(sw, r) if !ok { return } - adapter, upstreamModel, ok := g.resolveMediaRequest(sw, u, grp, requestedModel) + adapter, upstreamModel, ok := g.resolveMediaModel(sw, grp, requestedModel) if !ok { return } @@ -175,6 +197,15 @@ func (g *Gateway) handleAudioSpeech(w http.ResponseWriter, r *http.Request, u *u func (g *Gateway) handleAudioTranscriptions(w http.ResponseWriter, r *http.Request, u *user, grp *group) { sw := &statusTrackingWriter{ResponseWriter: w} + release, ok := g.acquireBodyAdmission(sw) + defer release() + if !ok { + return + } + if _, ok = g.admitRequest(sw, u, grp); !ok { + return + } + body, oversize, err := readCapped(r.Body, maxRequestBytes) if err != nil { writeOAIError(sw, http.StatusBadRequest, "invalid_request_error", "cannot read request body") @@ -222,7 +253,7 @@ func (g *Gateway) handleAudioTranscriptions(w http.ResponseWriter, r *http.Reque return } - adapter, upstreamModel, ok := g.resolveMediaRequest(sw, u, grp, requestedModel) + adapter, upstreamModel, ok := g.resolveMediaModel(sw, grp, requestedModel) if !ok { return } diff --git a/routes_unified.go b/routes_unified.go index 09d9678..8b32826 100644 --- a/routes_unified.go +++ b/routes_unified.go @@ -11,12 +11,49 @@ import ( "strconv" ) -// maxRequestBytes caps a client request body the unified routes decode -// into a map[string]any: 10MiB, generous for a chat or embeddings request -// (including an inline base64 image) while bounding memory against an -// oversized or malicious body. +// maxRequestBytes caps a client request body decoded into a +// map[string]any: 10MiB. As of the security review below, the unified +// chat/embeddings routes (runUnified) no longer use this constant +// directly — they use the narrower maxUnifiedRequestBytes instead. This +// constant remains unchanged, at its original value, for every OTHER +// caller that was never in this finding's scope: the media JSON routes' +// own decode (decodeMediaJSONRequest, routes_media.go — shared by +// images.generations and audio.speech), audio.transcriptions' multipart +// cap (readCapped, routes_media.go), and the MCP federation JSON-RPC +// routes (mcp_federation.go). Splitting the constant, rather than +// lowering this shared one, keeps those callers' documented body-size +// behavior (README's "Body limits" section) byte-for-byte unchanged. const maxRequestBytes = 10 << 20 +// maxUnifiedRequestBytes caps a client request body specifically for the +// unified chat/embeddings routes' own decode into map[string]any +// (security review finding 1c, 2026-08-22): 4MiB, down from the general +// maxRequestBytes (10MiB) above. The measured amplification — an +// encoding/json decode into map[string]any costs roughly 11x the wire +// body's own bytes in live heap — meant the original 10MiB cap alone let +// one request materialize ~110MiB just for this decode, before cacheKey's +// re-marshal, forwardJSON's 32MiB response buffer, and the cache capture +// writer are even counted (see acquireBodyAdmission's own doc comment, +// below, for how the semaphore and this cap now work together). +// +// 1MiB — the audit's own first suggestion — was evaluated and rejected: +// this constant's PRIOR doc comment (before this finding) said the +// original 10MiB was "generous for a chat or embeddings request +// (including an inline base64 image)", and that is a real, exercised +// code path, not a hypothetical one — translate_anthropic.go's +// imageContentPart and translate_gemini.go's own image handling both +// accept a "data:;base64," image_url content part inside a +// chat message, and a single moderately-sized photo commonly exceeds +// 1MiB once base64-encoded (roughly +33% over its raw bytes; a compressed +// phone photo alone is often 1-3MB raw). Capping at 1MiB would silently +// break that documented, tested capability for any real-world image. +// 4MiB keeps headroom for the single-inline-image case the original +// comment named, while still cutting the worst-case per-request decode +// amplification from ~110MiB to ~44MiB — a defensible middle ground, not +// the audit's suggested number, chosen because the smaller number would +// have broken a real feature this package ships. +const maxUnifiedRequestBytes = 4 << 20 + // gatewayAliasKey is the request map's internal-convention key carrying // the client-requested model id through to an adapter (ruling a, ALIAS // ECHO), distinct from "model" itself, which runUnified rewrites to the @@ -76,20 +113,46 @@ func (g *Gateway) handleEmbeddings(w http.ResponseWriter, r *http.Request, u *us }) } -// runUnified is the shared chat/embeddings pipeline: decode the request, -// resolve its model, enforce per-user and per-group limits, invoke the -// adapter via call, then account the resulting usage — even when call -// itself returned an error, so usage captured before a mid-stream failure -// still gets billed — before translating that error into a response. w is -// wrapped in its own statusTrackingWriter so a mid-stream adapter error -// (headers already sent) can be told apart from one that failed before any -// write. endpoint is cacheEndpointChat or cacheEndpointEmbeddings — one of -// cacheKey's key-material components (cache.go), so the two routes never -// collide into one cache entry. +// runUnified is the shared chat/embeddings pipeline: claim a body-admission +// slot, enforce per-user/per-group/total REQUEST-RATE limits, decode the +// request, resolve its model, invoke the adapter via call, then account +// the resulting usage — even when call itself returned an error, so usage +// captured before a mid-stream failure still gets billed — before +// translating that error into a response. w is wrapped in its own +// statusTrackingWriter so a mid-stream adapter error (headers already +// sent) can be told apart from one that failed before any write. endpoint +// is cacheEndpointChat or cacheEndpointEmbeddings — one of cacheKey's +// key-material components (cache.go), so the two routes never collide +// into one cache entry. +// +// ORDERING (security review finding 1a, 2026-08-22): admission and the +// rate-limit check both run BEFORE the body is ever read, not after — +// the previous ordering (decode, resolve model, THEN checkAndCount) meant +// a caller already over their requestsPerMinute paid the full cost of +// reading and json-decoding a body that was always going to be discarded. +// acquireBodyAdmission and admitRequest both need only u/grp, already +// available as this function's own parameters, so neither has any reason +// to wait for the body. Model resolution still runs AFTER decode — it +// genuinely needs the client's requested model id, which only exists once +// the body is parsed — so its own errors (unknown/denied model) are still +// reported after a successful admission+rate-check, exactly as before +// this fix; only the RATE-LIMIT check's position relative to the body +// moved. writeLimitViolation's own response body/headers are unchanged. func (g *Gateway) runUnified(w http.ResponseWriter, r *http.Request, u *user, grp *group, endpoint string, call adapterCall) { sw := &statusTrackingWriter{ResponseWriter: w} - body, err := io.ReadAll(io.LimitReader(r.Body, maxRequestBytes)) + release, ok := g.acquireBodyAdmission(sw) + defer release() + if !ok { + return + } + + scopes, ok := g.admitRequest(sw, u, grp) + if !ok { + return + } + + body, err := io.ReadAll(io.LimitReader(r.Body, maxUnifiedRequestBytes)) if err != nil { writeOAIError(sw, http.StatusBadRequest, "invalid_request_error", "cannot read request body") return @@ -114,12 +177,6 @@ func (g *Gateway) runUnified(w http.ResponseWriter, r *http.Request, u *user, gr return } - scopes := withTotalScope(buildLimitScopes(u, grp)) - if violation := g.limiter.checkAndCount(scopes); violation != nil { - writeLimitViolation(sw, violation) - return - } - req["model"] = upstreamModel // cacheable gates every cache step below on spec §2's scope: a @@ -270,22 +327,84 @@ func buildLimitScopes(u *user, grp *group) []limitScope { // withTotalScope returns scopes with the synthetic total scope // (totalScopeKind/totalScopeID, limits.go) appended — the single helper -// every metered route (runUnified above; resolveMediaRequest, -// routes_media.go; handlePassthrough, routes_passthrough.go; -// handleTargetProxy, mcp_a2a.go; handleMCPFederated, mcp_federation.go) -// calls around its own buildLimitScopes result, so none of them can forget -// it and none of them duplicate the scope literal. Deliberately not folded -// into buildLimitScopes itself: buildLimitScopes' own doc comment now -// explains why admin traffic needs no special-casing here at all — -// handleAdminAPI never calls either function, so there was never a real -// "admin traffic must not contribute to total" case for this split to -// guard against; the split is kept anyway because buildLimitScopes' -// {user, group} pair and the synthetic total scope are conceptually -// different additions, worth two names. +// every metered route (runUnified above, via admitRequest below; +// handleImagesGenerations/handleAudioSpeech/handleAudioTranscriptions, +// routes_media.go, via the same admitRequest; handlePassthrough, +// routes_passthrough.go; handleTargetProxy, mcp_a2a.go; handleMCPFederated, +// mcp_federation.go) calls around its own buildLimitScopes result, so none +// of them can forget it and none of them duplicate the scope literal. +// Deliberately not folded into buildLimitScopes itself: buildLimitScopes' +// own doc comment now explains why admin traffic needs no special-casing +// here at all — handleAdminAPI never calls either function, so there was +// never a real "admin traffic must not contribute to total" case for this +// split to guard against; the split is kept anyway because +// buildLimitScopes' {user, group} pair and the synthetic total scope are +// conceptually different additions, worth two names. func withTotalScope(scopes []limitScope) []limitScope { return append(scopes, limitScope{kind: totalScopeKind, id: totalScopeID, limits: nil}) } +// admitRequest enforces per-user/per-group/total REQUEST-RATE limits +// (checkAndCount) using only u and grp — never a request body — so every +// caller can, and now does, call this before reading or decoding +// anything (security review finding 1a, 2026-08-22): a caller already +// over budget is refused on the strength of who they are alone, never +// after paying the cost of reading and json-decoding a body that turns +// out to be discarded anyway. scopes is returned alongside ok so the +// caller can reuse the identical scope slice for its own later account +// call, rather than rebuilding (and risking scope-list drift from) a +// second buildLimitScopes/withTotalScope pair. A violation writes its +// response to sw itself (writeLimitViolation, unchanged body/headers from +// before this fix) and returns ok=false. +func (g *Gateway) admitRequest(sw *statusTrackingWriter, u *user, grp *group) (scopes []limitScope, ok bool) { + scopes = withTotalScope(buildLimitScopes(u, grp)) + if violation := g.limiter.checkAndCount(scopes); violation != nil { + writeLimitViolation(sw, violation) + return scopes, false + } + return scopes, true +} + +// bodyAdmissionRetryAfterSeconds is the Retry-After value +// acquireBodyAdmission writes on a 503: a short, fixed window rather than +// a computed one — the semaphore's own occupancy has no natural "when +// will a slot free" answer the way a rate-limit window does (unlike +// writeLimitViolation's retryAfter, which names a real window boundary), +// and the in-flight cap is expected to drain within single-digit seconds +// under normal request latency. +const bodyAdmissionRetryAfterSeconds = 2 + +// acquireBodyAdmission non-blockingly claims one of g.bodyAdmission's +// slots (security review finding 1b, 2026-08-22) — bounding how many +// unified/media requests may concurrently hold a decoded request body (and +// every downstream buffer it feeds: cacheKey's re-marshal, forwardJSON's +// response buffer, the cache capture writer) in memory at once, so an +// unbounded burst of concurrent requests cannot OOM the whole shared +// Traefik ingress process this plugin runs inside (package doc, +// llmgateway.go) merely by arriving faster than any one of them can be +// rejected. Mirrors limiter.spawnTokens' own non-blocking acquire shape +// (limits.go) exactly: on exhaustion this NEVER queues or blocks waiting +// for a slot — a blocking acquire would just replace one unbounded-growth +// vector (unbounded concurrent decodes) with another (an unbounded pile of +// goroutines blocked on a channel receive) — it fails immediately instead, +// writing a 503 with a Retry-After header to sw and returning a no-op +// release so the caller's own `defer release()` stays valid either way. +// +// release must be deferred by the caller immediately after this returns, +// BEFORE checking ok — every code path through the caller, including this +// function's own 503 branch, must release exactly the number of slots it +// acquired (zero, on the 503 path, via the no-op release). +func (g *Gateway) acquireBodyAdmission(sw *statusTrackingWriter) (release func(), ok bool) { + select { + case g.bodyAdmission <- struct{}{}: + return func() { <-g.bodyAdmission }, true + default: + sw.Header().Set("Retry-After", strconv.Itoa(bodyAdmissionRetryAfterSeconds)) + writeOAIError(sw, http.StatusServiceUnavailable, "server_error", "server is at capacity; try again shortly") + return func() {}, false + } +} + // unifiedCostMicros resolves the price to charge one request's usage // against: canonical ("provider/model") first, falling back to bare — the // upstream model id alone — only when canonical has no configured price at diff --git a/routes_unified_test.go b/routes_unified_test.go index 936bd9f..514d348 100644 --- a/routes_unified_test.go +++ b/routes_unified_test.go @@ -1898,3 +1898,318 @@ func TestHandleChat_ContextDeadlineExceeded_RecordsProviderFailure(t *testing.T) t.Errorf("provider fails/day = %d, want 1 — a real context-deadline timeout must count as a provider-health failure (SHOULD-1)", fails) } } + +// --- security review finding 1, 2026-08-22: admission before decode, +// body-admission semaphore, narrower unified body cap --- + +// readTrackingBody is an io.ReadCloser that records whether Read was ever +// invoked, wrapping a real reader so a request that DOES get read still +// decodes normally. Used to prove a request rejected before the body is +// read (finding 1a) never actually touches it. +type readTrackingBody struct { + r io.Reader + called bool +} + +func (b *readTrackingBody) Read(p []byte) (int, error) { + b.called = true + return b.r.Read(p) +} + +func (b *readTrackingBody) Close() error { return nil } + +// TestHandleChat_RequestLimitExceeded_BodyNeverRead is the GATE's own +// explicit requirement: a 429'd user's body must never be read at all, +// proving admitRequest now runs strictly before io.ReadAll in runUnified +// (finding 1a) rather than merely returning the same status code for a +// different reason. +func TestHandleChat_RequestLimitExceeded_BodyNeverRead(t *testing.T) { + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{ + "openai": {Type: "openai", APIKey: "k", Models: []string{"gpt-test"}}, + } + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{ + {Name: "alice", Group: "default", APIKey: "sk-alice", Limits: &LimitsConfig{RequestsPerMinute: 1}}, + }} + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + h, err := New(context.Background(), next, cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + + // First request consumes the one-per-minute budget; its body is a + // real, valid body so the request completes normally end to end + // (upstream is unreachable — irrelevant here, only the request + // counter matters). + req1 := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", map[string]any{"model": "gpt-test", "messages": []any{}}) + rec1 := httptest.NewRecorder() + h.ServeHTTP(rec1, req1) + + // Second request: now over budget. Its body is wrapped in + // readTrackingBody — if runUnified reads it before checking the + // limit, called flips to true. + tracked := &readTrackingBody{r: strings.NewReader(`{"model":"gpt-test","messages":[]}`)} + req2 := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", tracked) + req2.Header.Set("Authorization", "Bearer sk-alice") + rec2 := httptest.NewRecorder() + h.ServeHTTP(rec2, req2) + + if rec2.Code != http.StatusTooManyRequests { + t.Fatalf("status = %d, want 429, body=%s", rec2.Code, rec2.Body.String()) + } + if tracked.called { + t.Error("a 429'd request's body must never be read — admitRequest must run before io.ReadAll") + } +} + +// TestHandleImagesGenerations_RequestLimitExceeded_BodyNeverRead mirrors +// the unified-route test above for the media JSON routes, which share the +// identical admission-before-decode ordering (finding 1a). +func TestHandleImagesGenerations_RequestLimitExceeded_BodyNeverRead(t *testing.T) { + cfg := newMediaTestConfig("http://127.0.0.1:1", "img-test") + cfg.Users = &UsersConfig{Inline: []*UserConfig{ + {Name: "alice", Group: "default", APIKey: "sk-alice", Limits: &LimitsConfig{RequestsPerMinute: 1}}, + }} + gw := newMediaTestGateway(t, cfg) + + req1 := newUnifiedRequest(t, http.MethodPost, imagesGenerationsPath, "sk-alice", map[string]any{"model": "img-test", "prompt": "a cat"}) + rec1 := httptest.NewRecorder() + gw.ServeHTTP(rec1, req1) + + tracked := &readTrackingBody{r: strings.NewReader(`{"model":"img-test","prompt":"a cat"}`)} + req2 := httptest.NewRequest(http.MethodPost, imagesGenerationsPath, tracked) + req2.Header.Set("Authorization", "Bearer sk-alice") + rec2 := httptest.NewRecorder() + gw.ServeHTTP(rec2, req2) + + if rec2.Code != http.StatusTooManyRequests { + t.Fatalf("status = %d, want 429, body=%s", rec2.Code, rec2.Body.String()) + } + if tracked.called { + t.Error("a 429'd media request's body must never be read — admitRequest must run before decode") + } +} + +// TestAcquireBodyAdmission_ExhaustionReturns503WithRetryAfter is the +// GATE's own explicit requirement: on semaphore exhaustion, +// acquireBodyAdmission must refuse immediately with 503 and a +// Retry-After header, never block. gw.bodyAdmission is filled directly +// (same package — this test manipulates the exact channel newGateway +// constructed) to make exhaustion deterministic without needing real +// concurrent goroutines. +func TestAcquireBodyAdmission_ExhaustionReturns503WithRetryAfter(t *testing.T) { + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k", Models: []string{"gpt-test"}}} + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} + cfg.MaxInFlightBodyRequests = 1 + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + h, err := New(context.Background(), next, cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + gw, ok := h.(*Gateway) + if !ok { + t.Fatal("handler is not *Gateway") + } + if cap(gw.bodyAdmission) != 1 { + t.Fatalf("bodyAdmission cap = %d, want 1 (explicit override)", cap(gw.bodyAdmission)) + } + + // Occupy the only slot directly, simulating another in-flight request. + gw.bodyAdmission <- struct{}{} + defer func() { <-gw.bodyAdmission }() + + req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", map[string]any{"model": "gpt-test", "messages": []any{}}) + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + + if rec.Code != http.StatusServiceUnavailable { + t.Fatalf("status = %d, want 503, body=%s", rec.Code, rec.Body.String()) + } + if ra := rec.Header().Get("Retry-After"); ra == "" { + t.Error("want a Retry-After header on 503, got none") + } +} + +// TestNewGateway_BodyAdmissionCap_SelfTunesFromGOMAXPROCS_ExplicitOverrideWins +// pins both halves of the house rule (self-tuning over operator knobs, +// but an explicit override always wins): a zero-config MaxInFlightBodyRequests +// self-tunes to defaultBodyAdmissionCap(), and a positive value overrides +// it exactly, regardless of what GOMAXPROCS would otherwise infer. +func TestNewGateway_BodyAdmissionCap_SelfTunesFromGOMAXPROCS_ExplicitOverrideWins(t *testing.T) { + baseCfg := func() *Config { + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}} + cfg.Groups = map[string]*GroupConfig{"default": {}} + return cfg + } + + t.Run("zero value self-tunes", func(t *testing.T) { + cfg := baseCfg() + h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + gw := h.(*Gateway) + if got, want := cap(gw.bodyAdmission), defaultBodyAdmissionCap(); got != want { + t.Errorf("bodyAdmission cap = %d, want %d (defaultBodyAdmissionCap)", got, want) + } + }) + + t.Run("explicit override wins", func(t *testing.T) { + cfg := baseCfg() + cfg.MaxInFlightBodyRequests = 7 + h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + gw := h.(*Gateway) + if got := cap(gw.bodyAdmission); got != 7 { + t.Errorf("bodyAdmission cap = %d, want 7 (explicit override)", got) + } + }) +} + +// TestDefaultBodyAdmissionCap_ClampedToBounds pins defaultBodyAdmissionCap's +// own contract directly: its result always lies within +// [minBodyAdmissionCap, maxBodyAdmissionCap], regardless of the host's own +// GOMAXPROCS. +func TestDefaultBodyAdmissionCap_ClampedToBounds(t *testing.T) { + got := defaultBodyAdmissionCap() + if got < minBodyAdmissionCap || got > maxBodyAdmissionCap { + t.Errorf("defaultBodyAdmissionCap() = %d, want within [%d, %d]", got, minBodyAdmissionCap, maxBodyAdmissionCap) + } +} + +// TestHandleChat_ZeroConfig_BodyAdmissionDoesNotThrottleOrdinaryTraffic is +// the DEFAULT-PRESERVING gate's own explicit requirement: a config that +// never sets MaxInFlightBodyRequests must behave exactly as before this +// finding for ordinary, non-bursty traffic — a modest run of sequential +// requests must never see a 503 from the new semaphore. +func TestHandleChat_ZeroConfig_BodyAdmissionDoesNotThrottleOrdinaryTraffic(t *testing.T) { + const respBody = `{"id":"c1","choices":[{"message":{"role":"assistant","content":"hi"}}],"usage":{"prompt_tokens":1,"completion_tokens":1}}` + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(respBody)) + })) + defer srv.Close() + + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{ + "openai": {Type: "openai", BaseURL: srv.URL, APIKey: "k", Models: []string{"gpt-test"}}, + } + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} + // Deliberately no MaxInFlightBodyRequests set — every config that + // predates this finding. + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + h, err := New(context.Background(), next, cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + + for i := 0; i < 20; i++ { + body := map[string]any{"model": "gpt-test", "messages": []any{}} + req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + if rec.Code != http.StatusOK { + t.Fatalf("request %d: status = %d, want 200, body=%s", i, rec.Code, rec.Body.String()) + } + } +} + +// TestHandleChat_OversizedBody_TruncatedByNarrowerUnifiedCap proves the +// unified routes now enforce maxUnifiedRequestBytes (4MiB), narrower than +// the general maxRequestBytes (10MiB) media/MCP-federation routes still +// use (finding 1c): a body whose valid JSON only closes past 4MiB is +// silently truncated by the LimitReader and fails to decode, exactly the +// existing "invalid JSON body" contract an oversized body has always hit +// — no new status code, just a lower threshold. +func TestHandleChat_OversizedBody_TruncatedByNarrowerUnifiedCap(t *testing.T) { + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k", Models: []string{"gpt-test"}}} + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + h, err := New(context.Background(), next, cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + + // A padding field alone bigger than maxUnifiedRequestBytes, so the + // LimitReader cap lands before the object closes. + padding := strings.Repeat("A", maxUnifiedRequestBytes+1024) + body := `{"model":"gpt-test","padding":"` + padding + `","messages":[]}` + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(body)) + req.Header.Set("Authorization", "Bearer sk-alice") + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + + if rec.Code != http.StatusBadRequest { + t.Fatalf("status = %d, want 400 (body silently truncated, invalid JSON), body=%s", rec.Code, rec.Body.String()) + } +} + +// TestHandleChat_UnderNewUnifiedCap_InlineBase64Image_StillWorks proves +// the chosen 4MiB unified cap (finding 1c) preserves the documented +// single-inline-base64-image use case maxUnifiedRequestBytes' own doc +// comment names: a chat request carrying a realistic ~1.5MiB base64 +// image_url content part, comfortably under 4MiB, must decode and reach +// the adapter exactly as before this finding — proving 1MiB (the audit's +// first suggestion, rejected in the doc comment) would have been too +// tight for this real feature, while the chosen number is not. +func TestHandleChat_UnderNewUnifiedCap_InlineBase64Image_StillWorks(t *testing.T) { + var gotBody []byte + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotBody, _ = io.ReadAll(r.Body) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{"id":"c1","choices":[{"message":{"role":"assistant","content":"ok"}}],"usage":{"prompt_tokens":1,"completion_tokens":1}}`)) + })) + defer srv.Close() + + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{ + "openai": {Type: "openai", BaseURL: srv.URL, APIKey: "k", Models: []string{"gpt-test"}}, + } + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + h, err := New(context.Background(), next, cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + + // ~1.5MiB of base64 payload — a realistic single inline photo, well + // under maxUnifiedRequestBytes (4MiB) and well OVER the audit's + // rejected 1MiB suggestion, proving that number would have broken + // this exact case. + imageB64 := strings.Repeat("A", 1_500_000) + body := map[string]any{ + "model": "gpt-test", + "messages": []any{ + map[string]any{ + "role": "user", + "content": []any{ + map[string]any{"type": "text", "text": "what is this?"}, + map[string]any{"type": "image_url", "image_url": map[string]any{"url": "data:image/jpeg;base64," + imageB64}}, + }, + }, + }, + } + req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body) + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want 200 (a ~1.5MiB inline image must fit under the new 4MiB cap), body=%s", rec.Code, rec.Body.String()) + } + if !strings.Contains(string(gotBody), imageB64) { + t.Error("upstream did not receive the full inline image payload — the cap silently truncated a request under its own limit") + } +}