mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-14 05:18:54 +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:
+77
-1
@@ -9,6 +9,7 @@ import (
|
|||||||
"fmt"
|
"fmt"
|
||||||
"net/http"
|
"net/http"
|
||||||
"os"
|
"os"
|
||||||
|
"runtime"
|
||||||
)
|
)
|
||||||
|
|
||||||
// Config is the plugin's dynamic configuration, populated by Traefik from
|
// 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.
|
// above: its own Enabled field is the on/off signal.
|
||||||
Cache CacheConfig `json:"cache"`
|
Cache CacheConfig `json:"cache"`
|
||||||
PassthroughUnknown bool `json:"passthroughUnknown,omitempty"`
|
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.
|
// ProviderConfig describes one upstream LLM provider.
|
||||||
@@ -303,6 +320,45 @@ func CreateConfig() *Config {
|
|||||||
return &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.
|
// Gateway is the Traefik middleware handler.
|
||||||
type Gateway struct {
|
type Gateway struct {
|
||||||
next http.Handler
|
next http.Handler
|
||||||
@@ -329,7 +385,14 @@ type Gateway struct {
|
|||||||
// pooled connections through. nil when config.Redis is absent, same
|
// pooled connections through. nil when config.Redis is absent, same
|
||||||
// as buildRedisClient's own nil-for-unconfigured contract.
|
// as buildRedisClient's own nil-for-unconfigured contract.
|
||||||
redisClient *respClient
|
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
|
// 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}
|
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 config.Users != nil && config.Users.File != "" {
|
||||||
if err = attachUsersFile(auth, config.Users.File, g); err != nil {
|
if err = attachUsersFile(auth, config.Users.File, g); err != nil {
|
||||||
return nil, err
|
return nil, err
|
||||||
|
|||||||
+49
-18
@@ -55,24 +55,29 @@ func (g *Gateway) decodeMediaJSONRequest(sw *statusTrackingWriter, r *http.Reque
|
|||||||
return req, model, true
|
return req, model, true
|
||||||
}
|
}
|
||||||
|
|
||||||
// resolveMediaRequest runs the shared model-routing/authz/limits steps
|
// resolveMediaModel runs the shared model-routing/authz step every media
|
||||||
// every media route needs before calling its adapter (spec §3's "auth →
|
// route needs before calling its adapter (spec §3's "auth → model routing
|
||||||
// model routing (registry) → group authz → limits" — group authz is
|
// (registry) → group authz" — group authz is folded into
|
||||||
// folded into modelRegistry.resolve itself, the same as
|
// modelRegistry.resolve itself, the same as routes_unified.go's
|
||||||
// routes_unified.go's runUnified). On success it returns the resolved
|
// runUnified). On success it returns the resolved adapter and the
|
||||||
// adapter and the upstream model id to route to; ok reports whether the
|
// upstream model id to route to; ok reports whether the caller may
|
||||||
// caller may proceed — a resolve or limit failure has already written its
|
// proceed — a resolve failure has already written its response to sw.
|
||||||
// response to sw.
|
//
|
||||||
func (g *Gateway) resolveMediaRequest(sw *statusTrackingWriter, u *user, grp *group, requestedModel string) (adapter providerAdapter, upstreamModel string, ok bool) {
|
// 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)
|
adapter, upstreamModel, _, err := g.registry.resolve(requestedModel, grp)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
writeModelResolveError(sw, err)
|
writeModelResolveError(sw, err)
|
||||||
return nil, "", false
|
return nil, "", false
|
||||||
}
|
}
|
||||||
if violation := g.limiter.checkAndCount(withTotalScope(buildLimitScopes(u, grp))); violation != nil {
|
|
||||||
writeLimitViolation(sw, violation)
|
|
||||||
return nil, "", false
|
|
||||||
}
|
|
||||||
return adapter, upstreamModel, true
|
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
|
// Imagen (translate_gemini_images.go); anthropic always answers 501
|
||||||
// (provider_anthropic.go's imagesGeneration). Images are never cached
|
// (provider_anthropic.go's imagesGeneration). Images are never cached
|
||||||
// (spec §2) and never cost-accounted (spec §3) — only the request
|
// (spec §2) and never cost-accounted (spec §3) — only the request
|
||||||
// counters resolveMediaRequest's checkAndCount call already incremented
|
// counters admitRequest's checkAndCount call already incremented move.
|
||||||
// move.
|
|
||||||
func (g *Gateway) handleImagesGenerations(w http.ResponseWriter, r *http.Request, u *user, grp *group) {
|
func (g *Gateway) handleImagesGenerations(w http.ResponseWriter, r *http.Request, u *user, grp *group) {
|
||||||
sw := &statusTrackingWriter{ResponseWriter: w}
|
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)
|
req, requestedModel, ok := g.decodeMediaJSONRequest(sw, r)
|
||||||
if !ok {
|
if !ok {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
adapter, upstreamModel, ok := g.resolveMediaRequest(sw, u, grp, requestedModel)
|
adapter, upstreamModel, ok := g.resolveMediaModel(sw, grp, requestedModel)
|
||||||
if !ok {
|
if !ok {
|
||||||
return
|
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) {
|
func (g *Gateway) handleAudioSpeech(w http.ResponseWriter, r *http.Request, u *user, grp *group) {
|
||||||
sw := &statusTrackingWriter{ResponseWriter: w}
|
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)
|
req, requestedModel, ok := g.decodeMediaJSONRequest(sw, r)
|
||||||
if !ok {
|
if !ok {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
adapter, upstreamModel, ok := g.resolveMediaRequest(sw, u, grp, requestedModel)
|
adapter, upstreamModel, ok := g.resolveMediaModel(sw, grp, requestedModel)
|
||||||
if !ok {
|
if !ok {
|
||||||
return
|
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) {
|
func (g *Gateway) handleAudioTranscriptions(w http.ResponseWriter, r *http.Request, u *user, grp *group) {
|
||||||
sw := &statusTrackingWriter{ResponseWriter: w}
|
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)
|
body, oversize, err := readCapped(r.Body, maxRequestBytes)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
writeOAIError(sw, http.StatusBadRequest, "invalid_request_error", "cannot read request body")
|
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
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
adapter, upstreamModel, ok := g.resolveMediaRequest(sw, u, grp, requestedModel)
|
adapter, upstreamModel, ok := g.resolveMediaModel(sw, grp, requestedModel)
|
||||||
if !ok {
|
if !ok {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|||||||
+152
-33
@@ -11,12 +11,49 @@ import (
|
|||||||
"strconv"
|
"strconv"
|
||||||
)
|
)
|
||||||
|
|
||||||
// maxRequestBytes caps a client request body the unified routes decode
|
// maxRequestBytes caps a client request body decoded into a
|
||||||
// into a map[string]any: 10MiB, generous for a chat or embeddings request
|
// map[string]any: 10MiB. As of the security review below, the unified
|
||||||
// (including an inline base64 image) while bounding memory against an
|
// chat/embeddings routes (runUnified) no longer use this constant
|
||||||
// oversized or malicious body.
|
// 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
|
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
|
// gatewayAliasKey is the request map's internal-convention key carrying
|
||||||
// the client-requested model id through to an adapter (ruling a, ALIAS
|
// the client-requested model id through to an adapter (ruling a, ALIAS
|
||||||
// ECHO), distinct from "model" itself, which runUnified rewrites to the
|
// 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,
|
// runUnified is the shared chat/embeddings pipeline: claim a body-admission
|
||||||
// resolve its model, enforce per-user and per-group limits, invoke the
|
// slot, enforce per-user/per-group/total REQUEST-RATE limits, decode the
|
||||||
// adapter via call, then account the resulting usage — even when call
|
// request, resolve its model, invoke the adapter via call, then account
|
||||||
// itself returned an error, so usage captured before a mid-stream failure
|
// the resulting usage — even when call itself returned an error, so usage
|
||||||
// still gets billed — before translating that error into a response. w is
|
// captured before a mid-stream failure still gets billed — before
|
||||||
// wrapped in its own statusTrackingWriter so a mid-stream adapter error
|
// translating that error into a response. w is wrapped in its own
|
||||||
// (headers already sent) can be told apart from one that failed before any
|
// statusTrackingWriter so a mid-stream adapter error (headers already
|
||||||
// write. endpoint is cacheEndpointChat or cacheEndpointEmbeddings — one of
|
// sent) can be told apart from one that failed before any write. endpoint
|
||||||
// cacheKey's key-material components (cache.go), so the two routes never
|
// is cacheEndpointChat or cacheEndpointEmbeddings — one of cacheKey's
|
||||||
// collide into one cache entry.
|
// 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) {
|
func (g *Gateway) runUnified(w http.ResponseWriter, r *http.Request, u *user, grp *group, endpoint string, call adapterCall) {
|
||||||
sw := &statusTrackingWriter{ResponseWriter: w}
|
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 {
|
if err != nil {
|
||||||
writeOAIError(sw, http.StatusBadRequest, "invalid_request_error", "cannot read request body")
|
writeOAIError(sw, http.StatusBadRequest, "invalid_request_error", "cannot read request body")
|
||||||
return
|
return
|
||||||
@@ -114,12 +177,6 @@ func (g *Gateway) runUnified(w http.ResponseWriter, r *http.Request, u *user, gr
|
|||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
scopes := withTotalScope(buildLimitScopes(u, grp))
|
|
||||||
if violation := g.limiter.checkAndCount(scopes); violation != nil {
|
|
||||||
writeLimitViolation(sw, violation)
|
|
||||||
return
|
|
||||||
}
|
|
||||||
|
|
||||||
req["model"] = upstreamModel
|
req["model"] = upstreamModel
|
||||||
|
|
||||||
// cacheable gates every cache step below on spec §2's scope: a
|
// 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
|
// withTotalScope returns scopes with the synthetic total scope
|
||||||
// (totalScopeKind/totalScopeID, limits.go) appended — the single helper
|
// (totalScopeKind/totalScopeID, limits.go) appended — the single helper
|
||||||
// every metered route (runUnified above; resolveMediaRequest,
|
// every metered route (runUnified above, via admitRequest below;
|
||||||
// routes_media.go; handlePassthrough, routes_passthrough.go;
|
// handleImagesGenerations/handleAudioSpeech/handleAudioTranscriptions,
|
||||||
// handleTargetProxy, mcp_a2a.go; handleMCPFederated, mcp_federation.go)
|
// routes_media.go, via the same admitRequest; handlePassthrough,
|
||||||
// calls around its own buildLimitScopes result, so none of them can forget
|
// routes_passthrough.go; handleTargetProxy, mcp_a2a.go; handleMCPFederated,
|
||||||
// it and none of them duplicate the scope literal. Deliberately not folded
|
// mcp_federation.go) calls around its own buildLimitScopes result, so none
|
||||||
// into buildLimitScopes itself: buildLimitScopes' own doc comment now
|
// of them can forget it and none of them duplicate the scope literal.
|
||||||
// explains why admin traffic needs no special-casing here at all —
|
// Deliberately not folded into buildLimitScopes itself: buildLimitScopes'
|
||||||
// handleAdminAPI never calls either function, so there was never a real
|
// own doc comment now explains why admin traffic needs no special-casing
|
||||||
// "admin traffic must not contribute to total" case for this split to
|
// here at all — handleAdminAPI never calls either function, so there was
|
||||||
// guard against; the split is kept anyway because buildLimitScopes'
|
// never a real "admin traffic must not contribute to total" case for this
|
||||||
// {user, group} pair and the synthetic total scope are conceptually
|
// split to guard against; the split is kept anyway because
|
||||||
// different additions, worth two names.
|
// buildLimitScopes' {user, group} pair and the synthetic total scope are
|
||||||
|
// conceptually different additions, worth two names.
|
||||||
func withTotalScope(scopes []limitScope) []limitScope {
|
func withTotalScope(scopes []limitScope) []limitScope {
|
||||||
return append(scopes, limitScope{kind: totalScopeKind, id: totalScopeID, limits: nil})
|
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
|
// unifiedCostMicros resolves the price to charge one request's usage
|
||||||
// against: canonical ("provider/model") first, falling back to bare — the
|
// against: canonical ("provider/model") first, falling back to bare — the
|
||||||
// upstream model id alone — only when canonical has no configured price at
|
// upstream model id alone — only when canonical has no configured price at
|
||||||
|
|||||||
@@ -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)
|
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")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user