diff --git a/limits.go b/limits.go index 61ae8c9..013b6b7 100644 --- a/limits.go +++ b/limits.go @@ -1081,14 +1081,39 @@ 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, 1, minWindowTTL), - newCounterIncr(sc.kind, sc.id, metricReq, windowDay, now, 1, dayWindowTTL), - newCounterIncr(sc.kind, sc.id, metricReq, windowHour, now, 1, hourWindowTTL), + 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), ) } probes, reads := buildBudgetProbes(scopes, now) diff --git a/mcp_federation.go b/mcp_federation.go index d2386e7..d4ae414 100644 --- a/mcp_federation.go +++ b/mcp_federation.go @@ -66,6 +66,21 @@ const ( mcpSessionCloseTimeout = 5 * time.Second ) +// 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. +const mcpBackendResponseMaxBytes = 4 << 20 + // 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). @@ -237,6 +252,24 @@ func allowedMCPServerNames(cfg *Config, grp *group) []string { // other metered route — 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. 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") @@ -259,7 +292,17 @@ func (g *Gateway) handleMCPFederated(w http.ResponseWriter, r *http.Request, u * } scopes := withTotalScope(buildLimitScopes(u, grp)) - if violation := g.limiter.checkAndCount(scopes); violation != nil { + + 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 { writeLimitViolation(w, violation) return } @@ -270,7 +313,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, toolsListNames) case req.Method == "tools/call": g.mcpFederatedToolsCall(w, r, req, grp) case strings.HasPrefix(req.Method, "notifications/"): @@ -433,7 +476,7 @@ 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)) + respBody, err := io.ReadAll(io.LimitReader(httpResp.Body, mcpBackendResponseMaxBytes)) if err != nil { return nil, status, respSessionID, err } @@ -632,43 +675,64 @@ 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, the same bounded-concurrency shape +// limiter.spawnTokens already uses for its own goroutine fan-out +// (limits.go), applied here to gate each goroutine's actual outbound HTTP +// work rather than its launch. 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. 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 + // 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 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, +// 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. // // 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, 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, names []string) { var ( mu sync.Mutex wg sync.WaitGroup merged = make([]mcpTool, 0, len(names)) failed []string ) + sem := make(chan struct{}, mcpFederatedFanoutConcurrency) for _, name := range names { wg.Add(1) go func(name string) { defer wg.Done() + sem <- struct{}{} + defer func() { <-sem }() + ctx, cancel := context.WithTimeout(r.Context(), toolsListBackendTimeout) defer cancel() targetURL := g.cfg.MCPServers[name].URL diff --git a/mcp_federation_test.go b/mcp_federation_test.go index 7ee55a2..065dd32 100644 --- a/mcp_federation_test.go +++ b/mcp_federation_test.go @@ -4,6 +4,7 @@ import ( "bytes" "context" "encoding/json" + "fmt" "io" "net/http" "net/http/httptest" @@ -1214,6 +1215,196 @@ 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 1b: tools/list weighted by backend count --- + +// 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} + + 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/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) + } + + 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()) + } + + 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()) + } +} + +// 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()) + } + } +} + // --- MF3: legacy HTTP+SSE transport exclusion --- func TestIsLegacySSETransportURL(t *testing.T) {