mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-19 05:47:48 +00:00
fix(mcp): stop weighting tools/list budget, widen fan-out, split caps
Addresses critical finding 3 and important findings 3-4 from security review round 2, 2026-08-22. Critical finding 3 — the tools/list-per-backend weight charged against the caller's own requests-per-minute/day budget (security audit finding 1b) was not default-preserving: both live groups have an empty MCP allow-list (all 11 servers), so it silently cut every configured requests-per-minute budget by ~11x with no config change, and examples/kubernetes.yaml's shipped requests-per-minute: 10 would 429 the very first tools/list call ever made. Reverted: checkAndCountWeighted removed from limits.go, handleMCPFederated back to a plain checkAndCount. The real amplification stays fully visible to an operator via countTargetRequests' existing per-server 'mcp' target-scope counters — a separate scope from the user/group budget — rather than silently consuming tenant quota. The fan-out semaphore (below) is the actual protection against the amplification finding 1 described. Important finding 3 — the fan-out semaphore (4 concurrent backends) tripled worst-case latency (907ms measured at 300ms/server, was ~300ms) and pushed the hung-backend worst case to ceil(11/4)x20s=60s, past the ~30s timeout many MCP clients use. Raised to 8, and every backend in one tools/list call now shares ONE overall toolsListBackendTimeout deadline (fanoutCtx) instead of each getting a fresh one from its own staggered start — restoring this function's original 'bounds the whole call regardless of server count' invariant even with bounded concurrency. Important finding 4 — the 4MiB response cap (sized for the tools/list fan-out's memory-amplification risk) silently applied to tools/call too, where exactly one backend is ever contacted and a legitimate result (an image, an extracted document) can be much larger. tools/call now uses mcpBackendCallResponseMaxBytes (10MiB, the pre-finding-1c budget); tools/list keeps the 4MiB cap. A response that hits either cap is detected via a cap+1 read (never silently truncated into a confusing parse failure) and reported to the client as 'response too large'.
This commit is contained in:
+33
-22
@@ -1998,34 +1998,39 @@ func TestAdminTargets_AccessListsMatchEnforcement_AndCountersReflectTraffic(t *t
|
|||||||
// --- default-preserving: the live cluster's real shape (security audit, 2026-08-22) ---
|
// --- default-preserving: the live cluster's real shape (security audit, 2026-08-22) ---
|
||||||
|
|
||||||
// TestDefaultPreserving_LiveClusterShape_LoadsAndServesIdentically models
|
// TestDefaultPreserving_LiveClusterShape_LoadsAndServesIdentically models
|
||||||
// the live production shape this round's five fixes must not break
|
// the live production shape this round's fixes must not break (project
|
||||||
// (project brief): 5 "friend" users plus 4 "home" users (steve-cw and
|
// brief): 5 "friend" users plus 4 "home" users (steve-cw and kevinsandom
|
||||||
// kevinsandom are named live users; the rest are representative — real
|
// are named live users; the rest are representative — real names this
|
||||||
// names this test has no visibility into, which is exactly why finding
|
// test has no visibility into, which is exactly why finding 5's fix stops
|
||||||
// 5's fix stops short of restricting the name charset, see buildEntry's
|
// short of restricting the name charset, see buildEntry's own doc
|
||||||
// own doc comment, auth.go) and 11 MCP servers. A2A agents are out of
|
// comment, auth.go) and 11 MCP servers. A2A agents are out of scope here:
|
||||||
// scope here: none of this round's five findings touch mcp_a2a.go or any
|
// none of this round's fixes touch mcp_a2a.go or any A2A-specific code
|
||||||
// A2A-specific code path, so there is no mechanism by which they could
|
// path, so there is no mechanism by which they could regress agent
|
||||||
// regress agent handling.
|
// handling.
|
||||||
//
|
//
|
||||||
// It proves construction succeeds unchanged (finding 5's new empty/
|
// It proves construction succeeds unchanged (finding 5's new empty/
|
||||||
// duplicate-name checks accept every one of these real names), a normal
|
// duplicate-name checks accept every one of these real names), a
|
||||||
// federated tools/list call from an ordinary user still succeeds under a
|
// friends-group user's configured requests-per-minute budget still buys
|
||||||
// realistic per-minute limit (finding 1b's per-backend weighting does not
|
// EXACTLY that many tools/list calls against all 11 servers — not
|
||||||
// spuriously throttle ordinary usage — 11 servers is comfortably under a
|
// budget/11 — regressing security review round 2's critical finding 3
|
||||||
// 60/min budget), and the admin dashboard's usage poll still returns
|
// (weighting tools/list by backend count was reverted specifically
|
||||||
// exactly one row per user and per group in a single store round trip
|
// because it was not default-preserving at this exact shape: both live
|
||||||
// (finding 4's chunking is a no-op at 10 users, well under
|
// groups have an empty MCP allow-list, i.e. all 11 servers, so a naive
|
||||||
// adminUsageChunkScopes).
|
// 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) {
|
func TestDefaultPreserving_LiveClusterShape_LoadsAndServesIdentically(t *testing.T) {
|
||||||
friendNames := []string{"steve-cw", "kevinsandom", "friend3", "friend4", "friend5"}
|
friendNames := []string{"steve-cw", "kevinsandom", "friend3", "friend4", "friend5"}
|
||||||
homeNames := []string{"home1", "home2", "home3", "home4"}
|
homeNames := []string{"home1", "home2", "home3", "home4"}
|
||||||
|
const friendsRPM = 60
|
||||||
|
|
||||||
cfg := CreateConfig()
|
cfg := CreateConfig()
|
||||||
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}
|
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}
|
||||||
cfg.Admin = &AdminConfig{Enabled: true}
|
cfg.Admin = &AdminConfig{Enabled: true}
|
||||||
cfg.Groups = map[string]*GroupConfig{
|
cfg.Groups = map[string]*GroupConfig{
|
||||||
"friends": {Limits: &LimitsConfig{RequestsPerMinute: 60}},
|
"friends": {Limits: &LimitsConfig{RequestsPerMinute: friendsRPM}},
|
||||||
"home": {Limits: &LimitsConfig{RequestsPerMinute: 120}},
|
"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)
|
t.Fatalf("New: want the realistic live-cluster config to load unchanged, got error: %v", err)
|
||||||
}
|
}
|
||||||
|
|
||||||
rec := httptest.NewRecorder()
|
// Drives friendsRPM calls, not just one (security review round 2:
|
||||||
h.ServeHTTP(rec, newFederatedRequest(t, "sk-steve-cw", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "tools/list", ID: json.RawMessage("1")}))
|
// a single-call assertion cannot distinguish "budget buys N calls"
|
||||||
if rec.Code != http.StatusOK {
|
// from "budget buys N/11 calls" — exactly the blind spot that let the
|
||||||
t.Fatalf("tools/list status = %d, want 200, body=%s", rec.Code, rec.Body.String())
|
// 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 := httptest.NewRequest("GET", adminUsagePath, nil)
|
||||||
|
|||||||
@@ -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
|
// call completes, one request may still push a scope over its budget — an
|
||||||
// accepted trade-off, unchanged from before this round.
|
// accepted trade-off, unchanged from before this round.
|
||||||
func (l *limiter) checkAndCount(scopes []limitScope) *limitViolation {
|
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()
|
now := l.now()
|
||||||
|
|
||||||
entries := make([]counterIncr, 0, len(scopes)*checkAndCountKeysPerScope)
|
entries := make([]counterIncr, 0, len(scopes)*checkAndCountKeysPerScope)
|
||||||
for _, sc := range scopes {
|
for _, sc := range scopes {
|
||||||
entries = append(entries,
|
entries = append(entries,
|
||||||
newCounterIncr(sc.kind, sc.id, metricReq, windowMin, now, weight, minWindowTTL),
|
newCounterIncr(sc.kind, sc.id, metricReq, windowMin, now, 1, minWindowTTL),
|
||||||
newCounterIncr(sc.kind, sc.id, metricReq, windowDay, now, weight, dayWindowTTL),
|
newCounterIncr(sc.kind, sc.id, metricReq, windowDay, now, 1, dayWindowTTL),
|
||||||
newCounterIncr(sc.kind, sc.id, metricReq, windowHour, now, weight, hourWindowTTL),
|
newCounterIncr(sc.kind, sc.id, metricReq, windowHour, now, 1, hourWindowTTL),
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
probes, reads := buildBudgetProbes(scopes, now)
|
probes, reads := buildBudgetProbes(scopes, now)
|
||||||
|
|||||||
+152
-68
@@ -67,20 +67,51 @@ const (
|
|||||||
)
|
)
|
||||||
|
|
||||||
// mcpBackendResponseMaxBytes caps a single backend MCP server's HTTP
|
// mcpBackendResponseMaxBytes caps a single backend MCP server's HTTP
|
||||||
// response body doBackendJSONRPC will read into memory — previously
|
// response body for the tools/list FAN-OUT path only (mcpFederatedToolsList)
|
||||||
// maxRequestBytes (10MiB, routes_unified.go), a budget sized for a
|
// — previously maxRequestBytes (10MiB, routes_unified.go), a budget sized
|
||||||
// CLIENT's own chat/embeddings request body, not a backend's JSON-RPC
|
// for a CLIENT's own chat/embeddings request body, not a backend's
|
||||||
// reply (security audit finding 1c, 2026-08-22). 4MiB matches this
|
// JSON-RPC reply (security audit finding 1c, 2026-08-22). 4MiB matches
|
||||||
// codebase's other "generous but bounded" read caps for a non-primary-
|
// this codebase's other "generous but bounded" read caps for a non-
|
||||||
// request-body read (cache.go's maxAccountingTeeBytes is also 4MiB) —
|
// primary-request-body read (cache.go's maxAccountingTeeBytes is also
|
||||||
// comfortably above any tools/list schema listing or tools/call result
|
// 4MiB) — measured at 547KiB for a real tools/list response describing
|
||||||
// this gateway has actually seen live (typically low KB), while bounding
|
// 20 fat tools (security review round 2, 2026-08-22), comfortably below
|
||||||
// mcpFederatedFanoutConcurrency backends' worst-case concurrent memory to
|
// this cap — while bounding mcpFederatedFanoutConcurrency backends'
|
||||||
// 4 × 4MiB = 16MiB, against the unbounded ~110MiB (11 production servers
|
// worst-case concurrent memory to a small multiple of 4MiB, against the
|
||||||
// × up to 10MiB each, all fired at once) the security audit measured
|
// unbounded ~110MiB (11 production servers × up to 10MiB each, all fired
|
||||||
// before this fix.
|
// 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
|
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
|
// 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
|
// 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).
|
// 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 —
|
// checkAndCount runs once per incoming request, before method dispatch —
|
||||||
// covering the caller's own user/group/total scopes exactly like every
|
// 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
|
// inside mcpFederatedToolsList/mcpFederatedToolsCall (their own doc
|
||||||
// comments cover exactly what "attributed" means for each).
|
// comments cover exactly what "attributed" means for each).
|
||||||
//
|
//
|
||||||
// The weight charged to checkAndCount is 1 for every method EXCEPT
|
// tools/list's fan-out is deliberately NOT charged against the caller's
|
||||||
// tools/list (security audit finding 1b, 2026-08-22): a tools/list call
|
// own requests-per-minute/requests-per-day budget (security review round
|
||||||
// fans out to every server allowedMCPServerNames(grp) returns — up to the
|
// 2, 2026-08-22, critical finding 3 — reverting security audit finding
|
||||||
// whole MCP catalog, concurrently (mcpFederatedToolsList's own doc
|
// 1b's weighted-charge attempt): weighting checkAndCount by the number of
|
||||||
// comment) — so one incoming client request can cost this gateway many
|
// allowed MCP servers is not default-preserving — on this operator's own
|
||||||
// real outbound backend calls. Charging it as a single request against
|
// live fleet, both configured groups have an empty MCP allow-list (every
|
||||||
// the caller's requests-per-minute/requests-per-day budget undercounts
|
// server allowed), so EVERY tools/list call would have weighed 11,
|
||||||
// that real cost; weighting it by len(names) instead makes the budget
|
// silently turning a configured requests-per-minute: 60 into an effective
|
||||||
// reflect what the request actually does. This is computed and charged
|
// budget of 5 successful tools/list calls per minute, and
|
||||||
// exactly once, right here, before dispatch — never doubled by anything
|
// examples/kubernetes.yaml's shipped requests-per-minute: 10 would have
|
||||||
// mcpFederatedToolsList itself does (its own countTargetRequests call is a
|
// 429'd the very FIRST tools/list call ever made — with no config change
|
||||||
// SEPARATE scope kind, "mcp"/target-id, entirely disjoint from the
|
// on the operator's part. The fan-out's real amplification is still
|
||||||
// user/group/total scopes checkAndCountWeighted evaluates here; see
|
// fully visible to an operator without silently consuming tenant quota:
|
||||||
// countTargetRequests' own doc comment, limits.go). tools/call, by
|
// mcpFederatedToolsList's own countTargetRequests call attributes one
|
||||||
// contrast, always resolves to exactly ONE backend (resolveFederatedTool),
|
// request to EVERY server actually attempted, at the "mcp"/target-id
|
||||||
// so it keeps the plain weight of 1 — checkAndCountWeighted(scopes, 1) is
|
// scope (limits.go) — a SEPARATE scope kind from the user/group/total
|
||||||
// byte-for-byte checkAndCount's own existing behavior.
|
// 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) {
|
func (g *Gateway) handleMCPFederated(w http.ResponseWriter, r *http.Request, u *user, grp *group) {
|
||||||
if r.Header.Get("Upgrade") != "" {
|
if r.Header.Get("Upgrade") != "" {
|
||||||
writeOAIError(w, http.StatusNotImplemented, "invalid_request_error", "websocket/upgrade passthrough not supported")
|
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))
|
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)
|
writeLimitViolation(w, violation)
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
@@ -313,7 +339,7 @@ func (g *Gateway) handleMCPFederated(w http.ResponseWriter, r *http.Request, u *
|
|||||||
case req.Method == "ping":
|
case req.Method == "ping":
|
||||||
writeJSONRPCResult(w, req.ID, map[string]any{})
|
writeJSONRPCResult(w, req.ID, map[string]any{})
|
||||||
case req.Method == "tools/list":
|
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":
|
case req.Method == "tools/call":
|
||||||
g.mcpFederatedToolsCall(w, r, req, grp)
|
g.mcpFederatedToolsCall(w, r, req, grp)
|
||||||
case strings.HasPrefix(req.Method, "notifications/"):
|
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
|
// doBackendJSONRPC issues one JSON-RPC 2.0 POST to targetURL — an MCP
|
||||||
// server's own configured base URL, the same single endpoint
|
// server's own configured base URL, the same single endpoint
|
||||||
// handleTargetProxy reverse-proxies every method to — carrying sessionID
|
// handleTargetProxy reverse-proxies every method to — carrying sessionID
|
||||||
// as the Mcp-Session-Id request header when non-empty. It is the one place
|
// as the Mcp-Session-Id request header when non-empty, and reading at
|
||||||
// that actually builds and sends an HTTP request for this file; every
|
// 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,
|
// other function in the handshake-fallback machinery (mcpBackendCall,
|
||||||
// mcpBackendHandshake, mcpBackendCloseSession's own DELETE aside) goes
|
// mcpBackendHandshake, mcpBackendCloseSession's own DELETE aside) goes
|
||||||
// through this.
|
// through this.
|
||||||
@@ -439,7 +473,7 @@ type mcpToolsListResult struct {
|
|||||||
// that WAS valid JSON but not a real JSON-RPC response (see
|
// that WAS valid JSON but not a real JSON-RPC response (see
|
||||||
// parseBackendJSONRPC's own validity check) would have silently produced
|
// parseBackendJSONRPC's own validity check) would have silently produced
|
||||||
// an all-zero-fields response otherwise.
|
// 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)
|
paramsRaw, err := json.Marshal(params)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, 0, "", err
|
return nil, 0, "", err
|
||||||
@@ -476,10 +510,17 @@ func (g *Gateway) doBackendJSONRPC(ctx context.Context, targetURL, method string
|
|||||||
status = httpResp.StatusCode
|
status = httpResp.StatusCode
|
||||||
respSessionID = httpResp.Header.Get(mcpSessionHeader)
|
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 {
|
if err != nil {
|
||||||
return nil, status, respSessionID, err
|
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 {
|
if status < 200 || status >= 300 {
|
||||||
return nil, status, respSessionID, fmt.Errorf("upstream returned HTTP %d", status)
|
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
|
// the initialize call itself failed outright (network error, non-2xx, a
|
||||||
// JSON-RPC error response to initialize itself) — there is no retry of a
|
// JSON-RPC error response to initialize itself) — there is no retry of a
|
||||||
// failed handshake.
|
// 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{
|
initParams := map[string]any{
|
||||||
"protocolVersion": defaultMCPProtocolVersion,
|
"protocolVersion": defaultMCPProtocolVersion,
|
||||||
"capabilities": map[string]any{},
|
"capabilities": map[string]any{},
|
||||||
"clientInfo": map[string]any{"name": "traefik-llmgateway", "version": pluginVersion},
|
"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 {
|
if err != nil {
|
||||||
return "", err
|
return "", err
|
||||||
}
|
}
|
||||||
@@ -567,9 +608,14 @@ func (g *Gateway) mcpBackendCloseSession(ctx context.Context, targetURL, session
|
|||||||
// except the close itself — see mcpSessionCloseTimeout's own doc
|
// except the close itself — see mcpSessionCloseTimeout's own doc
|
||||||
// comment); callers derive it from context.WithTimeout(r.Context(), ...)
|
// comment); callers derive it from context.WithTimeout(r.Context(), ...)
|
||||||
// with toolsListBackendTimeout or toolsCallBackendTimeout (MF1, review
|
// with toolsListBackendTimeout or toolsCallBackendTimeout (MF1, review
|
||||||
// round 2, 2026-08-21).
|
// round 2, 2026-08-21). maxBytes is doBackendJSONRPC's own response-size
|
||||||
func (g *Gateway) mcpBackendCall(ctx context.Context, targetURL, method string, params any) (*jsonrpcResponse, error) {
|
// cap, passed straight through to every doBackendJSONRPC call this
|
||||||
resp, status, _, err := g.doBackendJSONRPC(ctx, targetURL, method, params, "")
|
// 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 {
|
if err == nil && resp.Error == nil {
|
||||||
return resp, nil
|
return resp, nil
|
||||||
}
|
}
|
||||||
@@ -579,7 +625,7 @@ func (g *Gateway) mcpBackendCall(ctx context.Context, targetURL, method string,
|
|||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
|
|
||||||
sessionID, initErr := g.mcpBackendHandshake(ctx, targetURL)
|
sessionID, initErr := g.mcpBackendHandshake(ctx, targetURL, maxBytes)
|
||||||
if initErr != nil {
|
if initErr != nil {
|
||||||
// The handshake itself failed: surface the bare attempt's own
|
// The handshake itself failed: surface the bare attempt's own
|
||||||
// outcome when it has one (a JSON-RPC error response is more
|
// 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 {
|
if retryErr != nil {
|
||||||
return nil, retryErr
|
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,
|
// 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
|
// one client call could hold up to 11 concurrent connections and up to
|
||||||
// 11× mcpBackendResponseMaxBytes of backend response bodies in memory at
|
// 11× mcpBackendResponseMaxBytes of backend response bodies in memory at
|
||||||
// once, entirely outside this gateway's own admission control. 4 keeps
|
// once, entirely outside this gateway's own admission control. Every
|
||||||
// every server eventually contacted — a goroutine blocks on the
|
// server is eventually contacted regardless of this bound — a goroutine
|
||||||
// semaphore, it is never dropped or skipped — while capping the worst
|
// blocks on the semaphore, it is never dropped or skipped. 8 (raised from
|
||||||
// case to 4 concurrent backends' cost, not the whole catalog's.
|
// an initial 4, security review round 2, 2026-08-22, important finding
|
||||||
const mcpFederatedFanoutConcurrency = 4
|
// 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
|
// mcpFederatedToolsList implements the federated "tools/list": it fans
|
||||||
// out one mcpBackendCall per server in names (handleMCPFederated's own
|
// out one mcpBackendCall per server in names (handleMCPFederated's own
|
||||||
// allowedMCPServerNames(grp) call, made once there so the SAME set that
|
// allowedMCPServerNames(grp) call, made once there so the exact set
|
||||||
// was weighed into checkAndCountWeighted is the set actually contacted),
|
// contacted is always the one the client actually asked to reach), at
|
||||||
// at most mcpFederatedFanoutConcurrency at a time — each under its own
|
// most mcpFederatedFanoutConcurrency at a time — every backend sharing
|
||||||
// toolsListBackendTimeout budget (context.WithTimeout off r.Context(),
|
// ONE overall toolsListBackendTimeout deadline (fanoutCtx, below) — then
|
||||||
// MF1) — then merges every reachable server's own tools into one result,
|
// merges every reachable server's own tools into one result,
|
||||||
// each tool's name prefixed "<serverName>_<toolName>", the exact
|
// each tool's name prefixed "<serverName>_<toolName>", the exact
|
||||||
// underscore-separator convention the old agentgateway this plugin
|
// underscore-separator convention the old agentgateway this plugin
|
||||||
// replaces used, and the one pugbot's mcpclient and agentkit have already
|
// 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))
|
merged = make([]mcpTool, 0, len(names))
|
||||||
failed []string
|
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)
|
sem := make(chan struct{}, mcpFederatedFanoutConcurrency)
|
||||||
for _, name := range names {
|
for _, name := range names {
|
||||||
wg.Add(1)
|
wg.Add(1)
|
||||||
@@ -752,10 +825,8 @@ func (g *Gateway) mcpFederatedToolsList(w http.ResponseWriter, r *http.Request,
|
|||||||
sem <- struct{}{}
|
sem <- struct{}{}
|
||||||
defer func() { <-sem }()
|
defer func() { <-sem }()
|
||||||
|
|
||||||
ctx, cancel := context.WithTimeout(r.Context(), toolsListBackendTimeout)
|
|
||||||
defer cancel()
|
|
||||||
targetURL := g.cfg.MCPServers[name].URL
|
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()
|
mu.Lock()
|
||||||
defer mu.Unlock()
|
defer mu.Unlock()
|
||||||
@@ -867,11 +938,24 @@ func (g *Gateway) mcpFederatedToolsCall(w http.ResponseWriter, r *http.Request,
|
|||||||
ctx, cancel := context.WithTimeout(r.Context(), toolsCallBackendTimeout)
|
ctx, cancel := context.WithTimeout(r.Context(), toolsCallBackendTimeout)
|
||||||
defer cancel()
|
defer cancel()
|
||||||
targetURL := g.cfg.MCPServers[serverName].URL
|
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)
|
g.limiter.countTargetRequest(targetKindMCP, serverName)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
g.logf("federated tools/call: server %q: %v", serverName, err)
|
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
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
+201
-83
@@ -1031,7 +1031,7 @@ func TestMcpBackendCall_HandshakeFallbackTrigger(t *testing.T) {
|
|||||||
t.Fatal("handler is not *Gateway")
|
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()
|
gotCalls := callCount.Load()
|
||||||
if tc.wantHandshake && gotCalls < 2 {
|
if tc.wantHandshake && gotCalls < 2 {
|
||||||
@@ -1060,12 +1060,134 @@ func TestMcpBackendCall_NetworkFailure_NeverAttemptsHandshake(t *testing.T) {
|
|||||||
t.Fatal("handler is not *Gateway")
|
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 {
|
if err == nil {
|
||||||
t.Fatal("want an error (nothing listening on 127.0.0.1:1)")
|
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
|
// TestHandleMCPFederated_ToolsList_HandshakeFallback_SucceedsAfterSessionRetry
|
||||||
// is MF3's core proof: a server that rejects a bare tools/list call
|
// is MF3's core proof: a server that rejects a bare tools/list call
|
||||||
// (readitall/fetch's real, live-probed behavior — sessionRequiredMockServer's
|
// (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
|
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
|
// TestHandleMCPFederated_ToolsList_NotWeightedByBackendCount_BudgetBuysNCalls
|
||||||
// proves a tools/list call charges checkAndCount len(allowedMCPServerNames)
|
// is the regression for security review round 2's critical finding 3:
|
||||||
// requests, not one: with 3 configured (and unrestricted) servers and a
|
// weighting checkAndCount by the number of allowed MCP servers (security
|
||||||
// requests-per-minute budget of 2, the very FIRST tools/list call must be
|
// audit finding 1b's original fix) was not default-preserving — on a
|
||||||
// rejected outright — and, because checkAndCountWeighted runs and fails
|
// config where the MCP allow-list is unrestricted (every group on the
|
||||||
// BEFORE the fan-out ever starts, none of the three servers may be
|
// live fleet), it silently turned a configured requests-per-minute budget
|
||||||
// contacted at all.
|
// into budget/serverCount successful tools/list calls, with NO config
|
||||||
func TestHandleMCPFederated_ToolsList_WeightedByBackendCount_RejectsWhenOverBudget(t *testing.T) {
|
// change on the operator's part (11 servers, requests-per-minute: 60 ->
|
||||||
alpha := newMockJSONRPCServer(t, nil)
|
// only 5 calls/minute actually succeeded; examples/kubernetes.yaml's
|
||||||
beta := newMockJSONRPCServer(t, nil)
|
// shipped requests-per-minute: 10 would 429 the very FIRST call ever
|
||||||
gamma := newMockJSONRPCServer(t, nil)
|
// made). With numServers (11, matching the live fleet) configured and a
|
||||||
cfg := newFederationTestConfig(alpha.srv.URL, beta.srv.URL, false)
|
// requests-per-minute budget of exactly rpmBudget, all rpmBudget calls
|
||||||
cfg.MCPServers["gamma"] = &TargetConfig{URL: gamma.srv.URL}
|
// must succeed and the next one must be rejected — proving the charge is
|
||||||
cfg.Groups["default"].Limits = &LimitsConfig{RequestsPerMinute: 2}
|
// 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")
|
h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw")
|
||||||
if err != nil {
|
if err != nil {
|
||||||
t.Fatalf("New: %v", err)
|
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()
|
rec := httptest.NewRecorder()
|
||||||
h.ServeHTTP(rec, newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "tools/list", ID: json.RawMessage("1")}))
|
h.ServeHTTP(rec, newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "tools/list", ID: json.RawMessage("1")}))
|
||||||
if rec.Code != http.StatusTooManyRequests {
|
if rec.Code != http.StatusOK {
|
||||||
t.Fatalf("status = %d, want 429 (weight=3 > limit=2), body=%s", rec.Code, rec.Body.String())
|
t.Fatalf("status = %d, want 200, 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()
|
now := time.Now()
|
||||||
h.ServeHTTP(rec1, newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "tools/list", ID: json.RawMessage("1")}))
|
userReq, ok := gw.limiter.getCounter("user", "alice", metricReq, windowDay, now)
|
||||||
if rec1.Code != http.StatusOK {
|
if !ok || userReq != 1 {
|
||||||
t.Fatalf("first call status = %d, want 200, body=%s", rec1.Code, rec1.Body.String())
|
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)
|
||||||
rec2 := httptest.NewRecorder()
|
if !ok || alphaReq != 1 {
|
||||||
h.ServeHTTP(rec2, newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "tools/list", ID: json.RawMessage("2")}))
|
t.Errorf("mcp/alpha req:day counter = %d (ok=%v), want 1", alphaReq, ok)
|
||||||
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())
|
|
||||||
}
|
}
|
||||||
}
|
betaReq, ok := gw.limiter.getCounter(targetKindMCP, "beta", metricReq, windowDay, now)
|
||||||
|
if !ok || betaReq != 1 {
|
||||||
// TestHandleMCPFederated_ToolsCall_WeightStaysOne_EvenWithManyServersConfigured
|
t.Errorf("mcp/beta req:day counter = %d (ok=%v), want 1", betaReq, ok)
|
||||||
// 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())
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user