diff --git a/admin_test.go b/admin_test.go index 74ce70e..1645e99 100644 --- a/admin_test.go +++ b/admin_test.go @@ -1998,34 +1998,39 @@ func TestAdminTargets_AccessListsMatchEnforcement_AndCountersReflectTraffic(t *t // --- default-preserving: the live cluster's real shape (security audit, 2026-08-22) --- // TestDefaultPreserving_LiveClusterShape_LoadsAndServesIdentically models -// the live production shape this round's five 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 five findings touch mcp_a2a.go or any -// A2A-specific code path, so there is no mechanism by which they could -// regress agent handling. +// 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 normal -// federated tools/list call from an ordinary user still succeeds under a -// realistic per-minute limit (finding 1b's per-backend weighting does not -// spuriously throttle ordinary usage — 11 servers is comfortably under a -// 60/min budget), 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). +// 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: 60}}, + "friends": {Limits: &LimitsConfig{RequestsPerMinute: friendsRPM}}, "home": {Limits: &LimitsConfig{RequestsPerMinute: 120}}, } @@ -2049,10 +2054,16 @@ func TestDefaultPreserving_LiveClusterShape_LoadsAndServesIdentically(t *testing t.Fatalf("New: want the realistic live-cluster config to load unchanged, got error: %v", err) } - rec := httptest.NewRecorder() - h.ServeHTTP(rec, newFederatedRequest(t, "sk-steve-cw", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "tools/list", ID: json.RawMessage("1")})) - if rec.Code != http.StatusOK { - t.Fatalf("tools/list status = %d, want 200, body=%s", rec.Code, rec.Body.String()) + // 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) diff --git a/limits.go b/limits.go index 013b6b7..61ae8c9 100644 --- a/limits.go +++ b/limits.go @@ -1081,39 +1081,14 @@ func buildBudgetProbes(scopes []limitScope, now time.Time) (probes []budgetProbe // call completes, one request may still push a scope over its budget — an // accepted trade-off, unchanged from before this round. func (l *limiter) checkAndCount(scopes []limitScope) *limitViolation { - return l.checkAndCountWeighted(scopes, 1) -} - -// checkAndCountWeighted is checkAndCount generalized to charge `weight` -// requests against every scope's requests-per-minute/requests-per-day -// counters instead of exactly one (security audit finding 1b, 2026-08-22): -// a route whose single incoming client request causes this gateway to -// issue `weight` outbound backend calls — today, only the federated MCP -// tools/list fan-out (mcp_federation.go's handleMCPFederated, weight = -// len(allowedMCPServerNames(grp))) — must charge its real cost, not -// undercount a fan-out as a single request. weight<=0 is treated as 1 -// (never zero-cost, never negative); checkAndCount(scopes) is exactly -// checkAndCountWeighted(scopes, 1), so every other call site's behavior is -// completely unchanged by this method existing. Semantics beyond the -// weighted delta are identical to checkAndCount's own doc comment in -// every other respect: one storeIncrAndGetMulti round trip, requests- -// per-minute/requests-per-day compared against the just-incremented -// (now weighted) count, token/cost budgets read and compared unweighted -// (a fan-out has no tokens or USD cost of its own — those probes come -// from buildBudgetProbes, untouched by weight), violation precedence -// unchanged. -func (l *limiter) checkAndCountWeighted(scopes []limitScope, weight int64) *limitViolation { - if weight <= 0 { - weight = 1 - } now := l.now() entries := make([]counterIncr, 0, len(scopes)*checkAndCountKeysPerScope) for _, sc := range scopes { entries = append(entries, - newCounterIncr(sc.kind, sc.id, metricReq, windowMin, now, weight, minWindowTTL), - newCounterIncr(sc.kind, sc.id, metricReq, windowDay, now, weight, dayWindowTTL), - newCounterIncr(sc.kind, sc.id, metricReq, windowHour, now, weight, hourWindowTTL), + newCounterIncr(sc.kind, sc.id, metricReq, windowMin, now, 1, minWindowTTL), + newCounterIncr(sc.kind, sc.id, metricReq, windowDay, now, 1, dayWindowTTL), + newCounterIncr(sc.kind, sc.id, metricReq, windowHour, now, 1, hourWindowTTL), ) } probes, reads := buildBudgetProbes(scopes, now) diff --git a/mcp_federation.go b/mcp_federation.go index 45e7782..25e3a6d 100644 --- a/mcp_federation.go +++ b/mcp_federation.go @@ -67,20 +67,51 @@ const ( ) // mcpBackendResponseMaxBytes caps a single backend MCP server's HTTP -// response body doBackendJSONRPC will read into memory — 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) — -// comfortably above any tools/list schema listing or tools/call result -// this gateway has actually seen live (typically low KB), while bounding -// mcpFederatedFanoutConcurrency backends' worst-case concurrent memory to -// 4 × 4MiB = 16MiB, against the unbounded ~110MiB (11 production servers -// × up to 10MiB each, all fired at once) the security audit measured -// before this fix. +// 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 + +// mcpResponseTooLargeMarker is a fixed, unique substring doBackendJSONRPC +// embeds in its own error message when a backend response hits its +// call's maxBytes cap (security review round 2, 2026-08-22, important +// finding 4) — detected via strings.Contains, deliberately NOT an +// errors.Is-checked sentinel. matchesSentinel's own doc comment +// (limits.go) proves errors.Is safe under Yaegi only for chains reaching +// a call site this package's own yaegi-check harness actually exercises +// at runtime under the real interpreter — federation's tools/call path +// is not one of those exercised paths today, so a plain string check +// trades a little idiom for zero interpreter risk here, rather than +// relying on reasoning about compiled-vs-interpreted-origin values that +// nothing in the gate actually verifies for this specific path. +const mcpResponseTooLargeMarker = "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). @@ -249,27 +280,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). // -// The weight charged to checkAndCount is 1 for every method EXCEPT -// tools/list (security audit finding 1b, 2026-08-22): a tools/list call -// fans out to every server allowedMCPServerNames(grp) returns — up to the -// whole MCP catalog, concurrently (mcpFederatedToolsList's own doc -// comment) — so one incoming client request can cost this gateway many -// real outbound backend calls. Charging it as a single request against -// the caller's requests-per-minute/requests-per-day budget undercounts -// that real cost; weighting it by len(names) instead makes the budget -// reflect what the request actually does. This is computed and charged -// exactly once, right here, before dispatch — never doubled by anything -// mcpFederatedToolsList itself does (its own countTargetRequests call is a -// SEPARATE scope kind, "mcp"/target-id, entirely disjoint from the -// user/group/total scopes checkAndCountWeighted evaluates here; see -// countTargetRequests' own doc comment, limits.go). tools/call, by -// contrast, always resolves to exactly ONE backend (resolveFederatedTool), -// so it keeps the plain weight of 1 — checkAndCountWeighted(scopes, 1) is -// byte-for-byte checkAndCount's own existing behavior. +// 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") @@ -292,17 +328,7 @@ func (g *Gateway) handleMCPFederated(w http.ResponseWriter, r *http.Request, u * } scopes := withTotalScope(buildLimitScopes(u, grp)) - - var toolsListNames []string - weight := int64(1) - if req.Method == "tools/list" { - toolsListNames = allowedMCPServerNames(g.cfg, grp) - if n := int64(len(toolsListNames)); n > weight { - weight = n - } - } - - if violation := g.limiter.checkAndCountWeighted(scopes, weight); violation != nil { + if violation := g.limiter.checkAndCount(scopes); violation != nil { writeLimitViolation(w, violation) return } @@ -313,7 +339,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, toolsListNames) + 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/"): @@ -413,8 +439,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 with mcpResponseTooLargeMarker in its message, 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. @@ -439,7 +473,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 @@ -476,10 +510,17 @@ 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, mcpBackendResponseMaxBytes)) + // 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 } + if int64(len(respBody)) > maxBytes { + return nil, status, respSessionID, fmt.Errorf("%s: exceeds %d bytes", mcpResponseTooLargeMarker, maxBytes) + } if status < 200 || status >= 300 { return nil, status, respSessionID, fmt.Errorf("upstream returned HTTP %d", status) } @@ -500,13 +541,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 } @@ -567,9 +608,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 } @@ -579,7 +625,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 @@ -600,7 +646,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 } @@ -685,19 +731,28 @@ func lastSSEDataLine(body []byte) []byte { // 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. 4 keeps -// every server eventually contacted — a goroutine blocks on the -// semaphore, it is never dropped or skipped — while capping the worst -// case to 4 concurrent backends' cost, not the whole catalog's. -const mcpFederatedFanoutConcurrency = 4 +// 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 in names (handleMCPFederated's own -// allowedMCPServerNames(grp) call, made once there so the SAME set that -// was weighed into checkAndCountWeighted is the set actually contacted), -// at most mcpFederatedFanoutConcurrency at a time — each under its own -// toolsListBackendTimeout budget (context.WithTimeout off r.Context(), -// MF1) — then merges every reachable server's own tools into one result, +// 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 @@ -725,6 +780,24 @@ func (g *Gateway) mcpFederatedToolsList(w http.ResponseWriter, r *http.Request, 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. + fanoutCtx, fanoutCancel := context.WithTimeout(r.Context(), toolsListBackendTimeout) + defer fanoutCancel() + sem := make(chan struct{}, mcpFederatedFanoutConcurrency) for _, name := range names { wg.Add(1) @@ -752,10 +825,8 @@ func (g *Gateway) mcpFederatedToolsList(w http.ResponseWriter, r *http.Request, sem <- struct{}{} defer func() { <-sem }() - ctx, cancel := context.WithTimeout(r.Context(), toolsListBackendTimeout) - defer cancel() 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() @@ -867,11 +938,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 strings.Contains(err.Error(), mcpResponseTooLargeMarker) { + 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 dd19b17..5bf0582 100644 --- a/mcp_federation_test.go +++ b/mcp_federation_test.go @@ -1031,7 +1031,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 { @@ -1060,12 +1060,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 with mcpResponseTooLargeMarker in its message — +// distinguishable via strings.Contains 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 !strings.Contains(err.Error(), mcpResponseTooLargeMarker) { + t.Errorf("err = %v, want it to contain %q", err, mcpResponseTooLargeMarker) + } +} + +// 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 @@ -1391,100 +1513,96 @@ func betaHost(rawURL string) (string, error) { return strings.TrimPrefix(rawURL, prefix), nil } -// --- security audit finding 1b: tools/list weighted by backend count --- +// --- security review round 2, critical finding 3: tools/list is NOT weighted --- -// TestHandleMCPFederated_ToolsList_WeightedByBackendCount_RejectsWhenOverBudget -// proves a tools/list call charges checkAndCount len(allowedMCPServerNames) -// requests, not one: with 3 configured (and unrestricted) servers and a -// requests-per-minute budget of 2, the very FIRST tools/list call must be -// rejected outright — and, because checkAndCountWeighted runs and fails -// BEFORE the fan-out ever starts, none of the three servers may be -// contacted at all. -func TestHandleMCPFederated_ToolsList_WeightedByBackendCount_RejectsWhenOverBudget(t *testing.T) { - alpha := newMockJSONRPCServer(t, nil) - beta := newMockJSONRPCServer(t, nil) - gamma := newMockJSONRPCServer(t, nil) - cfg := newFederationTestConfig(alpha.srv.URL, beta.srv.URL, false) - cfg.MCPServers["gamma"] = &TargetConfig{URL: gamma.srv.URL} - cfg.Groups["default"].Limits = &LimitsConfig{RequestsPerMinute: 2} +// 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.StatusTooManyRequests { - t.Fatalf("status = %d, want 429 (weight=3 > limit=2), body=%s", rec.Code, rec.Body.String()) - } - for name, srv := range map[string]*mockJSONRPCServer{"alpha": alpha, "beta": beta, "gamma": gamma} { - if srv.called.Load() != 0 { - t.Errorf("server %q was contacted %d times, want 0 (rejected before fan-out)", name, srv.called.Load()) - } - } -} - -// TestHandleMCPFederated_ToolsList_WeightedByBackendCount_AccumulatesAcrossCalls -// proves the weighted charge actually accumulates against the same -// requests-per-minute window: with 3 servers and a budget of 5, the first -// tools/list call (weight 3) succeeds, and the second (weight 3 more, -// total 6 > 5) is rejected. -func TestHandleMCPFederated_ToolsList_WeightedByBackendCount_AccumulatesAcrossCalls(t *testing.T) { - alpha := newMockJSONRPCServer(t, []mcpTool{{Name: "lookup"}}) - beta := newMockJSONRPCServer(t, []mcpTool{{Name: "search"}}) - gamma := newMockJSONRPCServer(t, []mcpTool{{Name: "fetch"}}) - cfg := newFederationTestConfig(alpha.srv.URL, beta.srv.URL, false) - cfg.MCPServers["gamma"] = &TargetConfig{URL: gamma.srv.URL} - cfg.Groups["default"].Limits = &LimitsConfig{RequestsPerMinute: 5} - - h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw") - if err != nil { - t.Fatalf("New: %v", err) + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) } - rec1 := httptest.NewRecorder() - h.ServeHTTP(rec1, newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "tools/list", ID: json.RawMessage("1")})) - if rec1.Code != http.StatusOK { - t.Fatalf("first call status = %d, want 200, body=%s", rec1.Code, rec1.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) } - - rec2 := httptest.NewRecorder() - h.ServeHTTP(rec2, newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "tools/list", ID: json.RawMessage("2")})) - if rec2.Code != http.StatusTooManyRequests { - t.Fatalf("second call status = %d, want 429 (3+3=6 > limit=5), body=%s", rec2.Code, rec2.Body.String()) + 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) } -} - -// TestHandleMCPFederated_ToolsCall_WeightStaysOne_EvenWithManyServersConfigured -// proves tools/call — which only ever resolves to ONE backend — never -// inherits tools/list's per-backend weighting: with 3 servers configured -// and a requests-per-minute budget of 2, two sequential tools/call -// requests must both succeed (weight 1 each, total 2, at the limit), where -// the equivalent tools/list weighting (3 each) would already have -// rejected the first one. -func TestHandleMCPFederated_ToolsCall_WeightStaysOne_EvenWithManyServersConfigured(t *testing.T) { - alpha := newMockJSONRPCServer(t, nil) - beta := newMockJSONRPCServer(t, nil) - gamma := newMockJSONRPCServer(t, nil) - cfg := newFederationTestConfig(alpha.srv.URL, beta.srv.URL, false) - cfg.MCPServers["gamma"] = &TargetConfig{URL: gamma.srv.URL} - cfg.Groups["default"].Limits = &LimitsConfig{RequestsPerMinute: 2} - - 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 <= 2; i++ { - rec := httptest.NewRecorder() - h.ServeHTTP(rec, newFederatedRequest(t, "sk-alice", jsonrpcRequest{ - JSONRPC: jsonrpcVersion, Method: "tools/call", ID: json.RawMessage(fmt.Sprintf("%d", i)), - Params: json.RawMessage(`{"name":"alpha_lookup"}`), - })) - if rec.Code != http.StatusOK { - t.Fatalf("call %d status = %d, want 200 (tools/call weight must stay 1), body=%s", i, rec.Code, rec.Body.String()) - } + 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) } }