diff --git a/admin.go b/admin.go index a7b5dd7..f24544b 100644 --- a/admin.go +++ b/admin.go @@ -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)] diff --git a/admin_test.go b/admin_test.go index e888b9d..1645e99 100644 --- a/admin_test.go +++ b/admin_test.go @@ -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)) + } +} diff --git a/auth.go b/auth.go index 96359b7..c128088 100644 --- a/auth.go +++ b/auth.go @@ -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 " (case-insensitive "bearer" prefix) or // "x-api-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 diff --git a/auth_test.go b/auth_test.go index 40cca70..39dba7d 100644 --- a/auth_test.go +++ b/auth_test.go @@ -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) + } +} diff --git a/logger.go b/logger.go index 1af1c0a..9426105 100644 --- a/logger.go +++ b/logger.go @@ -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. diff --git a/logger_test.go b/logger_test.go new file mode 100644 index 0000000..ff2434f --- /dev/null +++ b/logger_test.go @@ -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) + } +} diff --git a/mcp_federation.go b/mcp_federation.go index d2386e7..37dae6c 100644 --- a/mcp_federation.go +++ b/mcp_federation.go @@ -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 -// "_", 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 "_", 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 } diff --git a/mcp_federation_test.go b/mcp_federation_test.go index 7ee55a2..ed0e5a4 100644 --- a/mcp_federation_test.go +++ b/mcp_federation_test.go @@ -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(¤t, 1) + for { + old := atomic.LoadInt32(&peak) + if n <= old || atomic.CompareAndSwapInt32(&peak, old, n) { + break + } + } + <-release + atomic.AddInt32(¤t, -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(¤t) < 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(¤t); 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) { diff --git a/tools/yaegi-check/main.go b/tools/yaegi-check/main.go index ede1ebb..1261538 100644 --- a/tools/yaegi-check/main.go +++ b/tools/yaegi-check/main.go @@ -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 +// "_" 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"}]}`,