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:
2026-08-22 14:49:40 +01:00
parent bf33732123
commit f736b231a3
4 changed files with 593 additions and 52 deletions
+77 -1
View File
@@ -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
+49 -18
View File
@@ -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
}
+152 -33
View File
@@ -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:<type>;base64,<data>" 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
+315
View File
@@ -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")
}
}