mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-21 05:57:53 +00:00
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:
+49
-28
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user