fix(routes): release body-admission slot after decode, not whole handler

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.
This commit is contained in:
2026-08-22 14:49:57 +01:00
parent 60036c43df
commit c69210fd91
4 changed files with 425 additions and 146 deletions
+49 -28
View File
@@ -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