From c69210fd91016851bd1a3b360d6c5bd87bb52d47 Mon Sep 17 00:00:00 2001 From: Lukasz Raczylo Date: Sat, 22 Aug 2026 14:33:29 +0100 Subject: [PATCH] fix(routes): release body-admission slot after decode, not whole handler MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fix an availability regression the round-2 body-admission semaphore shipped (MUST-1, round 3, 2026-08-22 review): `defer release()` spanned the WHOLE handler in runUnified and all three media handlers, including the upstream round trip and response streaming. The stated goal was bounding decoded-body memory during the read+json.Unmarshal step; the shipped code instead bounded total in-flight requests, since an LLM completion is I/O-bound for seconds to minutes. Measured, zero-config, 400 concurrent chat requests: 220 succeeded, 180 got 503, on a cap a 2-vCPU pod would size at 16. Narrowed the held window to read+decode only, via new helpers that do the acquire/decode/release internally and return before the caller ever resolves a model or calls the adapter: - routes_unified.go: readAndDecodeUnifiedBody (runUnified no longer acquires/defers at all). - routes_media.go: decodeMediaJSONRequest now claims/releases its own slot (images.generations, audio.speech); a new readAdmittedCapped helper does the same around audio.transcriptions' readCapped call — narrower still, since that route never json.Unmarshal's into a map[string]any at all. Every release still goes through `defer` inside these narrow helpers, not a bare post-call statement, so a slot is released exactly once on every path including a panic unwinding through the read/decode itself. Since the held window shrank from "whole request" (seconds-minutes) to "one decode" (low milliseconds), raised the self-tuned default: bodyAdmissionCapPerCPU 8->24, floor 16->32, ceiling 128->256 — justified in llmgateway.go's doc comment against the measured adversarial-shape decode peak (~51.5MB LIVE heap, ~189MB TotalAlloc for a 4MiB adversarial body), not the old, no-longer-accurate ~170MB whole-pipeline estimate. Also clamped an explicit Config.MaxInFlightBodyRequests override to 10,000 (previously accepted verbatim, so e.g. a fat-fingered MaxInt would have silently disabled the bound). Corrected maxUnifiedRequestBytes' own DoS-math doc comment (SHOULD-3): the prior "~11x general case" framing was wrong — a realistic single inline base64 image measures ~1.0x (no amplification at all), while the adversarial many-small-values shape is what actually amplifies (12.3x LIVE / ~45x TotalAlloc, both now stated with real measured numbers instead of a single misleading multiplier). Documented two real, previously-silent behavior changes from finding 1a's admission-before-decode reordering (SHOULD-4): an over-budget caller whose request would otherwise have 400'd/404'd now gets 429 instead, and a malformed/unroutable request now consumes rate-limit budget it previously did not. Both are arguably improvements but were not visible from the diff alone. Replaced the sequential zero-config admission test, which could not actually exercise contention (MUST-2: the reviewer reran its exact shape with MaxInFlightBodyRequests=1 and it still passed), with a genuinely concurrent one: n = defaultBodyAdmissionCap()+20 real goroutines against a gated upstream, asserting (a) zero 503s across all n and (b) the semaphore holds zero slots while all n are parked in the upstream -- the regression's own regression test. --- llmgateway.go | 117 +++++++++++++----- routes_media.go | 77 +++++++----- routes_unified.go | 267 +++++++++++++++++++++++++++++------------ routes_unified_test.go | 110 +++++++++++++++-- 4 files changed, 425 insertions(+), 146 deletions(-) diff --git a/llmgateway.go b/llmgateway.go index 056b9d4..cf8ce96 100644 --- a/llmgateway.go +++ b/llmgateway.go @@ -63,20 +63,25 @@ type Config struct { 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. + // multipart requests (routes_unified.go's runUnified via + // readAndDecodeUnifiedBody; routes_media.go's decodeMediaJSONRequest/ + // readAdmittedCapped, used by handleImagesGenerations/ + // handleAudioSpeech/handleAudioTranscriptions) may concurrently be + // INSIDE their own read+decode step at once (security review finding + // 1b, 2026-08-22; scope narrowed to read+decode only in round 3 — see + // acquireBodyAdmission's own doc comment, routes_unified.go, for why + // this is deliberately NOT a cap on total in-flight requests). 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 — but is still clamped to + // maxExplicitBodyAdmissionCap (llmgateway.go): "wins" means it + // overrides the self-tuned value, not that an operator can silently + // disable the bound entirely with an unbounded number. MaxInFlightBodyRequests int `json:"maxInFlightBodyRequests,omitempty"` } @@ -321,24 +326,60 @@ func CreateConfig() *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. +// self-tuned result (security review finding 1b; numbers revised round +// 3, 2026-08-22 — see this const block's own history below for why). +// 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. +// +// REVISED (round 3, coordinator ruling — fixes an availability +// regression the round-2 numbers shipped alongside): round 2 sized these +// against a "whole request" hold duration (acquireBodyAdmission was held +// for the entire handler, including the upstream round trip) and a +// ~170MB-per-slot estimate that bundled the decode, the re-marshal, AND +// the response buffer together. Measured, zero-config, at round 2's +// numbers: 400 concurrent chat requests against a 2-vCPU pod (cap 16) +// saw 220 succeed and 180 get 503 — an LLM completion is I/O-bound for +// seconds to minutes, so a cap sized for a brief decode became a hard +// ceiling on TOTAL in-flight requests instead. +// +// acquireBodyAdmission's slot is now held ONLY for the read+decode step +// itself (see its own doc comment, routes_unified.go) — typically +// single-digit milliseconds for a body under maxUnifiedRequestBytes +// (4MiB), a reduction of roughly three to four orders of magnitude in +// HOLD DURATION versus round 2. The relevant per-slot memory figure is +// also now decode-only, not decode+response: measured at up to ~51.5MB +// LIVE heap (12.3x) for an adversarial 4MiB body shaped to maximize +// per-value allocation overhead, with ~189MB TotalAlloc (~45x) of mostly +// short-lived garbage collected within the same decode +// (maxUnifiedRequestBytes' own doc comment, routes_unified.go, has the +// full measurement). At the new ceiling (256), a SUSTAINED worst case — +// every slot simultaneously decoding the adversarial shape, which the +// now-brief hold duration makes far less likely to persist than round +// 2's multi-second window — peaks around 256 × 51.5MB ≈ 13.2GB LIVE +// heap; the cumulative TotalAlloc figure (256 × 189MB ≈ 48.4GB) is +// garbage-collection churn, not simultaneously resident memory, and is +// not the number to size a pod's memory limit against. Both the floor +// and ceiling were doubled from round 2 (16→32, 128→256) and the +// per-CPU multiplier tripled (8→24): a raise smaller than the hold- +// duration reduction warrants, deliberately conservative, since the +// worst-case LIVE-heap arithmetic above still needs to stay defensible +// on a modest node. const ( - minBodyAdmissionCap = 16 - maxBodyAdmissionCap = 128 + minBodyAdmissionCap = 32 + maxBodyAdmissionCap = 256 ) // 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 +// multiplier (revised round 3 — see minBodyAdmissionCap's own doc +// comment for the full before/after reasoning): a decode-scoped slot is +// held for milliseconds, not the seconds-to-minutes an I/O-bound LLM +// completion takes, so — unlike round 2's comment, which reasoned about +// upstream-latency headroom for a hold spanning the whole request — this +// multiplier is sized against CPU throughput for a brief, CPU-bound +// json.Unmarshal burst, not I/O parking. +const bodyAdmissionCapPerCPU = 24 // defaultBodyAdmissionCap returns the self-tuned in-flight cap // acquireBodyAdmission's semaphore is sized to when Config. @@ -359,6 +400,18 @@ func defaultBodyAdmissionCap() int { return n } +// maxExplicitBodyAdmissionCap clamps an explicit Config. +// MaxInFlightBodyRequests override (round 3, 2026-08-22 review — the +// round-2 version accepted any positive value verbatim, including +// something like math.MaxInt from a fat-fingered config, which would +// silently disable the bound entirely by making the semaphore +// effectively unbounded). The house rule ("explicit override always +// wins") governs the choice between self-tuning and pinning a value, not +// whether a pinned value gets any sanity check at all — an operator who +// genuinely needs more concurrency than this still gets a very high +// ceiling, just not an unbounded one reachable by typo. +const maxExplicitBodyAdmissionCap = 10_000 + // Gateway is the Traefik middleware handler. type Gateway struct { next http.Handler @@ -453,10 +506,16 @@ func newGateway(ctx context.Context, next http.Handler, config *Config, name str // 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. + // an operator who simply never set this field. An explicit override + // is still clamped to maxExplicitBodyAdmissionCap (round 3, 2026-08-22 + // review) — "wins" means it overrides the self-tuned VALUE, not that + // it can disable the bound entirely via an unbounded number. bodyAdmissionCap := defaultBodyAdmissionCap() if config.MaxInFlightBodyRequests > 0 { bodyAdmissionCap = config.MaxInFlightBodyRequests + if bodyAdmissionCap > maxExplicitBodyAdmissionCap { + bodyAdmissionCap = maxExplicitBodyAdmissionCap + } } g.bodyAdmission = make(chan struct{}, bodyAdmissionCap) diff --git a/routes_media.go b/routes_media.go index 633b2c2..da26ed7 100644 --- a/routes_media.go +++ b/routes_media.go @@ -28,16 +28,26 @@ const ( // of data. const maxMultipartModelFieldBytes = 4096 -// decodeMediaJSONRequest reads r.Body capped at maxRequestBytes (routes_ -// unified.go's cap, shared here) and decodes it as a JSON object, then -// extracts its "model" field. It mirrors runUnified's own decode path -// exactly (spec §3's "reuse the unified decode path") — used by both the -// images and audio-speech routes, whose request bodies are both single -// JSON objects carrying a top-level "model" field, unlike -// audio-transcriptions' multipart body. A read failure or invalid/missing -// JSON is a 400, already written to sw; ok reports whether the caller may -// proceed. +// decodeMediaJSONRequest claims a body-admission slot (acquireBodyAdmission), +// reads r.Body capped at maxRequestBytes (routes_unified.go's cap, shared +// here) and decodes it as a JSON object, then extracts its "model" +// field — releasing the slot before returning, on every path (security +// review finding 1b, round 3, 2026-08-22: the slot is scoped to this +// read+decode step alone, exactly like readAndDecodeUnifiedBody +// (routes_unified.go), never held across the adapter call that follows +// in the caller). It mirrors runUnified's own decode path exactly (spec +// §3's "reuse the unified decode path") — used by both the images and +// audio-speech routes, whose request bodies are both single JSON objects +// carrying a top-level "model" field, unlike audio-transcriptions' +// multipart body. A read failure or invalid/missing JSON is a 400, +// already written to sw; ok reports whether the caller may proceed. func (g *Gateway) decodeMediaJSONRequest(sw *statusTrackingWriter, r *http.Request) (req map[string]any, model string, ok bool) { + release, admitted := g.acquireBodyAdmission(sw) + defer release() + if !admitted { + return nil, "", false + } + body, err := io.ReadAll(io.LimitReader(r.Body, maxRequestBytes)) if err != nil { writeOAIError(sw, http.StatusBadRequest, "invalid_request_error", "cannot read request body") @@ -90,12 +100,7 @@ func (g *Gateway) resolveMediaModel(sw *statusTrackingWriter, grp *group, reques 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 { + if _, ok := g.admitRequest(sw, u, grp); !ok { return } @@ -132,12 +137,7 @@ 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 { + if _, ok := g.admitRequest(sw, u, grp); !ok { return } @@ -197,16 +197,14 @@ 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 { + if _, ok := g.admitRequest(sw, u, grp); !ok { return } - body, oversize, err := readCapped(r.Body, maxRequestBytes) + body, oversize, admitted, err := g.readAdmittedCapped(sw, r.Body, maxRequestBytes) + if !admitted { + return + } if err != nil { writeOAIError(sw, http.StatusBadRequest, "invalid_request_error", "cannot read request body") return @@ -281,6 +279,29 @@ func (g *Gateway) handleAudioTranscriptions(w http.ResponseWriter, r *http.Reque } } +// readAdmittedCapped claims a body-admission slot (acquireBodyAdmission), +// calls readCapped, and releases the slot before returning — on every +// path, including a panic unwinding through readCapped itself — used by +// handleAudioTranscriptions (security review finding 1b, round 3, +// 2026-08-22). This route never json.Unmarshal's into a map[string]any +// (its own "decode", extractMultipartModel's bounded scan, runs over +// bytes already fully read, AFTER this returns and the slot is already +// free), so the slot only ever needs to be held for the read itself — +// narrower even than readAndDecodeUnifiedBody/decodeMediaJSONRequest, +// which both hold it through an actual json.Unmarshal too. admitted +// reports whether a slot was claimed at all (false means +// acquireBodyAdmission already wrote its own 503 to sw); body/oversize/ +// err are readCapped's own results, valid only when admitted is true. +func (g *Gateway) readAdmittedCapped(sw *statusTrackingWriter, r io.Reader, limit int64) (body []byte, oversize, admitted bool, err error) { + release, admitted := g.acquireBodyAdmission(sw) + defer release() + if !admitted { + return nil, false, false, nil + } + body, oversize, err = readCapped(r, limit) + return body, oversize, true, err +} + // readCapped reads up to limit+1 bytes from r, reporting oversize=true // when more than limit bytes were available. Unlike // decodeMediaJSONRequest's plain io.LimitReader-and-decode (which diff --git a/routes_unified.go b/routes_unified.go index 8b32826..cd03a08 100644 --- a/routes_unified.go +++ b/routes_unified.go @@ -28,30 +28,49 @@ 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). +// maxRequestBytes (10MiB) above. +// +// MEASURED AMPLIFICATION (corrected, round 3, 2026-08-22 review — the +// round-2 doc comment's "~11x" figure was presented as the general case +// and is not): amplification is SHAPE-DEPENDENT, not a fixed multiplier. +// A realistic single-inline-base64-image chat payload (the documented +// case this cap is sized for, below) measures at roughly 1.0x — a large +// base64 string decodes to one large Go string, no material blow-up. +// The adversarial shape — a body packed with many small values instead +// of one large one (e.g. a huge flat array/object of short strings or +// numbers, each becoming its own heap-allocated map entry/interface +// value) — is what actually amplifies: measured on a 4MiB adversarial +// array at 12.3x LIVE heap (~51.5MB resident at the decode's peak) and +// ~45x cumulative allocation (~189MB TotalAlloc, mostly short-lived +// garbage the collector reclaims quickly, not simultaneously resident). +// The bound this cap and acquireBodyAdmission's semaphore actually +// provide, honestly stated: one held decode's LIVE heap peaks at roughly +// 51.5MB in the adversarial case (worst case measured, not a +// theoretical 11x-of-4MiB ~44MiB figure); defaultBodyAdmissionCap's own +// doc comment (llmgateway.go) does the resulting worst-case-concurrent +// arithmetic against that real number, not this one's old estimate. // // 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 +// this constant's PRE-finding-1c doc comment 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. +// break that documented, tested capability for any real-world image — +// and, per the measurement above, that case shows ~1.0x amplification +// regardless, so the ORIGINAL amplification argument for shrinking this +// cap never actually applied to it; 1MiB was rejected for breaking a +// real feature, not because the realistic case was itself dangerous. // 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. +// comment named, while still bounding the adversarial shape's LIVE-heap +// footprint to a fixed, known-small amount per held decode — 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 @@ -113,54 +132,66 @@ func (g *Gateway) handleEmbeddings(w http.ResponseWriter, r *http.Request, u *us }) } -// 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. +// runUnified is the shared chat/embeddings pipeline: enforce per-user/ +// per-group/total REQUEST-RATE limits, claim a body-admission slot to +// 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. +// ORDERING (security review finding 1a, 2026-08-22): the rate-limit +// check runs 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. +// admitRequest needs only u/grp, already available as this function's +// own parameters, so it has no 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 — a +// 429's body is byte-identical to before this finding. +// +// BEHAVIOR CHANGE (documented, not hidden — SHOULD-4, round 3, +// 2026-08-22 review): moving the rate-limit check ahead of decode has +// two real, previously-silent side effects for a caller who is BOTH +// over budget AND sending a request that would otherwise have failed +// decode/model-resolution: +// - A request that would have gotten 400 (invalid JSON, missing +// "model") or 404 (unknown model) now gets 429 instead, if the +// caller was already over budget — the rate-limit check runs first +// and returns before decode/resolution ever gets a chance to run. +// - A malformed or unroutable request that would NOT have consumed any +// rate-limit budget under the old ordering (decode/resolve failed +// before checkAndCount ran) now DOES consume it, since checkAndCount +// always runs first regardless of what the body turns out to +// contain. A tenant with a buggy client that sends malformed JSON +// repeatedly could previously retry indefinitely for free; it now +// burns real requestsPerMinute/Day budget on every attempt. +// +// Both are consequences of fixing finding 1a's actual bug (an +// already-over-budget caller paying the read+decode cost) and are +// considered acceptable — arguably improvements, since a client that +// cannot even form a valid request is now itself rate-limited rather +// than getting unlimited free retries — but are called out explicitly +// here since neither is visible from the diff alone. func (g *Gateway) runUnified(w http.ResponseWriter, r *http.Request, u *user, grp *group, endpoint string, call adapterCall) { sw := &statusTrackingWriter{ResponseWriter: w} - 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 - } - - var req map[string]any - if err = json.Unmarshal(body, &req); err != nil { - writeOAIError(sw, http.StatusBadRequest, "invalid_request_error", "invalid JSON body") + body, req, ok := g.readAndDecodeUnifiedBody(sw, r) + if !ok { return } @@ -356,6 +387,13 @@ func withTotalScope(scopes []limitScope) []limitScope { // second buildLimitScopes/withTotalScope pair. A violation writes its // response to sw itself (writeLimitViolation, unchanged body/headers from // before this fix) and returns ok=false. +// +// Shared verbatim by every unified/media route (runUnified above; +// handleImagesGenerations/handleAudioSpeech/handleAudioTranscriptions, +// routes_media.go) — see runUnified's own "BEHAVIOR CHANGE" doc-comment +// section for the two previously-silent side effects this reordering has +// for a caller who is both over budget and sending a malformed/ +// unroutable request; they apply identically here. 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 { @@ -376,24 +414,69 @@ 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. +// unified/media requests may concurrently be INSIDE the read+json-decode +// step (readAndDecodeUnifiedBody, decodeMediaJSONRequest, the readCapped +// call in handleAudioTranscriptions), which is where the measured +// amplification actually happens (an encoding/json decode into +// map[string]any costs up to ~12x the wire body's own bytes in live +// heap for an adversarial shape — see maxUnifiedRequestBytes' own doc +// comment for the measured figures), so an unbounded burst of +// concurrent decodes 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. // -// 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). +// SCOPE (round 3, 2026-08-22, coordinator ruling — fixes an availability +// regression the round-2 version shipped): the slot is held ONLY for the +// read+decode step, never for the rest of the request — model +// resolution, the cache lookup, the upstream round trip, or streaming +// the response back to the client. Round 2 acquired the slot at handler +// entry and released it via a single `defer` spanning the WHOLE handler, +// which measurably turned this into a hard ceiling on TOTAL concurrent +// in-flight LLM requests rather than concurrent decodes: an LLM +// completion is I/O-bound for seconds to minutes, so holding a slot for +// that whole span meant the self-tuned default (previously GOMAXPROCS * +// 8) became the gateway's real-world concurrency limit — measured at +// zero-config, 400 concurrent chat requests: 220 succeeded, 180 got 503, +// on a cap a 2-vCPU pod would size at 16. A request parked waiting on an +// upstream response holds no slot now — every caller acquires and +// releases around ONLY its own read+decode call, not around the whole +// handler. The upstream request being built, sent, and its response +// streamed back are unguarded by this semaphore: that path is I/O-bound +// and does not re-amplify memory the way the decode step does, and the +// response side already has its own bound (forwardJSON/forwardStream's +// maxResponseBytes, 32MiB, providers.go — unchanged by this finding). +// Because the held window shrank from "whole request" (seconds-minutes) +// to "one decode" (low milliseconds), the default cap below is sized +// much higher than round 2's — see defaultBodyAdmissionCap's own doc +// comment for the reasoning and the resulting worst-case bound. +// +// Mirrors limiter.spawnTokens' own non-blocking acquire shape +// (limits.go): 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, and that defer must be scoped narrowly +// around the read+decode step alone (see readAndDecodeUnifiedBody/ +// decodeMediaJSONRequest for the pattern) — never around the caller's +// entire handler. Every code path, including this function's own 503 +// branch, releases exactly the number of slots it acquired (zero, on +// the 503 path, via the no-op release); a `defer` guarantees this holds +// even on a panic unwinding through the caller, since Go runs deferred +// functions during a panic regardless of where recover() (if any) is +// registered further up the same goroutine's stack. +// +// handlePassthrough (routes_passthrough.go) deliberately does NOT +// acquire a slot at all: it streams the request/response body straight +// through via io.Copy (proxyUpstream) rather than buffering it into a +// map[string]any, so there is no decode-amplification event here for +// this semaphore to bound — the only body-inspection it ever does is +// peekPassthroughModel's own bounded 64KiB read, an intentional +// omission, not an oversight. func (g *Gateway) acquireBodyAdmission(sw *statusTrackingWriter) (release func(), ok bool) { select { case g.bodyAdmission <- struct{}{}: @@ -405,6 +488,38 @@ func (g *Gateway) acquireBodyAdmission(sw *statusTrackingWriter) (release func() } } +// readAndDecodeUnifiedBody claims a body-admission slot (acquireBodyAdmission), +// reads r's body capped at maxUnifiedRequestBytes, decodes it as a JSON +// object, and releases the slot before returning — on every path: a +// successful decode, a read failure, a decode failure, or admission +// exhaustion itself (security review finding 1b, round 3, 2026-08-22). +// This is deliberately the ENTIRE scope of what the semaphore guards for +// the unified routes: by the time this returns, req is fully decoded and +// nothing further in runUnified re-amplifies memory the way the decode +// itself does, so model resolution, the cache lookup, and the upstream +// round trip all run after the slot is already released. A read or +// decode failure has already written its 400 to sw; ok reports whether +// the caller may proceed. +func (g *Gateway) readAndDecodeUnifiedBody(sw *statusTrackingWriter, r *http.Request) (body []byte, req map[string]any, ok bool) { + release, admitted := g.acquireBodyAdmission(sw) + defer release() + if !admitted { + return nil, nil, false + } + + body, err := io.ReadAll(io.LimitReader(r.Body, maxUnifiedRequestBytes)) + if err != nil { + writeOAIError(sw, http.StatusBadRequest, "invalid_request_error", "cannot read request body") + return nil, nil, false + } + + if err = json.Unmarshal(body, &req); err != nil { + writeOAIError(sw, http.StatusBadRequest, "invalid_request_error", "invalid JSON body") + return nil, nil, false + } + return body, req, true +} + // 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 514d348..e005a80 100644 --- a/routes_unified_test.go +++ b/routes_unified_test.go @@ -2083,14 +2083,51 @@ func TestDefaultBodyAdmissionCap_ClampedToBounds(t *testing.T) { } } -// 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) { +// TestHandleChat_ZeroConfig_ConcurrentAboveDefaultCap_NoThrottling_SlotsReleasedBeforeUpstream +// is the DEFAULT-PRESERVING gate's own explicit requirement, replacing a +// SEQUENTIAL version that could not actually exercise it (MUST-2, round +// 3, 2026-08-22 review: the reviewer re-ran the sequential test's exact +// shape with MaxInFlightBodyRequests=1 and it still passed — a +// capacity-one semaphore is never contended by requests that arrive one +// at a time, so it was not evidence the semaphore leaves ordinary +// traffic unthrottled). +// +// This drives n = defaultBodyAdmissionCap()+20 requests genuinely +// CONCURRENTLY (real goroutines, real network round trips to an +// httptest.Server upstream — n is deliberately larger than the +// self-tuned default cap) against an upstream gated to block until every +// one of them has arrived, and asserts two things: +// +// 1. While all n requests are parked waiting on the gated upstream — +// more in-flight requests than admission slots exist — the +// semaphore (gw.bodyAdmission) holds ZERO slots. This is the +// regression's own regression test (MUST-1, round 3): before that +// fix, a slot was held for the WHOLE handler including this exact +// upstream wait, so n > cap concurrent parked requests would have +// left the semaphore fully occupied (and, with a large enough n +// relative to cap, some later arrivals would have found it +// exhausted). With the fix, a slot is released the instant +// read+decode finishes, strictly before the upstream call ever +// starts — by the time a request reaches the gated upstream, it is +// provably holding no slot at all. +// 2. Once every parked request completes, all n get 200 — none was +// ever refused with 503 merely because n other requests happened to +// be concurrently in flight, which is exactly what "zero-config +// ordinary traffic is not throttled" has to mean on an I/O-bound +// gateway where many requests are routinely in flight for seconds +// at a time. +func TestHandleChat_ZeroConfig_ConcurrentAboveDefaultCap_NoThrottling_SlotsReleasedBeforeUpstream(t *testing.T) { const respBody = `{"id":"c1","choices":[{"message":{"role":"assistant","content":"hi"}}],"usage":{"prompt_tokens":1,"completion_tokens":1}}` + + // gate blocks every upstream response until the test explicitly + // closes it — simulating an LLM completion parked on a slow + // upstream for a while, the exact scenario the MUST-1 regression + // measured (a real completion is I/O-bound for seconds to minutes). + gate := make(chan struct{}) + var arrived sync.WaitGroup srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + arrived.Done() // this request's admission slot is already released by now — see below + <-gate w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusOK) _, _ = w.Write([]byte(respBody)) @@ -2110,14 +2147,61 @@ func TestHandleChat_ZeroConfig_BodyAdmissionDoesNotThrottleOrdinaryTraffic(t *te if err != nil { t.Fatalf("New: %v", err) } + gw, ok := h.(*Gateway) + if !ok { + t.Fatal("handler is not *Gateway") + } - 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()) + n := cap(gw.bodyAdmission) + 20 + arrived.Add(n) + + statuses := make([]int, n) + var wg sync.WaitGroup + wg.Add(n) + for i := 0; i < n; i++ { + go func(i int) { + defer wg.Done() + 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) + statuses[i] = rec.Code + }(i) + } + + // Blocks until every one of the n requests has reached the gated + // upstream — which, by this call chain's own synchronous ordering + // (runUnified -> readAndDecodeUnifiedBody's deferred release -> + // registry.resolve -> the adapter's own client.Do), can only happen + // AFTER that request's admission slot has already been released. + // Bounded by a timeout, not a bare Wait(): if the regression this + // test guards against reappeared, fewer than n requests would ever + // reach the upstream at all (the rest would be refused with an + // immediate 503 — acquireBodyAdmission never blocks, see its own doc + // comment — and never call arrived.Done()), which would hang a bare + // Wait() forever instead of failing the test. + arrivedCh := make(chan struct{}) + go func() { + arrived.Wait() + close(arrivedCh) + }() + select { + case <-arrivedCh: + case <-time.After(10 * time.Second): + close(gate) // unstick any goroutine that did make it through, so the test can still finish reporting statuses + t.Fatal("timed out waiting for all n requests to reach the gated upstream — some are likely being refused with 503 before ever reaching it, which is the regression this test guards against") + } + + if held := len(gw.bodyAdmission); held != 0 { + t.Errorf("body-admission slots held while %d requests (cap=%d) are parked in the upstream = %d, want 0 — a request parked past decode must never still hold a slot", n, cap(gw.bodyAdmission), held) + } + + close(gate) + wg.Wait() + + for i, status := range statuses { + if status != http.StatusOK { + t.Errorf("request %d: status = %d, want 200 (cap=%d, n=%d concurrent) — ordinary traffic must never see 503 from the body-admission semaphore merely for being concurrent", i, status, cap(gw.bodyAdmission), n) } } }