merge: federated MCP auth and fan-out hardening (audit round 2)

This commit is contained in:
2026-08-22 15:38:59 +01:00
9 changed files with 1868 additions and 82 deletions
+58 -11
View File
@@ -631,16 +631,63 @@ func usageEntryView(su scopeUsage, limits *LimitsConfig) adminUsageEntryView {
// buildAdminUsage assembles adminUsageResponse: authStore.snapshot lists
// every currently active user and every configured group (names, group
// membership, and limits only — never a key or its digest), and
// limiter.currentUsage reads every one's current-window counters — plus
// the synthetic total scope's own (v0.2 data-layer task) — in ONE
// storeGetMulti round trip total (v0.2 final review wave, 2026-08-20; see
// limiter.currentUsage's own doc comment for the tradeoff this
// supersedes) — users, groups, and the total scope are concatenated into
// a single scopes slice before that one call, then the flat result is
// sliced back into the three response sections at the same split points,
// order preserved. Unlike buildLimitScopes (routes_unified.go), this
// never omits an entity for having nil limits — the dashboard shows usage
// for every user and group, limited or not.
// chunkedCurrentUsage (above) reads every one's current-window counters —
// plus the synthetic total scope's own (v0.2 data-layer task) — in
// batches of at most adminUsageChunkScopes scopes per storeGetMulti round
// trip (security audit finding 4, 2026-08-22, capping the v0.2 final
// review wave's own "fold everything into ONE round trip" design — see
// adminUsageChunkScopes' own doc comment for why) — users, groups, and the
// total scope are concatenated into a single scopes slice before
// chunking, then the flat result is sliced back into the three response
// sections at the same split points, order preserved. Unlike
// buildLimitScopes (routes_unified.go), this never omits an entity for
// having nil limits — the dashboard shows usage for every user and group,
// limited or not.
// adminUsageChunkScopes bounds how many scopes buildAdminUsage sends to
// limiter.currentUsage in a single storeGetMulti round trip (security
// audit finding 4, 2026-08-22). currentUsage's own doc comment (limits.go)
// already explains why every scope was folded into ONE round trip rather
// than one per scope — but "one round trip" for the WHOLE catalog means
// that round trip's own pipeline size, and how long it holds the shared,
// mutex-guarded Redis connection (respClient, resp.go), now scales with
// how many users and groups are configured. With 1,000 users that is
// ~8,000 GETs (usageKeysPerScope=8) in one pipelined call, during which
// every concurrent checkAndCount from live LLM traffic queues behind the
// same connection mutex — and the admin dashboard polls this endpoint
// every 5s. chunkedCurrentUsage (below) keeps each individual round trip's
// key count bounded, letting live traffic interleave between chunks,
// while a chunk still batches every scope inside it exactly as before,
// keeping most of currentUsage's own round-trip-collapsing win. 200
// scopes (1,600 keys per chunk) comfortably covers the live cluster's 5
// friends + 4 home users in a single chunk, so today's real deployment
// sees no behavior change at all — chunking only engages once a
// deployment's user+group count actually grows past it. One accepted,
// disclosed side effect: currentUsage's own doc comment notes a transient
// store error marks every scope in ONE call storeDown together; chunking
// narrows that blast radius to one chunk's worth of scopes instead of the
// whole response, which is a strict improvement, not a new risk.
const adminUsageChunkScopes = 200
// chunkedCurrentUsage calls limiter.currentUsage in batches of at most
// adminUsageChunkScopes scopes at a time, concatenating the results in
// the SAME order scopes was given — buildAdminUsage's own slicing of the
// flat result back into users/groups/total (below) is unaffected by, and
// unaware of, the chunking underneath. A scopes slice at or under the
// chunk size makes exactly one call, byte-for-byte what calling
// currentUsage directly would have done.
func (g *Gateway) chunkedCurrentUsage(scopes []limitScope) []scopeUsage {
out := make([]scopeUsage, 0, len(scopes))
for len(scopes) > 0 {
n := adminUsageChunkScopes
if n > len(scopes) {
n = len(scopes)
}
out = append(out, g.limiter.currentUsage(scopes[:n])...)
scopes = scopes[n:]
}
return out
}
func (g *Gateway) buildAdminUsage() adminUsageResponse {
userSummaries, groupSummaries := g.auth.snapshot()
@@ -653,7 +700,7 @@ func (g *Gateway) buildAdminUsage() adminUsageResponse {
}
scopes = append(scopes, limitScope{kind: totalScopeKind, id: totalScopeID, limits: nil})
allUsage := g.limiter.currentUsage(scopes)
allUsage := g.chunkedCurrentUsage(scopes)
userUsage := allUsage[:len(userSummaries)]
groupUsage := allUsage[len(userSummaries) : len(userSummaries)+len(groupSummaries)]
totalUsage := allUsage[len(userSummaries)+len(groupSummaries)]
+186
View File
@@ -6,6 +6,7 @@ import (
"crypto/sha256"
"encoding/hex"
"encoding/json"
"fmt"
"net/http"
"net/http/httptest"
"net/url"
@@ -1456,6 +1457,100 @@ func TestLimiterCurrentUsage_BatchesOneGetMultiCallForAllScopes(t *testing.T) {
}
}
// --- security audit finding 4: admin usage batch chunking ---
// TestChunkedCurrentUsage_SplitsIntoMultipleRoundTrips is the regression
// for finding 4: buildAdminUsage previously handed limiter.currentUsage
// the WHOLE scopes slice in one call, so one round trip's own pipeline
// size (and how long it holds the shared, mutex-guarded store connection)
// scaled with catalog size with no ceiling at all. With more scopes than
// adminUsageChunkScopes, chunkedCurrentUsage must issue more than one
// getMulti call — proving chunking actually happens, not merely that the
// constant exists.
func TestChunkedCurrentUsage_SplitsIntoMultipleRoundTrips(t *testing.T) {
t.Parallel()
fixedNow := time.Date(2026, 8, 22, 10, 0, 0, 0, time.UTC)
const n = adminUsageChunkScopes + 50 // over one chunk, under two
store := &countingMultiStore{values: make(map[string]int64, n)}
scopes := make([]limitScope, n)
for i := 0; i < n; i++ {
id := fmt.Sprintf("user%d", i)
scopes[i] = limitScope{kind: "user", id: id, limits: &LimitsConfig{}}
store.values[windowKey("user", id, metricReq, windowMin, fixedNow)] = int64(i)
}
l := newLimiter(store, true)
l.nowFn = func() time.Time { return fixedNow }
gw := &Gateway{limiter: l}
got := gw.chunkedCurrentUsage(scopes)
if store.getMultiCalls != 2 {
t.Errorf("getMultiCalls = %d, want 2 (%d scopes over one %d-scope chunk)", store.getMultiCalls, n, adminUsageChunkScopes)
}
if len(got) != n {
t.Fatalf("len(got) = %d, want %d", len(got), n)
}
// Order must be preserved across the chunk boundary — spot-check the
// first scope of the SECOND chunk, whose value only exists if the
// concatenation kept scopes[adminUsageChunkScopes] aligned with
// got[adminUsageChunkScopes].
boundary := got[adminUsageChunkScopes]
if boundary.id != fmt.Sprintf("user%d", adminUsageChunkScopes) || boundary.requestsPerMinute != int64(adminUsageChunkScopes) {
t.Errorf("scope at the chunk boundary = %+v, want id=%q requestsPerMinute=%d (order preserved across chunks)",
boundary, fmt.Sprintf("user%d", adminUsageChunkScopes), adminUsageChunkScopes)
}
}
// TestChunkedCurrentUsage_AtOrBelowChunkSize_MakesExactlyOneCall proves
// chunking is a no-op difference for any deployment at or under
// adminUsageChunkScopes — including the live cluster's 5 friends + 4 home
// users (project brief) — matching currentUsage's own pre-existing
// single-round-trip behavior exactly.
func TestChunkedCurrentUsage_AtOrBelowChunkSize_MakesExactlyOneCall(t *testing.T) {
t.Parallel()
fixedNow := time.Date(2026, 8, 22, 10, 0, 0, 0, time.UTC)
store := &countingMultiStore{values: map[string]int64{}}
scopes := []limitScope{
{kind: "user", id: "alice", limits: &LimitsConfig{}},
{kind: "user", id: "bob", limits: &LimitsConfig{}},
{kind: "group", id: "eng", limits: &LimitsConfig{}},
}
l := newLimiter(store, true)
l.nowFn = func() time.Time { return fixedNow }
gw := &Gateway{limiter: l}
got := gw.chunkedCurrentUsage(scopes)
if store.getMultiCalls != 1 {
t.Errorf("getMultiCalls = %d, want 1 (scope count under adminUsageChunkScopes)", store.getMultiCalls)
}
if len(got) != len(scopes) {
t.Fatalf("len(got) = %d, want %d", len(got), len(scopes))
}
}
// TestChunkedCurrentUsage_EmptyScopes proves the edge case a real caller
// never hits directly (buildAdminUsage always appends the synthetic total
// scope) still behaves like currentUsage's own empty-input contract:
// no call, no panic, empty result.
func TestChunkedCurrentUsage_EmptyScopes(t *testing.T) {
t.Parallel()
store := &countingMultiStore{values: map[string]int64{}}
l := newLimiter(store, true)
gw := &Gateway{limiter: l}
got := gw.chunkedCurrentUsage(nil)
if store.getMultiCalls != 0 {
t.Errorf("getMultiCalls = %d, want 0 for an empty scopes slice", store.getMultiCalls)
}
if len(got) != 0 {
t.Errorf("len(got) = %d, want 0", len(got))
}
}
// --- overview: redis lastErr carries a timestamp (folded item 4, 2026-08-20 review) ---
func TestAdminOverview_RedisLastErrAt(t *testing.T) {
@@ -1899,3 +1994,94 @@ func TestAdminTargets_AccessListsMatchEnforcement_AndCountersReflectTraffic(t *t
t.Errorf("bot2.Access = %v, want %v (restricted's own agents list excludes it)", bot2.Access, want)
}
}
// --- default-preserving: the live cluster's real shape (security audit, 2026-08-22) ---
// TestDefaultPreserving_LiveClusterShape_LoadsAndServesIdentically models
// the live production shape this round's fixes must not break (project
// brief): 5 "friend" users plus 4 "home" users (steve-cw and kevinsandom
// are named live users; the rest are representative — real names this
// test has no visibility into, which is exactly why finding 5's fix stops
// short of restricting the name charset, see buildEntry's own doc
// comment, auth.go) and 11 MCP servers. A2A agents are out of scope here:
// none of this round's fixes touch mcp_a2a.go or any A2A-specific code
// path, so there is no mechanism by which they could regress agent
// handling.
//
// It proves construction succeeds unchanged (finding 5's new empty/
// duplicate-name checks accept every one of these real names), a
// friends-group user's configured requests-per-minute budget still buys
// EXACTLY that many tools/list calls against all 11 servers — not
// budget/11 — regressing security review round 2's critical finding 3
// (weighting tools/list by backend count was reverted specifically
// because it was not default-preserving at this exact shape: both live
// groups have an empty MCP allow-list, i.e. all 11 servers, so a naive
// per-backend weight would have silently cut every configured
// requests-per-minute budget by ~11x here), and the admin dashboard's
// usage poll still returns exactly one row per user and per group in a
// single store round trip (finding 4's chunking is a no-op at 10 users,
// well under adminUsageChunkScopes).
func TestDefaultPreserving_LiveClusterShape_LoadsAndServesIdentically(t *testing.T) {
friendNames := []string{"steve-cw", "kevinsandom", "friend3", "friend4", "friend5"}
homeNames := []string{"home1", "home2", "home3", "home4"}
const friendsRPM = 60
cfg := CreateConfig()
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}
cfg.Admin = &AdminConfig{Enabled: true}
cfg.Groups = map[string]*GroupConfig{
"friends": {Limits: &LimitsConfig{RequestsPerMinute: friendsRPM}},
"home": {Limits: &LimitsConfig{RequestsPerMinute: 120}},
}
cfg.MCPServers = make(map[string]*TargetConfig, 11)
for i := 0; i < 11; i++ {
srv := newMockJSONRPCServer(t, []mcpTool{{Name: "lookup"}})
cfg.MCPServers[fmt.Sprintf("mcp%d", i)] = &TargetConfig{URL: srv.srv.URL}
}
inline := []*UserConfig{{Name: "admin1", Group: "home", Admin: true, APIKey: "sk-admin1"}}
for _, name := range friendNames {
inline = append(inline, &UserConfig{Name: name, Group: "friends", APIKey: "sk-" + name})
}
for _, name := range homeNames {
inline = append(inline, &UserConfig{Name: name, Group: "home", APIKey: "sk-" + name})
}
cfg.Users = &UsersConfig{Inline: inline}
h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw")
if err != nil {
t.Fatalf("New: want the realistic live-cluster config to load unchanged, got error: %v", err)
}
// Drives friendsRPM calls, not just one (security review round 2:
// a single-call assertion cannot distinguish "budget buys N calls"
// from "budget buys N/11 calls" — exactly the blind spot that let the
// weighted-charge regression ship undetected).
for i := 1; i <= friendsRPM; i++ {
rec := httptest.NewRecorder()
h.ServeHTTP(rec, newFederatedRequest(t, "sk-steve-cw", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "tools/list", ID: json.RawMessage(fmt.Sprintf("%d", i))}))
if rec.Code != http.StatusOK {
t.Fatalf("tools/list call %d/%d status = %d, want 200 (budget must buy all %d calls against 11 servers), body=%s", i, friendsRPM, rec.Code, friendsRPM, rec.Body.String())
}
}
usageReq := httptest.NewRequest("GET", adminUsagePath, nil)
usageReq.Header.Set("Authorization", "Bearer sk-admin1")
usageRec := httptest.NewRecorder()
h.ServeHTTP(usageRec, usageReq)
if usageRec.Code != http.StatusOK {
t.Fatalf("admin usage status = %d, want 200, body=%s", usageRec.Code, usageRec.Body.String())
}
var usage adminUsageResponse
if err := json.Unmarshal(usageRec.Body.Bytes(), &usage); err != nil {
t.Fatalf("decode admin usage: %v", err)
}
wantUsers := len(friendNames) + len(homeNames) + 1 // + admin1
if len(usage.Users) != wantUsers {
t.Errorf("admin usage users = %d, want %d", len(usage.Users), wantUsers)
}
if len(usage.Groups) != 2 {
t.Errorf("admin usage groups = %d, want 2", len(usage.Groups))
}
}
+342 -22
View File
@@ -4,6 +4,7 @@ import (
"crypto/sha256"
"crypto/subtle"
"fmt"
"net"
"net/http"
"net/url"
"path"
@@ -187,6 +188,164 @@ type authEntry struct {
digest [32]byte
}
// authFailureWindow and authFailureLimit bound how many failed identify
// attempts one source IP may make within a window before this package
// treats it as "throttled" for OBSERVABILITY purposes only (security
// audit finding 3, 2026-08-22; revised by security review round 2,
// 2026-08-22, critical finding 1 — see identify's own doc comment for
// why this is observability, not enforcement, post-revision). 20
// failures per minute is generous headroom for a legitimate client with
// a stale or mistyped key retried a few times, while still flagging a
// source failing fast enough to be worth an operator's attention.
const (
authFailureWindow = time.Minute
authFailureLimit = 20
// authFailureLogEvery rate-limits BOTH logAuthEvent's routine
// failed-auth line and its distinct throttle-engaged line (logger.go)
// — the identical technique limiter.logStoreError (limits.go) already
// applies to store-error lines, sharing that function's own
// storeErrorLogEvery interval for cadence consistency.
authFailureLogEvery = storeErrorLogEvery
)
// authFailureMapCap bounds authFailureTracker's total tracked-source-IP
// count across both its generations (security review round 2,
// 2026-08-22, important finding 1): a distributed attacker sending
// failed attempts from many distinct IPs could otherwise grow the
// tracker's memory without any limit, and a naive full-map sweep to
// expire old entries would hold its lock for a duration that scales with
// tracked-IP count (measured, on the general-purpose memoryStore this
// field used before this revision: ~83ms mutex-held at 1M tracked IPs,
// blocking every identify call for that duration). authFailureTracker's
// own two-generation rotation (its doc comment below) bounds total memory
// to at most 2×authFailureMapCap entries with an O(1) rotation, never a
// scan.
const authFailureMapCap = 50_000
// authFailureEntry is one source IP's failure count within
// authFailureTracker, expiring authFailureWindow after its first failure
// in the current generation.
type authFailureEntry struct {
expiry time.Time
count int64
}
// authFailureTracker is a bounded, two-generation per-source-IP failed-
// auth-attempt counter (security review round 2, 2026-08-22, important
// finding 1) — deliberately NOT the general-purpose memoryStore
// (limits.go): memoryStore's map has no size cap at all and its
// sweepLocked does a full, mutex-held O(n) scan, which is fine for the
// bounded-by-configuration scope space checkAndCount/account use (one
// entry per configured user/group) but wrong for a value with no
// configured upper bound like source IP.
//
// current is the active generation every increment/read goes to first;
// previous is the immediately-prior generation, consulted as a read-only
// fallback so an IP's count does not spuriously read as zero the instant
// after a rotation. Rotation — current becomes previous, a fresh empty
// map becomes current — is a two-pointer swap (O(1), never a scan),
// triggered on inserting a brand-new IP once current is already at
// authFailureMapCap; total memory is bounded to at most
// 2×authFailureMapCap entries. An IP whose entry is displaced by a
// rotation simply starts a fresh count on its next failure — the same
// outcome a normal authFailureWindow rollover already produces for any
// tracked IP; this is accepted, lossy-under-extreme-load behavior, the
// same trade-off limiter.spawnTokens' own doc comment (limits.go)
// accepts for a different bounded resource.
type authFailureTracker struct {
current map[string]*authFailureEntry
previous map[string]*authFailureEntry
nowFn func() time.Time
mu sync.Mutex
}
// newAuthFailureTracker returns an empty tracker using nowFn for its
// clock — a func value, not a copied time.Time, so a caller that later
// swaps what nowFn points to (authStore's own fakeClock-in-tests
// convention) is honored on every subsequent call.
func newAuthFailureTracker(nowFn func() time.Time) *authFailureTracker {
return &authFailureTracker{current: make(map[string]*authFailureEntry), nowFn: nowFn}
}
// get returns ip's current failure count, or 0 if it is untracked or its
// window has expired. Checks current first, then previous.
func (t *authFailureTracker) get(ip string) int64 {
now := t.nowFn()
t.mu.Lock()
defer t.mu.Unlock()
if e, ok := t.current[ip]; ok && now.Before(e.expiry) {
return e.count
}
if e, ok := t.previous[ip]; ok && now.Before(e.expiry) {
return e.count
}
return 0
}
// increment records one failure for ip and returns its new count for the
// current window, creating a fresh entry in current if ip is untracked in
// current or its own entry there already expired — a plain fixed-window
// counter, matching memoryStore.incrBy's own contract (limits.go) in
// spirit. Rotates current -> previous before inserting a brand-new ip
// once current is already at authFailureMapCap — see this type's own doc
// comment for why that bounds total memory without ever scanning.
func (t *authFailureTracker) increment(ip string) int64 {
now := t.nowFn()
t.mu.Lock()
defer t.mu.Unlock()
if e, ok := t.current[ip]; ok {
if now.Before(e.expiry) {
e.count++
return e.count
}
e.count = 1
e.expiry = now.Add(authFailureWindow)
return 1
}
if len(t.current) >= authFailureMapCap {
t.previous = t.current
t.current = make(map[string]*authFailureEntry, authFailureMapCap)
}
t.current[ip] = &authFailureEntry{count: 1, expiry: now.Add(authFailureWindow)}
return 1
}
// logGate rate-limits a single log line to once per authFailureLogEvery,
// counting how many calls were suppressed in between so the next actually
// -written line can report them (security review round 2, 2026-08-22,
// important finding 6): a single global rate-limit gate would otherwise
// make several DISTINCT attacker IPs failing within one interval look
// identical to a single noisy one — three separate IPs each failing once
// in a 30s window previously produced only ONE log line total, with no
// sign two more existed. Reporting the suppressed count in the next line
// restores that visibility without tracking a per-IP gate (an LRU of
// recently-logged IPs was the documented alternative; a suppressed count
// is simpler and gives an operator the same "something else happened"
// signal). Two independent zero-value-usable instances back authStore's
// routine-failure and throttle-engaged lines (below) so a burst of one
// can never suppress, or be suppressed by, the other. The very first call
// on a fresh gate always logs (lastLogAt's zero value is more than
// authFailureLogEvery before any real now).
type logGate struct {
lastLogAt time.Time
suppressed int64
mu sync.Mutex
}
func (g *logGate) shouldLog(now time.Time) (log bool, suppressed int64) {
g.mu.Lock()
defer g.mu.Unlock()
if !g.lastLogAt.IsZero() && now.Sub(g.lastLogAt) < authFailureLogEvery {
g.suppressed++
return false, 0
}
suppressed = g.suppressed
g.suppressed = 0
g.lastLogAt = now
return true, suppressed
}
// authStore identifies requests by API key and resolves them to a user and
// their group. Raw API keys are never retained after construction — only
// their SHA-256 digests.
@@ -198,16 +357,30 @@ type authEntry struct {
// overwrite each other. buildMu is never held at the same time as mu, so a
// rebuild in progress never blocks identify's readers. reloadMu (Task 4)
// serializes maybeReload's throttle check and the reload it may trigger.
//
// failures is identify's own per-source-IP failed-attempt counter
// (security audit finding 3, 2026-08-22; failureLog/throttleLog added by
// security review round 2, 2026-08-22) — see identify's own doc comment
// for why this is pure observability, never enforcement, and
// authFailureTracker's own doc comment for why it is not the
// general-purpose memoryStore (limits.go). It is authStore's own
// dedicated instance: auth failures are a distinct concept from LLM
// traffic limits, and authStore must work identically whether or not a
// Gateway even has a limiter wired yet. failureLog/throttleLog back
// logAuthEvent's two distinct rate-limited lines (logger.go).
type authStore struct {
lastModTime time.Time
lastCheck time.Time
log gatewayLogger
failures *authFailureTracker
usersFile *usersFile
nowFn func() time.Time
groups map[string]*group
inline map[[32]byte]*authEntry // built once by newAuthStore; never mutated afterward
byDigest map[[32]byte]*authEntry // inline entries plus the current file-sourced set; guarded by mu
usersFile *usersFile // nil when Users.File is not configured; maybeReload no-ops
log gatewayLogger // set alongside usersFile; unused when usersFile is nil
nowFn func() time.Time // injected for tests; defaults to time.Now
lastCheck time.Time // guarded by reloadMu
lastModTime time.Time // guarded by reloadMu
fileUserCount int // guarded by mu; size of the current file-sourced user set
byDigest map[[32]byte]*authEntry
inline map[[32]byte]*authEntry
failureLog logGate
throttleLog logGate
fileUserCount int
mu sync.RWMutex
buildMu sync.Mutex
reloadMu sync.Mutex
@@ -223,6 +396,12 @@ func newAuthStore(cfg *Config) (*authStore, error) {
byDigest: make(map[[32]byte]*authEntry),
nowFn: time.Now,
}
// Delegate through a.nowFn rather than copying it: a test that
// overrides a.nowFn AFTER construction (the fakeClock convention
// newReloadableAuthStore already uses, users_file_test.go) must still
// control failures' own clock — this closure reads a.nowFn fresh on
// every call, so it always sees whichever func is currently assigned.
a.failures = newAuthFailureTracker(func() time.Time { return a.nowFn() })
for name, gc := range cfg.Groups {
if gc == nil {
return nil, fmt.Errorf("llmgateway: group %q: config must not be nil", name)
@@ -265,11 +444,22 @@ func newAuthStore(cfg *Config) (*authStore, error) {
}
if cfg.Users != nil {
// seenNames rejects two INLINE users sharing the same name
// (security audit finding 5, 2026-08-22) — scoped to inline only,
// not cross-checked against a later file-sourced set: replaceFileUsers
// intentionally allows a file user to override an inline user of the
// same name (its own doc comment), so "duplicate name" only ever
// means "duplicate within the same source".
seenNames := make(map[string]bool, len(cfg.Users.Inline))
for _, uc := range cfg.Users.Inline {
entry, err := a.buildEntry(uc)
if err != nil {
return nil, err
}
if seenNames[entry.user.name] {
return nil, fmt.Errorf("llmgateway: duplicate user name %q", entry.user.name)
}
seenNames[entry.user.name] = true
if _, exists := a.inline[entry.digest]; exists {
return nil, fmt.Errorf("llmgateway: duplicate API key for user %q", uc.Name)
}
@@ -283,10 +473,33 @@ func newAuthStore(cfg *Config) (*authStore, error) {
// buildEntry resolves uc's group and API key into an authEntry. The API key
// goes through resolveSecret before digesting, so env:/file: references
// work the same as for provider keys.
//
// uc.Name must be non-empty (security audit finding 5, 2026-08-22):
// counters are keyed on a user's name (windowKey's "kind:id:..." shape,
// limits.go, called with id=u.name at routes_unified.go's own call site),
// so an empty name would collapse every such user's requests-per-minute/
// day, token, and cost counters onto the single shared key
// "llmgw:user::...", each silently overwriting the others' admission
// checks and usage — never a legitimate configuration, always a config
// mistake worth failing construction over. This check is deliberately
// NOT extended to configNamePattern's full character restriction
// (providers.go): that pattern's own doc comment ties it specifically to
// being embedded as a URL PATH SEGMENT (passthroughRoute's
// "/{provider}/*", targetRoute's "/mcp/{name}/*" and "/a2a/{name}/*") — a
// concern that does not apply to a user name, which is never used as a
// route path segment anywhere in this package, only as a counter-key
// component and a JSON field in admin views. Restricting the character
// set beyond "non-empty" would risk rejecting a real, currently-working
// operator-chosen name this round has no visibility into, for a route-
// safety reason that has no bearing on users at all — see this finding's
// own report for the full reasoning.
func (a *authStore) buildEntry(uc *UserConfig) (*authEntry, error) {
if uc == nil {
return nil, fmt.Errorf("llmgateway: user config entry must not be nil")
}
if uc.Name == "" {
return nil, fmt.Errorf("llmgateway: user config entry must have a non-empty name")
}
grp, ok := a.groups[uc.Group]
if !ok {
return nil, fmt.Errorf("llmgateway: user %q references unknown group %q", uc.Name, uc.Group)
@@ -345,11 +558,21 @@ func (a *authStore) replaceFileUsers(us []*UserConfig) error {
}
next[digest] = entry
}
// seenFileNames rejects two FILE users sharing the same name (security
// audit finding 5, 2026-08-22) — scoped to this call's own us slice
// only, never checked against a.inline: a file user overriding an
// inline user of the same name is this function's own documented,
// intentional behavior (doc comment above), not a duplicate.
seenFileNames := make(map[string]bool, len(us))
for _, uc := range us {
entry, err := a.buildEntry(uc)
if err != nil {
return err
}
if seenFileNames[entry.user.name] {
return fmt.Errorf("llmgateway: duplicate user name %q", entry.user.name)
}
seenFileNames[entry.user.name] = true
if _, exists := next[entry.digest]; exists {
return fmt.Errorf("llmgateway: duplicate API key for user %q", uc.Name)
}
@@ -363,28 +586,125 @@ func (a *authStore) replaceFileUsers(us []*UserConfig) error {
return nil
}
// clientIP returns the source address identify's own throttle keys on:
// r.RemoteAddr with any ":port" suffix stripped — net/http's own "address
// as seen by this server's TCP accept", never a client-supplied header
// (security audit finding 3, 2026-08-22). This plugin runs AS a Traefik
// middleware inside the shared ingress process, not behind some second,
// untrusted proxy of its own, so RemoteAddr is the same trust boundary
// logAuthEvent's own RemoteAddr-only logging (logger.go) already assumes,
// and the same one routes_passthrough.go's dangerousClientHeaderPrefixes
// strip (X-Forwarded-*, X-Auth-*) exists to keep a client from spoofing on
// the OUTBOUND side — deliberately NOT read here either, for the same
// reason. A RemoteAddr with no ":port" (net.SplitHostPort's own error
// case — not expected from net/http in practice, but not assumed) is used
// as-is rather than discarded, so a throttle key still exists for
// whatever value RemoteAddr carries.
func clientIP(r *http.Request) string {
host, _, err := net.SplitHostPort(r.RemoteAddr)
if err != nil {
return r.RemoteAddr
}
return host
}
// authThrottled reports whether ip has reached authFailureLimit failed
// identify attempts within the current authFailureWindow — a read-only
// observability signal for logAuthEvent's distinct throttle-engaged line
// (logger.go). identify itself never consults this (security review
// round 2, 2026-08-22, critical finding 1 — see identify's own doc
// comment for why).
func (a *authStore) authThrottled(ip string) bool {
return a.failures.get(ip) >= authFailureLimit
}
// recordAuthFailure increments ip's failure count for the current
// authFailureWindow (authFailureTracker.increment's own fixed-window
// contract). Pure observability post security review round 2's critical
// finding 1: it never influences identify's own return value (see
// identify's doc comment) — only authThrottled's read, for
// logAuthEvent's distinct throttle-engaged line.
func (a *authStore) recordAuthFailure(ip string) {
a.failures.increment(ip)
}
// shouldLogAuthFailure reports whether the next ROUTINE failed-auth log
// line (logAuthEvent, logger.go) should actually be written, and how many
// prior failure events were suppressed since the last line actually
// wrote (logGate's own doc comment covers the suppressed-count
// reasoning) — the same rate-limiting technique limiter.logStoreError
// (limits.go) already applies to store-error lines, applied here so an
// attacker driving unlimited failed-auth attempts can never drive
// unlimited synchronous stderr writes on the shared Traefik ingress
// (security audit finding 3, 2026-08-22). Successful-auth logging is
// deliberately left unbounded: it requires an already-valid API key, so
// its volume is bounded by legitimate authenticated traffic — itself
// already governed by checkAndCount's requests-per-minute/day limits
// elsewhere — not by an unauthenticated caller's attempt rate, which is
// exactly what this finding is about bounding.
func (a *authStore) shouldLogAuthFailure() (log bool, suppressed int64) {
return a.failureLog.shouldLog(a.nowFn())
}
// shouldLogThrottleEngaged is shouldLogAuthFailure's counterpart for the
// distinct "this source IP is currently throttled" line (security review
// round 2, 2026-08-22, important finding 5) — its own separate rate-limit
// gate (logGate's own doc comment), so a burst of routine failure lines
// can never suppress, or be suppressed by, this higher-signal line.
func (a *authStore) shouldLogThrottleEngaged() (log bool, suppressed int64) {
return a.throttleLog.shouldLog(a.nowFn())
}
// identify resolves r's presented API key to a user and their group. It
// reads "Authorization: Bearer <key>" (case-insensitive "bearer" prefix) or
// "x-api-key: <key>"; Bearer wins when both are present. The presented key
// is looked up by SHA-256 digest, then verified with a constant-time
// compare — the key itself is never retained.
//
// The presented key is verified FIRST, UNCONDITIONALLY (security review
// round 2, 2026-08-22, critical finding 1): a request carrying a VALID
// key always succeeds, before any per-source-IP throttle state is even
// read. The prior revision of this function checked authThrottled BEFORE
// verifying the key at all — but this plugin runs as a Traefik middleware
// behind a Kubernetes Service, and with externalTrafficPolicy: Cluster
// (this operator's own live deployment), r.RemoteAddr is the NODE's own
// SNAT address, shared by every external caller reaching that node, not
// a per-tenant address. Confirmed against this operator's own 10 live
// users: 20 failed requests from anywhere on the internet, sharing that
// node, locked out every one of them, including the admin dashboard's own
// poll. Verifying the key first makes that impossible — a legitimate,
// already-valid request can never be rejected by another caller's failed
// attempts, from the same address or not.
//
// Every failure path below — no key presented, key not found, or a
// (practically unreachable, see subtle.ConstantTimeCompare's own call
// site) digest mismatch — records one failure for the caller's source IP
// via recordAuthFailure; logAuthEvent (logger.go) separately decides, per
// call, whether to emit a routine failure line or a distinct throttle-
// engaged line by consulting authThrottled itself. No branch in this
// function ever changes its own return value based on throttle state —
// throttling is pure observability now (bounding the log's write volume
// and the tracker's own memory; authFailureTracker's doc comment,
// above), never enforcement: an attacker with no valid key gains nothing
// from this change (every failure path here already returned false
// either way, throttled or not), and the raw per-attempt cost this
// function pays (one SHA-256 digest, one map lookup, one
// ConstantTimeCompare) is already cheap enough — nanoseconds — that
// gating it early bought negligible protection against the cost itself,
// only the false-positive risk the paragraph above describes.
func (a *authStore) identify(r *http.Request) (*user, *group, bool) {
key, ok := presentedKey(r)
if !ok {
return nil, nil, false
if ok {
digest := sha256.Sum256([]byte(key))
a.mu.RLock()
entry, found := a.byDigest[digest]
a.mu.RUnlock()
if found && subtle.ConstantTimeCompare(digest[:], entry.digest[:]) == 1 {
return entry.user, entry.group, true
}
}
digest := sha256.Sum256([]byte(key))
a.mu.RLock()
entry, ok := a.byDigest[digest]
a.mu.RUnlock()
if !ok {
return nil, nil, false
}
if subtle.ConstantTimeCompare(digest[:], entry.digest[:]) != 1 {
return nil, nil, false
}
return entry.user, entry.group, true
a.recordAuthFailure(clientIP(r))
return nil, nil, false
}
// presentedKey extracts the API key from r's Authorization Bearer header or
+312
View File
@@ -3,6 +3,7 @@ package traefikllmgateway
import (
"fmt"
"net/http/httptest"
"strings"
"sync"
"testing"
"time"
@@ -577,3 +578,314 @@ func TestAuthStore_Identify_NotBlockedByBuildMu(t *testing.T) {
t.Fatal("identify blocked while buildMu is held — reader lock must not depend on the writer-serialization lock")
}
}
// --- security audit finding 3: per-source-IP auth-failure throttling ---
func TestClientIP_StripsPort(t *testing.T) {
r := httptest.NewRequest("POST", "/v1/chat/completions", nil)
r.RemoteAddr = "203.0.113.9:54321"
if got := clientIP(r); got != "203.0.113.9" {
t.Errorf("clientIP = %q, want %q", got, "203.0.113.9")
}
}
func TestClientIP_MalformedRemoteAddr_UsedAsIs(t *testing.T) {
r := httptest.NewRequest("POST", "/v1/chat/completions", nil)
r.RemoteAddr = "not-a-host-port"
if got := clientIP(r); got != "not-a-host-port" {
t.Errorf("clientIP = %q, want the raw RemoteAddr back unchanged", got)
}
}
// TestAuthStore_Identify_ValidKeyAlwaysSucceeds_EvenWhenIPThrottled is the
// core regression for security review round 2's critical finding 1: this
// plugin runs as a Traefik middleware behind a Kubernetes Service — with
// externalTrafficPolicy: Cluster, r.RemoteAddr is the NODE's own SNAT
// address, shared by every external caller reaching that node, not a
// per-tenant one. The prior revision of identify checked the per-IP
// throttle BEFORE verifying the key, so 20 failed requests from anywhere
// sharing that node locked out every legitimate tenant too. After well
// past authFailureLimit failed attempts from one source IP, a SUBSEQUENT
// VALID key from that EXACT same IP must still succeed immediately.
func TestAuthStore_Identify_ValidKeyAlwaysSucceeds_EvenWhenIPThrottled(t *testing.T) {
a, err := newAuthStore(testAuthCfg()) // inline user "a" / sk-secret, group eng
if err != nil {
t.Fatal(err)
}
const sharedIP = "203.0.113.9:54321"
for i := 0; i < authFailureLimit*2; i++ { // well past the threshold
r := httptest.NewRequest("POST", "/v1/chat/completions", nil)
r.RemoteAddr = sharedIP
r.Header.Set("x-api-key", "wrong-key")
a.identify(r)
}
// Sanity: the tracker really did cross the threshold — otherwise this
// test would not actually exercise the throttled-IP path at all.
if !a.authThrottled("203.0.113.9") {
t.Fatal("test setup: expected the tracker to be over authFailureLimit for the shared IP")
}
valid := httptest.NewRequest("POST", "/v1/chat/completions", nil)
valid.RemoteAddr = sharedIP
valid.Header.Set("Authorization", "Bearer sk-secret")
u, grp, ok := a.identify(valid)
if !ok || u.name != "a" || grp.name != "eng" {
t.Fatalf("want a valid key to succeed even from a throttled source IP, got user=%v group=%v ok=%v", u, grp, ok)
}
}
// TestAuthStore_AuthThrottled_TracksPerIPIndependently proves
// authFailureTracker's own counting (used only by logAuthEvent's distinct
// throttle-engaged line now, logger.go — never by identify's return
// value, see identify's own doc comment) is still scoped per source IP:
// one IP crossing the threshold must never affect a different IP's own
// count.
func TestAuthStore_AuthThrottled_TracksPerIPIndependently(t *testing.T) {
a, err := newAuthStore(testAuthCfg())
if err != nil {
t.Fatal(err)
}
for i := 0; i < authFailureLimit; i++ {
r := httptest.NewRequest("POST", "/v1/chat/completions", nil)
r.RemoteAddr = "203.0.113.9:1"
r.Header.Set("x-api-key", "wrong-key")
a.identify(r)
}
if !a.authThrottled("203.0.113.9") {
t.Error("want 203.0.113.9 throttled after authFailureLimit failures")
}
if a.authThrottled("198.51.100.5") {
t.Error("want a different, never-failing IP to remain unthrottled")
}
}
// TestAuthStore_AuthThrottled_ResetsAfterWindowExpires proves
// authFailureTracker is a fixed-window counter, not a permanent record:
// once authFailureWindow has elapsed, a source IP's earlier failures no
// longer count toward authThrottled.
func TestAuthStore_AuthThrottled_ResetsAfterWindowExpires(t *testing.T) {
a, err := newAuthStore(testAuthCfg())
if err != nil {
t.Fatal(err)
}
clock := &fakeClock{now: time.Unix(1_700_000_000, 0)}
a.nowFn = clock.Now
for i := 0; i < authFailureLimit; i++ {
r := httptest.NewRequest("POST", "/v1/chat/completions", nil)
r.RemoteAddr = "203.0.113.9:1"
r.Header.Set("x-api-key", "wrong-key")
a.identify(r)
}
if !a.authThrottled("203.0.113.9") {
t.Fatal("want throttled before the window elapses")
}
clock.Advance(authFailureWindow + time.Second)
if a.authThrottled("203.0.113.9") {
t.Error("want the throttle cleared once authFailureWindow has elapsed")
}
}
// --- authFailureTracker: direct unit tests ---
func TestAuthFailureTracker_IncrementAndGet(t *testing.T) {
now := time.Unix(1_700_000_000, 0)
tr := newAuthFailureTracker(func() time.Time { return now })
if got := tr.get("1.2.3.4"); got != 0 {
t.Errorf("get on untracked ip = %d, want 0", got)
}
for i := int64(1); i <= 5; i++ {
if got := tr.increment("1.2.3.4"); got != i {
t.Errorf("increment #%d = %d, want %d", i, got, i)
}
}
if got := tr.get("1.2.3.4"); got != 5 {
t.Errorf("get after 5 increments = %d, want 5", got)
}
}
func TestAuthFailureTracker_WindowExpiry(t *testing.T) {
clock := &fakeClock{now: time.Unix(1_700_000_000, 0)}
tr := newAuthFailureTracker(clock.Now)
tr.increment("1.2.3.4")
tr.increment("1.2.3.4")
if got := tr.get("1.2.3.4"); got != 2 {
t.Fatalf("get = %d, want 2", got)
}
clock.Advance(authFailureWindow + time.Second)
if got := tr.get("1.2.3.4"); got != 0 {
t.Errorf("get after window expiry = %d, want 0", got)
}
if got := tr.increment("1.2.3.4"); got != 1 {
t.Errorf("increment after window expiry = %d, want 1 (fresh window)", got)
}
}
// TestAuthFailureTracker_BoundsMemoryViaRotation is the regression for
// security review round 2's important finding 1: the tracker must never
// grow past 2×authFailureMapCap total entries — filling it past
// authFailureMapCap distinct IPs must trigger an O(1) generation
// rotation, not unbounded growth.
func TestAuthFailureTracker_BoundsMemoryViaRotation(t *testing.T) {
now := time.Unix(1_700_000_000, 0)
tr := newAuthFailureTracker(func() time.Time { return now })
for i := 0; i < authFailureMapCap+10; i++ {
tr.increment(fmt.Sprintf("10.%d.%d.%d", i/65536, (i/256)%256, i%256))
}
tr.mu.Lock()
total := len(tr.current) + len(tr.previous)
currentLen := len(tr.current)
tr.mu.Unlock()
if total > 2*authFailureMapCap {
t.Errorf("tracked entries = %d, want <= %d (2x cap)", total, 2*authFailureMapCap)
}
if currentLen > authFailureMapCap {
t.Errorf("current generation = %d entries, want <= %d", currentLen, authFailureMapCap)
}
}
// TestAuthStore_Identify_FewFailedAttempts_NeverThrottled is the good-path
// counterpart: a handful of failed attempts (well under authFailureLimit
// — a human retrying a mistyped key a couple of times) must never trip
// the throttle.
func TestAuthStore_Identify_FewFailedAttempts_NeverThrottled(t *testing.T) {
a, err := newAuthStore(testAuthCfg())
if err != nil {
t.Fatal(err)
}
const ip = "203.0.113.9:54321"
for i := 0; i < 3; i++ {
r := httptest.NewRequest("POST", "/v1/chat/completions", nil)
r.RemoteAddr = ip
r.Header.Set("x-api-key", "wrong-key")
a.identify(r)
}
r := httptest.NewRequest("POST", "/v1/chat/completions", nil)
r.RemoteAddr = ip
r.Header.Set("Authorization", "Bearer sk-secret")
if _, _, ok := a.identify(r); !ok {
t.Fatal("want a valid key to succeed after only 3 prior failures (well under authFailureLimit)")
}
}
// TestAuthStore_Identify_SuccessfulAuthDoesNotRecordFailure proves a
// successful identify never itself counts toward the failure budget —
// otherwise a legitimate, high-traffic user could eventually throttle
// themselves.
func TestAuthStore_Identify_SuccessfulAuthDoesNotRecordFailure(t *testing.T) {
a, err := newAuthStore(testAuthCfg())
if err != nil {
t.Fatal(err)
}
const ip = "203.0.113.9:54321"
r := httptest.NewRequest("POST", "/v1/chat/completions", nil)
r.RemoteAddr = ip
r.Header.Set("Authorization", "Bearer sk-secret")
for i := 0; i < authFailureLimit+5; i++ {
if _, _, ok := a.identify(r); !ok {
t.Fatalf("attempt %d: want success (valid key), got failure", i)
}
}
}
// --- security audit finding 5: duplicate/empty user names ---
func TestNewAuthStore_EmptyUserName_ReturnsError(t *testing.T) {
cfg := &Config{
Providers: map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}},
Groups: map[string]*GroupConfig{"eng": {}},
Users: &UsersConfig{Inline: []*UserConfig{{Name: "", Group: "eng", APIKey: "sk-secret"}}},
}
if _, err := newAuthStore(cfg); err == nil {
t.Fatal("want error for an empty user name")
}
}
func TestNewAuthStore_DuplicateUserName_ReturnsError(t *testing.T) {
cfg := &Config{
Providers: map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}},
Groups: map[string]*GroupConfig{"eng": {}},
Users: &UsersConfig{Inline: []*UserConfig{
{Name: "dup", Group: "eng", APIKey: "sk-one"},
{Name: "dup", Group: "eng", APIKey: "sk-two"},
}},
}
_, err := newAuthStore(cfg)
if err == nil {
t.Fatal("want error for two inline users sharing the same name")
}
if !strings.Contains(err.Error(), "duplicate user name") {
t.Errorf("error = %q, want it to mention the duplicate-name check specifically (not the API-key one)", err.Error())
}
}
func TestAuthStore_ReplaceFileUsers_EmptyUserName_ReturnsError(t *testing.T) {
a, err := newAuthStore(testAuthCfg())
if err != nil {
t.Fatal(err)
}
if err := a.replaceFileUsers([]*UserConfig{{Name: "", Group: "eng", APIKey: "sk-file1"}}); err == nil {
t.Fatal("want error for a file user with an empty name")
}
}
func TestAuthStore_ReplaceFileUsers_DuplicateUserNameWithinFile_ReturnsError(t *testing.T) {
a, err := newAuthStore(testAuthCfg())
if err != nil {
t.Fatal(err)
}
err = a.replaceFileUsers([]*UserConfig{
{Name: "dup", Group: "eng", APIKey: "sk-file1"},
{Name: "dup", Group: "eng", APIKey: "sk-file2"},
})
if err == nil {
t.Fatal("want error for two file users sharing the same name")
}
if !strings.Contains(err.Error(), "duplicate user name") {
t.Errorf("error = %q, want it to mention the duplicate-name check specifically (not the API-key one)", err.Error())
}
// The store must still be usable — the rejected replace must not have
// partially mutated it.
r := httptest.NewRequest("POST", "/v1/chat/completions", nil)
r.Header.Set("Authorization", "Bearer sk-secret")
if _, _, ok := a.identify(r); !ok {
t.Fatal("want the original inline user still resolvable after a rejected duplicate-name replace")
}
}
// TestAuthStore_ReplaceFileUsers_SameNameAsInline_DifferentKey_OverrideAllowed
// proves the new duplicate-name check does NOT reject the intentional
// cross-source override replaceFileUsers' own doc comment describes: a
// file user of the same name as an inline user, even with a different API
// key, must still win — this is a feature (rotating/promoting a
// statically-configured user through the hot-reloadable file), not a
// duplicate.
func TestAuthStore_ReplaceFileUsers_SameNameAsInline_DifferentKey_OverrideAllowed(t *testing.T) {
a, err := newAuthStore(testAuthCfg()) // inline user "a" / sk-secret, group eng
if err != nil {
t.Fatal(err)
}
if err := a.replaceFileUsers([]*UserConfig{{Name: "a", Group: "eng", APIKey: "sk-file-a"}}); err != nil {
t.Fatalf("replaceFileUsers: want the same-name override allowed, got error: %v", err)
}
inlineReq := httptest.NewRequest("POST", "/v1/chat/completions", nil)
inlineReq.Header.Set("Authorization", "Bearer sk-secret")
if _, _, ok := a.identify(inlineReq); ok {
t.Fatal("want the inline key overridden once the file defines a same-named user")
}
fileReq := httptest.NewRequest("POST", "/v1/chat/completions", nil)
fileReq.Header.Set("Authorization", "Bearer sk-file-a")
u, _, ok := a.identify(fileReq)
if !ok || u.name != "a" {
t.Fatalf("want the file-sourced user resolvable under the overridden name, got %v ok=%v", u, ok)
}
}
+46 -7
View File
@@ -19,20 +19,59 @@ func (g *Gateway) errorf(format string, args ...any) {
}
// logAuthEvent logs one authentication attempt, at the point one of
// authStore.identify's six call sites resolves it. A failure logs the
// method, path, and remote address the request came from; a success logs
// the resolved user's name alongside the same method and path. Neither
// line ever includes the presented API key itself — identify resolves a
// key to a *user and hands this function only that user's name, never the
// key it authenticated.
// authStore.identify's six call sites resolves it. A success logs the
// resolved user's name alongside the request's method and path. Neither
// line this function ever writes includes the presented API key itself —
// identify resolves a key to a *user and hands this function only that
// user's name, never the key it authenticated.
//
// A failure logs one of two distinct lines (security review round 2,
// 2026-08-22, important finding 5): if the caller's source IP
// (clientIP) is currently authThrottled, a distinct "throttle engaged"
// line naming the IP and the limit it exceeded — an operator seeing this
// knows immediately that a SOURCE, not just an individual request, needs
// attention, which the routine per-attempt line alone does not convey.
// Otherwise the routine failure line (method, path, remote address).
// Both lines are independently rate-limited (authStore.shouldLog
// ThrottleEngaged / shouldLogAuthFailure, auth.go — the identical
// technique limiter.logStoreError, limits.go, already applies to
// store-error lines) so an unauthenticated caller driving unlimited
// failed attempts can never drive unlimited synchronous stderr writes on
// the shared Traefik ingress — a slow log pipe would otherwise stall
// request handling for every tenant, not just the one failing auth. Each
// gate reports how many events it suppressed since its last written line
// (security review round 2, important finding 6): a single global gate
// would otherwise make several DISTINCT attacker IPs failing within one
// interval indistinguishable from one noisy IP — suppressedSuffix
// appends that count to the line when non-zero. The success line stays
// completely unbounded — see shouldLogAuthFailure's own doc comment for
// why that is safe.
func (g *Gateway) logAuthEvent(ok bool, userName string, r *http.Request) {
if !ok {
g.logf("auth failed: %s %q from %s", r.Method, r.URL.Path, r.RemoteAddr)
ip := clientIP(r)
if g.auth.authThrottled(ip) {
if log, suppressed := g.auth.shouldLogThrottleEngaged(); log {
g.logf("auth throttle: source %s exceeded %d failed attempts in the last %s%s", ip, authFailureLimit, authFailureWindow, suppressedSuffix(suppressed))
}
return
}
if log, suppressed := g.auth.shouldLogAuthFailure(); log {
g.logf("auth failed: %s %q from %s%s", r.Method, r.URL.Path, r.RemoteAddr, suppressedSuffix(suppressed))
}
return
}
g.logf("auth ok: user %q %s %q", userName, r.Method, r.URL.Path)
}
// suppressedSuffix formats n (a logGate's own suppressed-event count,
// auth.go) as a log-line suffix, or "" when there is nothing to report.
func suppressedSuffix(n int64) string {
if n <= 0 {
return ""
}
return fmt.Sprintf(" (%d more suppressed since last log)", n)
}
// authEventUserName extracts u's name for logAuthEvent's userName
// argument, or "" when u is nil — identify's failure path, which never
// resolves a user.
+182
View File
@@ -0,0 +1,182 @@
package traefikllmgateway
import (
"bytes"
"context"
"io"
"net/http"
"net/http/httptest"
"os"
"strings"
"testing"
"time"
)
// newTestGatewayForLogger builds a minimal *Gateway via the real New()
// constructor (the only supported construction path — see this package's
// other test files, none of which build a bare &Gateway{} literal) for
// logAuthEvent's own tests below.
func newTestGatewayForLogger(t *testing.T) *Gateway {
t.Helper()
cfg := CreateConfig()
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}
h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw")
if err != nil {
t.Fatalf("New: %v", err)
}
gw, ok := h.(*Gateway)
if !ok {
t.Fatal("handler is not *Gateway")
}
return gw
}
// captureStderr redirects os.Stderr for the duration of fn (the same
// os.Pipe-based idiom this package's other test files already use for
// logf/errorf, e.g. TestGateway_Logf_WritesInfoPrefix, llmgateway_test.go)
// and returns everything written to it.
func captureStderr(t *testing.T, fn func()) string {
t.Helper()
origStderr := os.Stderr
r, w, err := os.Pipe()
if err != nil {
t.Fatalf("os.Pipe: %v", err)
}
os.Stderr = w
fn()
_ = w.Close()
os.Stderr = origStderr
var buf bytes.Buffer
if _, err := io.Copy(&buf, r); err != nil {
t.Fatalf("io.Copy: %v", err)
}
return buf.String()
}
// TestLogAuthEvent_Failure_RateLimited is the regression for the missing
// rate limit on logAuthEvent's failure line (security audit finding 3,
// 2026-08-22): an attacker driving many failed-auth attempts in quick
// succession must produce ONE log line per authFailureLogEvery, not one
// per attempt — otherwise the synchronous os.Stderr write itself becomes
// an amplification vector on the shared Traefik ingress.
func TestLogAuthEvent_Failure_RateLimited(t *testing.T) {
gw := newTestGatewayForLogger(t)
clock := &fakeClock{now: time.Unix(1_700_000_000, 0)}
gw.auth.nowFn = clock.Now
r := httptest.NewRequest("POST", "/v1/chat/completions", nil)
r.RemoteAddr = "203.0.113.9:1234"
out := captureStderr(t, func() {
for i := 0; i < 5; i++ {
gw.logAuthEvent(false, "", r)
}
})
if n := strings.Count(out, "auth failed:"); n != 1 {
t.Fatalf("auth-failed lines = %d, want exactly 1 (5 attempts within one authFailureLogEvery window), output=%q", n, out)
}
}
// TestLogAuthEvent_Failure_LogsAgainAfterIntervalElapses proves the rate
// limit is a cadence, not a one-shot silence: once authFailureLogEvery has
// actually elapsed, the next failure is logged again.
func TestLogAuthEvent_Failure_LogsAgainAfterIntervalElapses(t *testing.T) {
gw := newTestGatewayForLogger(t)
clock := &fakeClock{now: time.Unix(1_700_000_000, 0)}
gw.auth.nowFn = clock.Now
r := httptest.NewRequest("POST", "/v1/chat/completions", nil)
r.RemoteAddr = "203.0.113.9:1234"
out := captureStderr(t, func() {
gw.logAuthEvent(false, "", r)
clock.Advance(authFailureLogEvery + time.Second)
gw.logAuthEvent(false, "", r)
})
if n := strings.Count(out, "auth failed:"); n != 2 {
t.Fatalf("auth-failed lines = %d, want exactly 2 (interval elapsed between the two attempts), output=%q", n, out)
}
}
// TestLogAuthEvent_Success_NeverRateLimited proves successful-auth logging
// is deliberately left unbounded (shouldLogAuthFailure's own doc comment,
// auth.go, justifies why): it takes an already-valid key to reach, so its
// volume is bounded by legitimate traffic elsewhere, not by an
// unauthenticated caller's attempt rate.
func TestLogAuthEvent_Success_NeverRateLimited(t *testing.T) {
gw := newTestGatewayForLogger(t)
clock := &fakeClock{now: time.Unix(1_700_000_000, 0)}
gw.auth.nowFn = clock.Now
r := httptest.NewRequest("POST", "/v1/chat/completions", nil)
out := captureStderr(t, func() {
for i := 0; i < 5; i++ {
gw.logAuthEvent(true, "alice", r)
}
})
if n := strings.Count(out, "auth ok:"); n != 5 {
t.Fatalf("auth-ok lines = %d, want 5 (success logging must never be rate-limited), output=%q", n, out)
}
}
// TestLogAuthEvent_ThrottleEngaged_DistinctLine proves logAuthEvent emits
// a distinct "auth throttle:" line, not the routine "auth failed:" one,
// once a source IP has crossed authFailureLimit (security review round
// 2, 2026-08-22, important finding 5) — an operator seeing this line
// knows a SOURCE, not just one request, needs attention.
func TestLogAuthEvent_ThrottleEngaged_DistinctLine(t *testing.T) {
gw := newTestGatewayForLogger(t)
clock := &fakeClock{now: time.Unix(1_700_000_000, 0)}
gw.auth.nowFn = clock.Now
r := httptest.NewRequest("POST", "/v1/chat/completions", nil)
r.RemoteAddr = "203.0.113.9:1234"
for i := 0; i < authFailureLimit; i++ {
gw.auth.recordAuthFailure(clientIP(r))
}
out := captureStderr(t, func() {
gw.logAuthEvent(false, "", r)
})
if !strings.Contains(out, "auth throttle:") {
t.Errorf("want an 'auth throttle:' line once the IP has crossed the limit, got %q", out)
}
if strings.Contains(out, "auth failed:") {
t.Errorf("want the routine 'auth failed:' line suppressed in favor of the distinct throttle line, got %q", out)
}
}
// TestLogAuthEvent_Failure_SuppressedCountReported proves the routine
// failure line's rate-limit gate reports how many events it suppressed
// (security review round 2, 2026-08-22, important finding 6): a single
// global gate would otherwise make several rapid failures within one
// window indistinguishable from a single one — the next line actually
// written must name how many were collapsed into it.
func TestLogAuthEvent_Failure_SuppressedCountReported(t *testing.T) {
gw := newTestGatewayForLogger(t)
clock := &fakeClock{now: time.Unix(1_700_000_000, 0)}
gw.auth.nowFn = clock.Now
r := httptest.NewRequest("POST", "/v1/chat/completions", nil)
r.RemoteAddr = "203.0.113.9:1234"
captureStderr(t, func() {
gw.logAuthEvent(false, "", r) // logged (first call on a fresh gate)
gw.logAuthEvent(false, "", r) // suppressed #1
gw.logAuthEvent(false, "", r) // suppressed #2
})
clock.Advance(authFailureLogEvery + time.Second)
out := captureStderr(t, func() {
gw.logAuthEvent(false, "", r)
})
if !strings.Contains(out, "2 more suppressed since last log") {
t.Errorf("want the next line to report 2 suppressed events, got %q", out)
}
}
+241 -40
View File
@@ -4,6 +4,7 @@ import (
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
@@ -66,6 +67,65 @@ const (
mcpSessionCloseTimeout = 5 * time.Second
)
// mcpBackendResponseMaxBytes caps a single backend MCP server's HTTP
// response body for the tools/list FAN-OUT path only (mcpFederatedToolsList)
// — previously maxRequestBytes (10MiB, routes_unified.go), a budget sized
// for a CLIENT's own chat/embeddings request body, not a backend's
// JSON-RPC reply (security audit finding 1c, 2026-08-22). 4MiB matches
// this codebase's other "generous but bounded" read caps for a non-
// primary-request-body read (cache.go's maxAccountingTeeBytes is also
// 4MiB) — measured at 547KiB for a real tools/list response describing
// 20 fat tools (security review round 2, 2026-08-22), comfortably below
// this cap — while bounding mcpFederatedFanoutConcurrency backends'
// worst-case concurrent memory to a small multiple of 4MiB, against the
// unbounded ~110MiB (11 production servers × up to 10MiB each, all fired
// at once) the security audit measured before this fix. This cap is
// deliberately NOT used for tools/call — see mcpBackendCallResponseMaxBytes,
// whose own doc comment covers why the fan-out memory argument does not
// apply there.
const mcpBackendResponseMaxBytes = 4 << 20
// mcpBackendCallResponseMaxBytes caps a single backend's response for the
// tools/call path (mcpFederatedToolsCall) — kept at maxRequestBytes
// (10MiB), the SAME budget this file used everywhere before finding 1c,
// deliberately NOT shrunk to mcpBackendResponseMaxBytes (security review
// round 2, 2026-08-22, important finding 4): tools/call always resolves
// to exactly ONE backend (resolveFederatedTool), so the fan-out memory-
// amplification argument that justifies the fan-out path's smaller cap
// does not hold here at all — a single tool result can legitimately be
// much larger than a tools/list schema listing (an extracted document, an
// image or other blob a tool returns), and silently truncating one into
// a generic parse failure would be a real, avoidable regression for any
// tool whose legitimate results sit between 4MiB and 10MiB.
const mcpBackendCallResponseMaxBytes = maxRequestBytes
// errMCPResponseTooLarge is the sentinel doBackendJSONRPC wraps when a
// backend response hits its call's maxBytes cap (security review round 2,
// 2026-08-22, important finding 4), letting the tools/call handler report
// "response too large" instead of a generic "upstream error".
//
// errors.Is is the correct match here, and it is Yaegi-safe. An earlier
// revision used a string marker plus strings.Contains out of caution; a
// harness on yaegi v0.16.1 (the version tools/yaegi-check pins, same
// stdlib.Symbols set) interpreting this shape found errors.Is correct
// through shallow, deep, and multi-%w wrapping, and correctly false for
// an unrelated error. The reason is that declaring the VARIABLE in
// interpreted code does not make its VALUE interpreted: errors.New
// returns a compiled *errors.errorString and fmt.Errorf a compiled
// *fmt.wrapError, so the whole chain errors.Is walks is compiled. This
// is the same shape retry.go:isTransient has run in production since
// v0.2.0, and errRequestBuildFailed (providers.go) is already matched
// this way through a multi-%w wrap.
//
// Neither documented Yaegi trap applies: matchesSentinel's TRIP-WIRE
// (limits.go) is about a plugin-declared error TYPE, which errors.New
// does not create, and the comma-ok trap needs a plugin-declared
// interface. A string check was also strictly worse than it looked —
// err.Error() is itself an interpreted-to-compiled method dispatch, so
// it verified nothing extra, and matching on text that partly comes
// from the backend admits false positives a sentinel cannot have.
var errMCPResponseTooLarge = errors.New("mcp backend response too large")
// jsonrpcRequest is one JSON-RPC 2.0 request/notification, decoded from
// the client's POST body and re-encoded (with a synthetic id) for this
// gateway's own outbound calls to a backend MCP server (mcpBackendCall).
@@ -234,9 +294,32 @@ func allowedMCPServerNames(cfg *Config, grp *group) []string {
//
// checkAndCount runs once per incoming request, before method dispatch —
// covering the caller's own user/group/total scopes exactly like every
// other metered route — with per-target attribution added separately
// other metered route, charging the plain weight of one request
// regardless of method — with per-target attribution added separately
// inside mcpFederatedToolsList/mcpFederatedToolsCall (their own doc
// comments cover exactly what "attributed" means for each).
//
// tools/list's fan-out is deliberately NOT charged against the caller's
// own requests-per-minute/requests-per-day budget (security review round
// 2, 2026-08-22, critical finding 3 — reverting security audit finding
// 1b's weighted-charge attempt): weighting checkAndCount by the number of
// allowed MCP servers is not default-preserving — on this operator's own
// live fleet, both configured groups have an empty MCP allow-list (every
// server allowed), so EVERY tools/list call would have weighed 11,
// silently turning a configured requests-per-minute: 60 into an effective
// budget of 5 successful tools/list calls per minute, and
// examples/kubernetes.yaml's shipped requests-per-minute: 10 would have
// 429'd the very FIRST tools/list call ever made — with no config change
// on the operator's part. The fan-out's real amplification is still
// fully visible to an operator without silently consuming tenant quota:
// mcpFederatedToolsList's own countTargetRequests call attributes one
// request to EVERY server actually attempted, at the "mcp"/target-id
// scope (limits.go) — a SEPARATE scope kind from the user/group/total
// scopes checkAndCount evaluates here — so the dashboard's targets view
// already shows the true per-server cost. The fan-out's outbound HTTP/
// memory cost is bounded instead by mcpFederatedFanoutConcurrency's
// semaphore (mcpFederatedToolsList's own doc comment), which is the
// actual protection against the amplification finding 1 described.
func (g *Gateway) handleMCPFederated(w http.ResponseWriter, r *http.Request, u *user, grp *group) {
if r.Header.Get("Upgrade") != "" {
writeOAIError(w, http.StatusNotImplemented, "invalid_request_error", "websocket/upgrade passthrough not supported")
@@ -270,7 +353,7 @@ func (g *Gateway) handleMCPFederated(w http.ResponseWriter, r *http.Request, u *
case req.Method == "ping":
writeJSONRPCResult(w, req.ID, map[string]any{})
case req.Method == "tools/list":
g.mcpFederatedToolsList(w, r, req, grp)
g.mcpFederatedToolsList(w, r, req, allowedMCPServerNames(g.cfg, grp))
case req.Method == "tools/call":
g.mcpFederatedToolsCall(w, r, req, grp)
case strings.HasPrefix(req.Method, "notifications/"):
@@ -370,8 +453,16 @@ type mcpToolsListResult struct {
// doBackendJSONRPC issues one JSON-RPC 2.0 POST to targetURL — an MCP
// server's own configured base URL, the same single endpoint
// handleTargetProxy reverse-proxies every method to — carrying sessionID
// as the Mcp-Session-Id request header when non-empty. It is the one place
// that actually builds and sends an HTTP request for this file; every
// as the Mcp-Session-Id request header when non-empty, and reading at
// most maxBytes of the response body (mcpBackendResponseMaxBytes for the
// tools/list fan-out, mcpBackendCallResponseMaxBytes for tools/call —
// their own doc comments cover why the two differ; security review round
// 2, 2026-08-22, important finding 4). A response body that hits maxBytes
// is reported by wrapping errMCPResponseTooLarge, not silently truncated
// into whatever partial (and likely invalid) JSON happened to fit.
//
// It is the one place that actually builds and sends an HTTP request for
// this file; every
// other function in the handshake-fallback machinery (mcpBackendCall,
// mcpBackendHandshake, mcpBackendCloseSession's own DELETE aside) goes
// through this.
@@ -396,7 +487,7 @@ type mcpToolsListResult struct {
// that WAS valid JSON but not a real JSON-RPC response (see
// parseBackendJSONRPC's own validity check) would have silently produced
// an all-zero-fields response otherwise.
func (g *Gateway) doBackendJSONRPC(ctx context.Context, targetURL, method string, params any, sessionID string) (resp *jsonrpcResponse, status int, respSessionID string, err error) {
func (g *Gateway) doBackendJSONRPC(ctx context.Context, targetURL, method string, params any, sessionID string, maxBytes int64) (resp *jsonrpcResponse, status int, respSessionID string, err error) {
paramsRaw, err := json.Marshal(params)
if err != nil {
return nil, 0, "", err
@@ -433,13 +524,26 @@ func (g *Gateway) doBackendJSONRPC(ctx context.Context, targetURL, method string
status = httpResp.StatusCode
respSessionID = httpResp.Header.Get(mcpSessionHeader)
respBody, err := io.ReadAll(io.LimitReader(httpResp.Body, maxRequestBytes))
// Reads one byte past maxBytes so a response that hits the cap is
// DETECTABLE (len(respBody) > maxBytes) rather than silently truncated
// and handed to parseBackendJSONRPC to fail confusingly (security
// review round 2, 2026-08-22, important finding 4).
respBody, err := io.ReadAll(io.LimitReader(httpResp.Body, maxBytes+1))
if err != nil {
return nil, status, respSessionID, err
}
// Status is checked BEFORE the size cap: a backend that fails with a
// large HTML error page is diagnosed as "upstream returned HTTP 500",
// which is actionable, rather than as "too large", which sends the
// operator hunting a payload-size problem that does not exist. An
// oversized SUCCESS response still reports the cap, which is the case
// the cap exists for.
if status < 200 || status >= 300 {
return nil, status, respSessionID, fmt.Errorf("upstream returned HTTP %d", status)
}
if int64(len(respBody)) > maxBytes {
return nil, status, respSessionID, fmt.Errorf("%w: exceeds %d bytes", errMCPResponseTooLarge, maxBytes)
}
parsed, err := parseBackendJSONRPC(httpResp.Header.Get("Content-Type"), respBody)
if err != nil {
@@ -457,13 +561,13 @@ func (g *Gateway) doBackendJSONRPC(ctx context.Context, targetURL, method string
// the initialize call itself failed outright (network error, non-2xx, a
// JSON-RPC error response to initialize itself) — there is no retry of a
// failed handshake.
func (g *Gateway) mcpBackendHandshake(ctx context.Context, targetURL string) (sessionID string, err error) {
func (g *Gateway) mcpBackendHandshake(ctx context.Context, targetURL string, maxBytes int64) (sessionID string, err error) {
initParams := map[string]any{
"protocolVersion": defaultMCPProtocolVersion,
"capabilities": map[string]any{},
"clientInfo": map[string]any{"name": "traefik-llmgateway", "version": pluginVersion},
}
resp, _, respSessionID, err := g.doBackendJSONRPC(ctx, targetURL, "initialize", initParams, "")
resp, _, respSessionID, err := g.doBackendJSONRPC(ctx, targetURL, "initialize", initParams, "", maxBytes)
if err != nil {
return "", err
}
@@ -524,9 +628,14 @@ func (g *Gateway) mcpBackendCloseSession(ctx context.Context, targetURL, session
// except the close itself — see mcpSessionCloseTimeout's own doc
// comment); callers derive it from context.WithTimeout(r.Context(), ...)
// with toolsListBackendTimeout or toolsCallBackendTimeout (MF1, review
// round 2, 2026-08-21).
func (g *Gateway) mcpBackendCall(ctx context.Context, targetURL, method string, params any) (*jsonrpcResponse, error) {
resp, status, _, err := g.doBackendJSONRPC(ctx, targetURL, method, params, "")
// round 2, 2026-08-21). maxBytes is doBackendJSONRPC's own response-size
// cap, passed straight through to every doBackendJSONRPC call this
// function makes (bare attempt, handshake, retry) — callers pass
// mcpBackendResponseMaxBytes or mcpBackendCallResponseMaxBytes depending
// on which path they are (security review round 2, 2026-08-22, important
// finding 4).
func (g *Gateway) mcpBackendCall(ctx context.Context, targetURL, method string, params any, maxBytes int64) (*jsonrpcResponse, error) {
resp, status, _, err := g.doBackendJSONRPC(ctx, targetURL, method, params, "", maxBytes)
if err == nil && resp.Error == nil {
return resp, nil
}
@@ -536,7 +645,7 @@ func (g *Gateway) mcpBackendCall(ctx context.Context, targetURL, method string,
return nil, err
}
sessionID, initErr := g.mcpBackendHandshake(ctx, targetURL)
sessionID, initErr := g.mcpBackendHandshake(ctx, targetURL, maxBytes)
if initErr != nil {
// The handshake itself failed: surface the bare attempt's own
// outcome when it has one (a JSON-RPC error response is more
@@ -557,7 +666,7 @@ func (g *Gateway) mcpBackendCall(ctx context.Context, targetURL, method string,
}()
}
retryResp, _, _, retryErr := g.doBackendJSONRPC(ctx, targetURL, method, params, sessionID)
retryResp, _, _, retryErr := g.doBackendJSONRPC(ctx, targetURL, method, params, sessionID, maxBytes)
if retryErr != nil {
return nil, retryErr
}
@@ -632,47 +741,126 @@ func lastSSEDataLine(body []byte) []byte {
return last
}
// mcpFederatedFanoutConcurrency bounds how many of names' backend servers
// mcpFederatedToolsList's own fan-out below contacts at once — a
// buffered-channel semaphore gating each goroutine's actual outbound HTTP
// work rather than its launch. It resembles limiter.spawnTokens
// (limits.go) only in using a buffered channel: spawnTokens is a
// non-blocking select with a `default:` that DROPS the work, whereas this
// one blocks until a slot frees and must never skip a server. Do not
// treat the two as interchangeable. Before this bound (security audit finding
// 1a, 2026-08-22), one incoming tools/list request spawned one goroutine
// per allowed server with NO cap at all: with the 11 production servers,
// one client call could hold up to 11 concurrent connections and up to
// 11× mcpBackendResponseMaxBytes of backend response bodies in memory at
// once, entirely outside this gateway's own admission control. Every
// server is eventually contacted regardless of this bound — a goroutine
// blocks on the semaphore, it is never dropped or skipped. 8 (raised from
// an initial 4, security review round 2, 2026-08-22, important finding
// 3): 4 measurably tripled worst-case wall-clock latency for an
// unrestricted group against all 11 servers (907ms measured at 300ms/
// server, versus ~300ms pre-bound) and, against hung backends, pushed the
// worst case to ceil(11/4)×toolsListBackendTimeout = 60s — well past the
// ~30s timeout many real MCP clients use, holding the ServeHTTP goroutine
// the whole time. 8 nearly halves both: ceil(11/8)=2 batches. The fanoutCtx
// deadline below (not this constant alone) is what actually re-bounds the
// worst case back down near toolsListBackendTimeout regardless of
// concurrency or backend count — see its own doc comment.
const mcpFederatedFanoutConcurrency = 8
// mcpFederatedToolsList implements the federated "tools/list": it fans
// out one mcpBackendCall per server allowedMCPServerNames(grp) permits,
// concurrently — each under its own toolsListBackendTimeout budget
// (context.WithTimeout off r.Context(), MF1) — then merges every
// reachable server's own tools into one result, each tool's name prefixed
// "<serverName>_<toolName>", the exact underscore-separator convention the
// old agentgateway this plugin replaces used, and the one pugbot's
// mcpclient and agentkit have already persisted into their own tool-id
// databases (e.g. "brave-search_brave_web_search" — verified live). A
// server that errors, times out, or returns an unparsable/invalid result
// is logged and skipped, not surfaced as a whole-call failure — UNLESS
// every single attempted server failed, in which case an empty tools list
// would misleadingly look like "this caller's group has no MCP access" —
// see the loud-failure branch below.
// out one mcpBackendCall per server in names (handleMCPFederated's own
// allowedMCPServerNames(grp) call, made once there so the exact set
// contacted is always the one the client actually asked to reach), at
// most mcpFederatedFanoutConcurrency at a time — every backend sharing
// ONE overall toolsListBackendTimeout deadline (fanoutCtx, below) — then
// merges every reachable server's own tools into one result,
// each tool's name prefixed "<serverName>_<toolName>", the exact
// underscore-separator convention the old agentgateway this plugin
// replaces used, and the one pugbot's mcpclient and agentkit have already
// persisted into their own tool-id databases (e.g.
// "brave-search_brave_web_search" — verified live). A server that errors,
// times out, panics, or returns an unparsable/invalid result is logged
// and skipped, not surfaced as a whole-call failure — UNLESS every single
// attempted server failed, in which case an empty tools list would
// misleadingly look like "this caller's group has no MCP access" — see
// the loud-failure branch below.
//
// countTargetRequests (limits.go) attributes one request to EVERY server
// ATTEMPTED here, in ONE batched call after wg.Wait() — "attempted", not
// "reached" or "succeeded": a server this gateway dialed and got a
// response (or a timeout, or a connection refusal) from still had a real
// request sent to it and a real slot of this gateway's outbound capacity
// spent on it, which is what these counters exist to track. A single
// federated tools/list call can move several targets' own counters, not
// just one — unlike handleTargetProxy's always-exactly-one-target shape
// (mcp_a2a.go).
func (g *Gateway) mcpFederatedToolsList(w http.ResponseWriter, r *http.Request, req jsonrpcRequest, grp *group) {
names := allowedMCPServerNames(g.cfg, grp)
// response (or a timeout, a connection refusal, or a panic) from still had
// a real request sent to it and a real slot of this gateway's outbound
// capacity spent on it, which is what these counters exist to track. A
// single federated tools/list call can move several targets' own
// counters, not just one — unlike handleTargetProxy's always-exactly-
// one-target shape (mcp_a2a.go).
func (g *Gateway) mcpFederatedToolsList(w http.ResponseWriter, r *http.Request, req jsonrpcRequest, names []string) {
var (
mu sync.Mutex
wg sync.WaitGroup
merged = make([]mcpTool, 0, len(names))
failed []string
)
// fanoutCtx bounds the WHOLE fan-out to toolsListBackendTimeout from
// when THIS call started — not each individual backend's own start
// (security review round 2, 2026-08-22, important finding 3). With
// mcpFederatedFanoutConcurrency backends running at once, a later
// batch's goroutines only begin once an earlier one's semaphore slot
// frees; giving each one its own FRESH toolsListBackendTimeout budget
// (this function's pre-bound design, and its own initial post-
// semaphore revision) let total wall-clock grow with
// ceil(len(names)/mcpFederatedFanoutConcurrency) batches. Every
// backend below shares this ONE already-ticking context instead —
// context.WithTimeout always resolves to the earlier of a parent's
// existing deadline and its own duration (proven by
// TestHandleMCPFederated_ToolsList_SlowBackend_BoundedByRequestContext),
// so deriving fanoutCtx from r.Context() here preserves that same
// property for the request's own deadline, on top of this one.
//
// The accepted trade: a later batch inherits the REMAINING budget, not
// a fresh one, so if every backend is slow enough to consume most of
// toolsListBackendTimeout on its own, servers in later batches are cut
// off and their tools are missing from the merged list. Measured: 11
// servers at 11s each returns 8 of 11 tools after 20s, where a per-
// backend budget returned 11 of 11 after 11s. This is deliberate —
// the alternative is a worst case of ceil(11/8) x 20s = 60s for a
// single tools/list — and it degrades to a partial list plus a logged
// failure line per starved server, never a hung request. Real
// tools/list calls answer in well under a second (at 300ms/server all
// 11 return), so this path needs a pathological backend to trigger.
fanoutCtx, fanoutCancel := context.WithTimeout(r.Context(), toolsListBackendTimeout)
defer fanoutCancel()
sem := make(chan struct{}, mcpFederatedFanoutConcurrency)
for _, name := range names {
wg.Add(1)
go func(name string) {
defer wg.Done()
ctx, cancel := context.WithTimeout(r.Context(), toolsListBackendTimeout)
defer cancel()
// Unrecovered-panic guard (security audit finding 2,
// 2026-08-22): this goroutine runs off the request's own
// goroutine, so a panic here has no ServeHTTP caller to
// unwind into and would crash the whole shared Traefik
// process — the same reasoning limiter.spawn (limits.go) and
// modelRegistry.refreshProvider/captureModelMetadata
// (registry.go) already apply to their own off-request
// goroutines. A panicking backend must degrade to "that
// server failed", exactly like an HTTP error from it, never
// a process-wide crash.
defer func() {
if rec := recover(); rec != nil {
g.logf("federated tools/list: server %q panicked: %v", name, rec)
mu.Lock()
failed = append(failed, name)
mu.Unlock()
}
}()
sem <- struct{}{}
defer func() { <-sem }()
targetURL := g.cfg.MCPServers[name].URL
resp, err := g.mcpBackendCall(ctx, targetURL, "tools/list", struct{}{})
resp, err := g.mcpBackendCall(fanoutCtx, targetURL, "tools/list", struct{}{}, mcpBackendResponseMaxBytes)
mu.Lock()
defer mu.Unlock()
@@ -784,11 +972,24 @@ func (g *Gateway) mcpFederatedToolsCall(w http.ResponseWriter, r *http.Request,
ctx, cancel := context.WithTimeout(r.Context(), toolsCallBackendTimeout)
defer cancel()
targetURL := g.cfg.MCPServers[serverName].URL
resp, err := g.mcpBackendCall(ctx, targetURL, "tools/call", mcpToolCallParams{Name: toolName, Arguments: params.Arguments})
// mcpBackendCallResponseMaxBytes (10MiB), not the fan-out's smaller
// mcpBackendResponseMaxBytes (security review round 2, 2026-08-22,
// important finding 4): tools/call always resolves to exactly ONE
// backend, so the fan-out memory-amplification argument does not
// apply here — see mcpBackendCallResponseMaxBytes' own doc comment.
resp, err := g.mcpBackendCall(ctx, targetURL, "tools/call", mcpToolCallParams{Name: toolName, Arguments: params.Arguments}, mcpBackendCallResponseMaxBytes)
g.limiter.countTargetRequest(targetKindMCP, serverName)
if err != nil {
g.logf("federated tools/call: server %q: %v", serverName, err)
writeJSONRPCErrorResponse(w, req.ID, jsonrpcInternalError, "upstream error")
// A response that hit mcpBackendCallResponseMaxBytes gets its own
// specific message, not the generic "upstream error" a truncated-
// then-unparsable body would otherwise produce (security review
// round 2, 2026-08-22, important finding 4).
if errors.Is(err, errMCPResponseTooLarge) {
writeJSONRPCErrorResponse(w, req.ID, jsonrpcInternalError, "response too large")
} else {
writeJSONRPCErrorResponse(w, req.ID, jsonrpcInternalError, "upstream error")
}
return
}
+395 -2
View File
@@ -4,10 +4,13 @@ import (
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"net/http/httptest"
"slices"
"strings"
"sync"
"sync/atomic"
"testing"
@@ -1029,7 +1032,7 @@ func TestMcpBackendCall_HandshakeFallbackTrigger(t *testing.T) {
t.Fatal("handler is not *Gateway")
}
_, _ = gw.mcpBackendCall(context.Background(), srv.URL, "tools/list", struct{}{})
_, _ = gw.mcpBackendCall(context.Background(), srv.URL, "tools/list", struct{}{}, mcpBackendResponseMaxBytes)
gotCalls := callCount.Load()
if tc.wantHandshake && gotCalls < 2 {
@@ -1058,12 +1061,134 @@ func TestMcpBackendCall_NetworkFailure_NeverAttemptsHandshake(t *testing.T) {
t.Fatal("handler is not *Gateway")
}
_, err = gw.mcpBackendCall(context.Background(), "http://127.0.0.1:1", "tools/list", struct{}{})
_, err = gw.mcpBackendCall(context.Background(), "http://127.0.0.1:1", "tools/list", struct{}{}, mcpBackendResponseMaxBytes)
if err == nil {
t.Fatal("want an error (nothing listening on 127.0.0.1:1)")
}
}
// --- security review round 2, important finding 4: dual response caps ---
// TestMcpBackendCall_ResponseExceedsCap_ReturnsDistinctTooLargeError is a
// fast, small-scale unit test of doBackendJSONRPC's own maxBytes
// mechanism, using an artificially tiny cap so the test needs no
// multi-megabyte payload: a backend response over the cap must be
// reported by wrapping errMCPResponseTooLarge — distinguishable via
// errors.Is from an ordinary network/parse failure — never silently
// truncated into a confusing parse error.
func TestMcpBackendCall_ResponseExceedsCap_ReturnsDistinctTooLargeError(t *testing.T) {
const tinyCap = 64
big := strings.Repeat("x", tinyCap*4) // well past tinyCap once JSON-encoded
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
var req jsonrpcRequest
_ = json.NewDecoder(r.Body).Decode(&req)
result, _ := json.Marshal(map[string]string{"data": big})
w.Header().Set("Content-Type", "application/json")
_ = json.NewEncoder(w).Encode(jsonrpcResponse{ID: req.ID, Result: result})
}))
defer srv.Close()
cfg := newFederationTestConfig(srv.URL, "http://beta.invalid", false)
h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw")
if err != nil {
t.Fatalf("New: %v", err)
}
gw, ok := h.(*Gateway)
if !ok {
t.Fatal("handler is not *Gateway")
}
_, err = gw.mcpBackendCall(context.Background(), srv.URL, "tools/list", struct{}{}, tinyCap)
if err == nil {
t.Fatal("want an error for a response exceeding the cap")
}
if !errors.Is(err, errMCPResponseTooLarge) {
t.Errorf("err = %v, want it to wrap %v", err, errMCPResponseTooLarge)
}
}
// TestMcpBackendCall_ResponseAtOrUnderCap_Succeeds proves the cap+1 read
// technique does not false-positive: a response exactly AT the cap must
// still succeed, not be mistaken for oversized.
func TestMcpBackendCall_ResponseAtOrUnderCap_Succeeds(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
var req jsonrpcRequest
_ = json.NewDecoder(r.Body).Decode(&req)
result, _ := json.Marshal(map[string]bool{"ok": true})
w.Header().Set("Content-Type", "application/json")
_ = json.NewEncoder(w).Encode(jsonrpcResponse{ID: req.ID, Result: result})
}))
defer srv.Close()
cfg := newFederationTestConfig(srv.URL, "http://beta.invalid", false)
h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw")
if err != nil {
t.Fatalf("New: %v", err)
}
gw, ok := h.(*Gateway)
if !ok {
t.Fatal("handler is not *Gateway")
}
// A generous cap comfortably above this tiny response — must succeed
// cleanly, not be flagged as too-large.
resp, err := gw.mcpBackendCall(context.Background(), srv.URL, "tools/list", struct{}{}, 4096)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if resp.Error != nil {
t.Fatalf("resp.Error = %+v, want nil", resp.Error)
}
}
// TestHandleMCPFederated_ToolsCall_ResponseBetween4And10MiB_Succeeds is
// the end-to-end regression for important finding 4: a real tool result
// between mcpBackendResponseMaxBytes (4MiB, the fan-out-only cap) and
// mcpBackendCallResponseMaxBytes (10MiB, tools/call's own cap) — legal
// under the pre-finding-1c budget, and a realistic size for an image or
// extracted-document tool result — must succeed on the tools/call path,
// proving that path was never shrunk to the fan-out's smaller cap.
func TestHandleMCPFederated_ToolsCall_ResponseBetween4And10MiB_Succeeds(t *testing.T) {
const payloadBytes = 6 << 20 // 6MiB: over the 4MiB fan-out cap, under the 10MiB tools/call cap
big := strings.Repeat("y", payloadBytes)
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
var req jsonrpcRequest
_ = json.NewDecoder(r.Body).Decode(&req)
if req.Method == "tools/call" {
result, _ := json.Marshal(map[string]string{"data": big})
w.Header().Set("Content-Type", "application/json")
_ = json.NewEncoder(w).Encode(jsonrpcResponse{ID: req.ID, Result: result})
return
}
w.Header().Set("Content-Type", "application/json")
_ = json.NewEncoder(w).Encode(jsonrpcResponse{ID: req.ID, Result: json.RawMessage(`{}`)})
}))
defer srv.Close()
cfg := newFederationTestConfig(srv.URL, "http://beta.invalid", false)
h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw")
if err != nil {
t.Fatalf("New: %v", err)
}
rec := httptest.NewRecorder()
h.ServeHTTP(rec, newFederatedRequest(t, "sk-alice", jsonrpcRequest{
JSONRPC: jsonrpcVersion, Method: "tools/call", ID: json.RawMessage("1"),
Params: json.RawMessage(`{"name":"alpha_lookup"}`),
}))
if rec.Code != http.StatusOK {
t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String())
}
var got jsonrpcResponse
if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil {
t.Fatalf("decode: %v", err)
}
if got.Error != nil {
t.Fatalf("error = %+v, want nil (a 6MiB result must succeed on the tools/call path)", got.Error)
}
}
// TestHandleMCPFederated_ToolsList_HandshakeFallback_SucceedsAfterSessionRetry
// is MF3's core proof: a server that rejects a bare tools/list call
// (readitall/fetch's real, live-probed behavior — sessionRequiredMockServer's
@@ -1214,6 +1339,274 @@ func TestHandleMCPFederated_ToolsList_AllServersFail_ReturnsInternalError(t *tes
}
}
// --- security audit finding 1a: fan-out concurrency bound ---
// TestHandleMCPFederated_ToolsList_FanoutBoundedConcurrency proves
// mcpFederatedFanoutConcurrency actually caps how many backend servers a
// single tools/list call contacts AT ONCE — not just that it eventually
// contacts all of them. numServers is comfortably above the cap so the
// first wave of goroutines through the semaphore must stall on it: the
// test blocks every mock server's handler until every one of them
// observes exactly mcpFederatedFanoutConcurrency requests in flight
// simultaneously, then releases them all. If the fan-out were unbounded
// (the pre-fix behavior), every server's handler would receive its
// request immediately and this test would deadlock waiting for
// concurrency to reach a ceiling nothing ever enforces — a real
// regression here hangs the test until it times out, it does not
// silently pass.
func TestHandleMCPFederated_ToolsList_FanoutBoundedConcurrency(t *testing.T) {
const numServers = mcpFederatedFanoutConcurrency + 4
var (
current int32
peak int32
release = make(chan struct{})
)
cfg := CreateConfig()
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}
cfg.MCPServers = make(map[string]*TargetConfig, numServers)
for i := 0; i < numServers; i++ {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
n := atomic.AddInt32(&current, 1)
for {
old := atomic.LoadInt32(&peak)
if n <= old || atomic.CompareAndSwapInt32(&peak, old, n) {
break
}
}
<-release
atomic.AddInt32(&current, -1)
var req jsonrpcRequest
_ = json.NewDecoder(r.Body).Decode(&req)
result, _ := json.Marshal(mcpToolsListResult{})
w.Header().Set("Content-Type", "application/json")
_ = json.NewEncoder(w).Encode(jsonrpcResponse{ID: req.ID, Result: result})
}))
t.Cleanup(srv.Close)
cfg.MCPServers[fmt.Sprintf("srv%d", i)] = &TargetConfig{URL: srv.URL}
}
cfg.Groups = map[string]*GroupConfig{"default": {}}
cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}}
h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw")
if err != nil {
t.Fatalf("New: %v", err)
}
done := make(chan struct{})
rec := httptest.NewRecorder()
go func() {
h.ServeHTTP(rec, newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "tools/list", ID: json.RawMessage("1")}))
close(done)
}()
deadline := time.After(5 * time.Second)
for atomic.LoadInt32(&current) < mcpFederatedFanoutConcurrency {
select {
case <-deadline:
t.Fatal("concurrency never reached mcpFederatedFanoutConcurrency — fan-out may be serialized instead of bounded")
case <-time.After(time.Millisecond):
}
}
// Give any wrongly-unbounded extra goroutines a chance to also reach
// the handler before we sample — if the bound were missing, ALL
// numServers requests would already be in flight by now.
time.Sleep(20 * time.Millisecond)
if got := atomic.LoadInt32(&current); got > mcpFederatedFanoutConcurrency {
t.Fatalf("current in-flight backends = %d, want <= %d (fan-out is not bounded)", got, mcpFederatedFanoutConcurrency)
}
close(release)
select {
case <-done:
case <-time.After(5 * time.Second):
t.Fatal("request never completed after releasing all backends")
}
if peak != mcpFederatedFanoutConcurrency {
t.Errorf("peak concurrent backends = %d, want exactly %d", peak, mcpFederatedFanoutConcurrency)
}
if rec.Code != http.StatusOK {
t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String())
}
}
// --- security audit finding 2: unrecovered panic in the fan-out ---
// panicOnHostRoundTripper panics for any request to panicHost — simulating
// an unrecovered interpreter-level panic reached deep inside one backend's
// own call path (mcpBackendCall -> doBackendJSONRPC -> g.targetClient.Do)
// — and delegates every other request to next unchanged. This exercises
// the REAL code path a genuine Yaegi interpreter panic would take,
// through the actual HTTP client the gateway uses, rather than adding any
// test-only hook to production code.
type panicOnHostRoundTripper struct {
next http.RoundTripper
panicHost string
}
func (rt panicOnHostRoundTripper) RoundTrip(r *http.Request) (*http.Response, error) {
if r.URL.Host == rt.panicHost {
panic("simulated interpreter-level panic mid-flight")
}
return rt.next.RoundTrip(r)
}
// TestHandleMCPFederated_ToolsList_PanickingBackendDegradesNotCrashes is
// the regression for the missing recover() in mcpFederatedToolsList's own
// fan-out goroutine (security audit finding 2, 2026-08-22): that goroutine
// runs off the request's own goroutine, so an unrecovered panic there has
// no ServeHTTP caller to unwind into and would crash the whole shared
// Traefik process, not just fail one backend. Reaching the end of this
// test at all is part of the proof: if the recover() were missing, beta's
// panicking RoundTrip call below would already have crashed the test
// binary — exactly the registry.go precedent this fix follows
// (TestModelRegistry_RefreshProvider_AdapterPanics_RecoversAndReleasesInFlight).
func TestHandleMCPFederated_ToolsList_PanickingBackendDegradesNotCrashes(t *testing.T) {
alpha := newMockJSONRPCServer(t, []mcpTool{{Name: "lookup"}})
beta := newMockJSONRPCServer(t, []mcpTool{{Name: "search"}}) // never actually reached; RoundTrip panics first
cfg := newFederationTestConfig(alpha.srv.URL, beta.srv.URL, false)
h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw")
if err != nil {
t.Fatalf("New: %v", err)
}
gw, ok := h.(*Gateway)
if !ok {
t.Fatal("handler is not *Gateway")
}
betaURL, err := betaHost(beta.srv.URL)
if err != nil {
t.Fatalf("parse beta URL: %v", err)
}
gw.targetClient.Transport = panicOnHostRoundTripper{panicHost: betaURL, next: http.DefaultTransport}
rec := httptest.NewRecorder()
h.ServeHTTP(rec, newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "tools/list", ID: json.RawMessage("1")}))
if rec.Code != http.StatusOK {
t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String())
}
var got jsonrpcResponse
if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil {
t.Fatalf("decode: %v", err)
}
if got.Error != nil {
t.Fatalf("error = %+v, want nil (a panicking server degrades the aggregate, it does not fail the call)", got.Error)
}
var result mcpToolsListResult
if err := json.Unmarshal(got.Result, &result); err != nil {
t.Fatalf("decode result: %v", err)
}
if len(result.Tools) != 1 || result.Tools[0].Name != "alpha_lookup" {
t.Errorf("tools = %+v, want exactly alpha_lookup (beta panicked, skipped)", result.Tools)
}
}
// betaHost extracts the host:port a RoundTripper sees on r.URL.Host for
// requests aimed at rawURL — httptest.Server URLs are already in that
// exact "http://host:port" form, so this just strips the scheme.
func betaHost(rawURL string) (string, error) {
const prefix = "http://"
if !strings.HasPrefix(rawURL, prefix) {
return "", fmt.Errorf("unexpected test server URL shape: %q", rawURL)
}
return strings.TrimPrefix(rawURL, prefix), nil
}
// --- security review round 2, critical finding 3: tools/list is NOT weighted ---
// TestHandleMCPFederated_ToolsList_NotWeightedByBackendCount_BudgetBuysNCalls
// is the regression for security review round 2's critical finding 3:
// weighting checkAndCount by the number of allowed MCP servers (security
// audit finding 1b's original fix) was not default-preserving — on a
// config where the MCP allow-list is unrestricted (every group on the
// live fleet), it silently turned a configured requests-per-minute budget
// into budget/serverCount successful tools/list calls, with NO config
// change on the operator's part (11 servers, requests-per-minute: 60 ->
// only 5 calls/minute actually succeeded; examples/kubernetes.yaml's
// shipped requests-per-minute: 10 would 429 the very FIRST call ever
// made). With numServers (11, matching the live fleet) configured and a
// requests-per-minute budget of exactly rpmBudget, all rpmBudget calls
// must succeed and the next one must be rejected — proving the charge is
// exactly 1 per call, regardless of how many backends that call fans out
// to.
func TestHandleMCPFederated_ToolsList_NotWeightedByBackendCount_BudgetBuysNCalls(t *testing.T) {
const numServers = 11 // matches the live fleet's own MCP server count
const rpmBudget = 5
cfg := CreateConfig()
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}
cfg.MCPServers = make(map[string]*TargetConfig, numServers)
for i := 0; i < numServers; i++ {
srv := newMockJSONRPCServer(t, []mcpTool{{Name: "lookup"}})
cfg.MCPServers[fmt.Sprintf("srv%d", i)] = &TargetConfig{URL: srv.srv.URL}
}
cfg.Groups = map[string]*GroupConfig{"default": {Limits: &LimitsConfig{RequestsPerMinute: rpmBudget}}}
cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}}
h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw")
if err != nil {
t.Fatalf("New: %v", err)
}
for i := 1; i <= rpmBudget; i++ {
rec := httptest.NewRecorder()
h.ServeHTTP(rec, newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "tools/list", ID: json.RawMessage(fmt.Sprintf("%d", i))}))
if rec.Code != http.StatusOK {
t.Fatalf("call %d/%d status = %d, want 200 (budget must buy exactly rpmBudget calls, not rpmBudget/numServers), body=%s", i, rpmBudget, rec.Code, rec.Body.String())
}
}
rec := httptest.NewRecorder()
h.ServeHTTP(rec, newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "tools/list", ID: json.RawMessage(`"over"`)}))
if rec.Code != http.StatusTooManyRequests {
t.Fatalf("call %d status = %d, want 429 (budget exhausted), body=%s", rpmBudget+1, rec.Code, rec.Body.String())
}
}
// TestHandleMCPFederated_ToolsList_AmplificationVisibleAtTargetScopeOnly
// proves the fan-out's real per-server cost is still fully visible to an
// operator without silently consuming tenant quota (security review round
// 2, 2026-08-22, critical finding 3's own ruling): after ONE tools/list
// call against 2 servers, the user's own requests-per-minute counter
// reads exactly 1 (never weighted), while EACH server's own "mcp"
// target-scope counter (countTargetRequests, limits.go) reads exactly 1
// — the real amplification, attributed where an operator can see it.
func TestHandleMCPFederated_ToolsList_AmplificationVisibleAtTargetScopeOnly(t *testing.T) {
alpha := newMockJSONRPCServer(t, []mcpTool{{Name: "lookup"}})
beta := newMockJSONRPCServer(t, []mcpTool{{Name: "search"}})
cfg := newFederationTestConfig(alpha.srv.URL, beta.srv.URL, false)
h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw")
if err != nil {
t.Fatalf("New: %v", err)
}
gw, ok := h.(*Gateway)
if !ok {
t.Fatal("handler is not *Gateway")
}
rec := httptest.NewRecorder()
h.ServeHTTP(rec, newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "tools/list", ID: json.RawMessage("1")}))
if rec.Code != http.StatusOK {
t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String())
}
now := time.Now()
userReq, ok := gw.limiter.getCounter("user", "alice", metricReq, windowDay, now)
if !ok || userReq != 1 {
t.Errorf("user req:day counter = %d (ok=%v), want exactly 1 (never weighted by backend count)", userReq, ok)
}
alphaReq, ok := gw.limiter.getCounter(targetKindMCP, "alpha", metricReq, windowDay, now)
if !ok || alphaReq != 1 {
t.Errorf("mcp/alpha req:day counter = %d (ok=%v), want 1", alphaReq, ok)
}
betaReq, ok := gw.limiter.getCounter(targetKindMCP, "beta", metricReq, windowDay, now)
if !ok || betaReq != 1 {
t.Errorf("mcp/beta req:day counter = %d (ok=%v), want 1", betaReq, ok)
}
}
// --- MF3: legacy HTTP+SSE transport exclusion ---
func TestIsLegacySSETransportURL(t *testing.T) {
+106
View File
@@ -24,6 +24,7 @@ package main
import (
"bufio"
"bytes"
"context"
"encoding/json"
"fmt"
@@ -226,6 +227,30 @@ func run() error {
}))
defer slowUpstream.Close()
// mcpProbeUpstream answers federation's outbound tools/call with a
// body deliberately larger than mcpBackendCallResponseMaxBytes, so
// exerciseHandler's POST /mcp probe below drives doBackendJSONRPC's
// oversize branch — and therefore its errors.Is(err,
// errMCPResponseTooLarge) match — under the REAL interpreter.
//
// This exists because a security review round objected that no gate
// exercised federation's tools/call path interpreted, which left
// errors.Is there justified by reasoning rather than by evidence.
// Reasoning is exactly what the Yaegi traps in this codebase have
// repeatedly defeated, so the path is now covered instead of argued.
mcpProbeUpstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"jsonrpc":"2.0","id":1,"result":{"content":[{"type":"text","text":"`))
chunk := bytes.Repeat([]byte("x"), 1<<20)
for written := 0; written <= mcpProbeOversizeBytes; written += len(chunk) {
if _, writeErr := w.Write(chunk); writeErr != nil {
return
}
}
_, _ = w.Write([]byte(`"}]}}`))
}))
defer mcpProbeUpstream.Close()
// builtinLookupContextTokens (review fix, SHOULD-6): read directly
// out of the generated pricing_data_gen.go rather than hardcoding a
// number, so this assertion self-updates across a `make
@@ -252,6 +277,7 @@ func run() error {
// runs correctly under Yaegi too, not just the config-override
// one testDataWantModel already exercises.
`"modelMeta":{"` + testDataWantModel + `":{"contextTokens":` + yaegiMetaContextTokens + `,"inputCostPerMTokMicroUsd":1250000,"outputCostPerMTokMicroUsd":10000000}},` +
`"mcpServers":{"` + mcpProbeServerName + `":{"url":"` + mcpProbeUpstream.URL + `"}},` +
`"admin":{"enabled":true},"users":{"inline":[{"name":"tester","group":"default","apiKey":"` + testDataUserAPIKey + `"},{"name":"admin1","group":"default","apiKey":"` + attemptAccountingAdminAPIKey + `","admin":true}]}}`
if err = json.Unmarshal([]byte(attemptAccountingOverride), cfgVal.Interface()); err != nil {
return fmt.Errorf("decode attempt-accounting harness override into the interpreted Config: %w", err)
@@ -298,6 +324,32 @@ func run() error {
// the config-override path testDataWantModel already exercises.
const builtinLookupModelID = "gpt-4o"
// mcpProbeServerName is the federated MCP server run() configures against
// mcpProbeUpstream, and mcpProbeToolName is a tool id carrying its
// "<server>_" prefix so resolveFederatedTool routes POST /mcp's
// tools/call there. The tool needs no tools/list entry — resolution is
// pure prefix matching on the configured, group-allowed server names.
const (
mcpProbeServerName = "probe"
mcpProbeToolName = mcpProbeServerName + "_big"
)
// mcpProbeOversizeBytes is how much filler mcpProbeUpstream writes: over
// mcpBackendCallResponseMaxBytes (maxRequestBytes, 10MiB) so the response
// trips doBackendJSONRPC's cap. Kept as its own named value rather than
// importing the plugin's const, because this harness deliberately builds
// nothing from the plugin package at compile time — everything it asserts
// about the plugin must come from the interpreter.
const mcpProbeOversizeBytes = 11 << 20
// mcpTooLargeWantMessage is the exact JSON-RPC error message
// mcpFederatedToolsCall writes when errors.Is matches
// errMCPResponseTooLarge. "upstream error" here instead means errors.Is
// returned false under the interpreter for a chain that is true when
// compiled — which is precisely the class of Yaegi divergence this
// harness exists to catch, and which no compiled test can see.
const mcpTooLargeWantMessage = "response too large"
// readBuiltinContextTokens reads repoRoot/pricing_data_gen.go and
// extracts modelID's own ContextTokens value directly out of the
// generated source (review fix, SHOULD-6) — rather than hardcoding a
@@ -401,6 +453,10 @@ func exerciseHandler(handler http.Handler, builtinContextTokens int) error {
return err
}
if err := exerciseFederatedTooLarge(handler); err != nil {
return err
}
unauthedReq := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
unauthedRec := httptest.NewRecorder()
handler.ServeHTTP(unauthedRec, unauthedReq)
@@ -439,6 +495,56 @@ func exerciseHandler(handler http.Handler, builtinContextTokens int) error {
// a compiled concrete value's method set in Yaegi. matchesSentinel now
// uses real errors.Is instead (its own doc comment in limits.go has the
// full account); this harness is what caught the difference.
// exerciseFederatedTooLarge drives POST /mcp (federation's tools/call
// path) against mcpProbeUpstream, whose response deliberately exceeds
// mcpBackendCallResponseMaxBytes, and asserts the interpreted plugin
// answers with mcpTooLargeWantMessage rather than the generic "upstream
// error".
//
// What this actually pins is errors.Is under the interpreter. A security
// review round replaced that errors.Is with a strings.Contains marker on
// the grounds that no gate exercised this path interpreted — true at the
// time, and the honest objection. The reasoning for errors.Is is sound
// (errors.New and fmt.Errorf both return COMPILED values, so declaring
// the sentinel variable in interpreted code leaves the whole chain
// errors.Is walks compiled, which is why retry.go's isTransient has
// worked in production since v0.2.0), and a standalone yaegi v0.16.1
// harness confirmed it. But this codebase's own history — matchesSentinel
// (limits.go) — is a case where sound reasoning about the reflect
// boundary was wrong and only an interpreted probe caught it. So the
// path is covered here rather than argued in a comment.
//
// A failure surfaces as "upstream error": errors.Is returning false
// interpreted for a chain that is true compiled. No `go test` row can
// see that difference.
func exerciseFederatedTooLarge(handler http.Handler) error {
body := `{"jsonrpc":"2.0","id":7,"method":"tools/call","params":{"name":"` + mcpProbeToolName + `","arguments":{}}}`
req := httptest.NewRequest(http.MethodPost, "/mcp", strings.NewReader(body))
req.Header.Set("Authorization", "Bearer "+testDataUserAPIKey)
req.Header.Set("Content-Type", "application/json")
rec := httptest.NewRecorder()
handler.ServeHTTP(rec, req)
if rec.Code != http.StatusOK {
return fmt.Errorf("POST /mcp tools/call: status = %d, want 200 (JSON-RPC reports the failure in the body, not the status), body=%s", rec.Code, rec.Body.String())
}
var resp struct {
Error *struct {
Message string `json:"message"`
} `json:"error"`
}
if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil {
return fmt.Errorf("decode POST /mcp tools/call body: %w (body=%s)", err, rec.Body.String())
}
if resp.Error == nil {
return fmt.Errorf("POST /mcp tools/call returned no JSON-RPC error, want %q — the probe upstream's oversized response should have tripped the cap: %s", mcpTooLargeWantMessage, rec.Body.String())
}
if resp.Error.Message != mcpTooLargeWantMessage {
return fmt.Errorf("POST /mcp tools/call error message = %q, want %q — errors.Is(err, errMCPResponseTooLarge) did not match under the interpreter (mcp_federation.go)", resp.Error.Message, mcpTooLargeWantMessage)
}
return nil
}
func exerciseAttemptAccounting(handler http.Handler) error {
chatReq := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(
`{"model":"`+testDataWantModel+`","messages":[{"role":"user","content":"hi"}]}`,