mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-18 05:47:38 +00:00
fix(routes): admit before decode, bound in-flight bodies, narrow unified cap
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.
This commit is contained in:
+49
-18
@@ -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
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user