mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-20 05:56:56 +00:00
fix(mcp): bound federation tools/list fan-out and weight admission by backend count
Fixes finding 1 from the 2026-08-22 security audit: mcpFederatedToolsList spawned one unbounded goroutine per allowed MCP server, each able to issue up to 4 sequential HTTP calls reading up to 10MiB per response, while checkAndCount charged the caller only 1 request regardless of how many backends the fan-out actually contacted. With the 11 production servers, one client tools/list call could drive up to 44 outbound requests and hold ~110MiB concurrently, all for the cost of a single request against the caller's own rate limit. - Bound the fan-out to mcpFederatedFanoutConcurrency (4) concurrent backends via a buffered-channel semaphore, the same shape limiter.spawnTokens already uses. - Add limiter.checkAndCountWeighted so a tools/list call charges len(allowedMCPServerNames) requests against the caller's requests-per-minute/day budget instead of undercounting the fan-out as one; every other route keeps weight 1 via checkAndCount's own unchanged wrapper. - Shrink the per-backend response cap from maxRequestBytes (10MiB, sized for a client's chat/embeddings body) to mcpBackendResponseMaxBytes (4MiB, sized for a JSON-RPC reply), bounding worst-case concurrent memory to 4 x 4MiB = 16MiB.
This commit is contained in:
@@ -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
|
// 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, 1, minWindowTTL),
|
newCounterIncr(sc.kind, sc.id, metricReq, windowMin, now, weight, minWindowTTL),
|
||||||
newCounterIncr(sc.kind, sc.id, metricReq, windowDay, now, 1, dayWindowTTL),
|
newCounterIncr(sc.kind, sc.id, metricReq, windowDay, now, weight, dayWindowTTL),
|
||||||
newCounterIncr(sc.kind, sc.id, metricReq, windowHour, now, 1, hourWindowTTL),
|
newCounterIncr(sc.kind, sc.id, metricReq, windowHour, now, weight, hourWindowTTL),
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
probes, reads := buildBudgetProbes(scopes, now)
|
probes, reads := buildBudgetProbes(scopes, now)
|
||||||
|
|||||||
+89
-25
@@ -66,6 +66,21 @@ const (
|
|||||||
mcpSessionCloseTimeout = 5 * time.Second
|
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
|
// 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).
|
||||||
@@ -237,6 +252,24 @@ func allowedMCPServerNames(cfg *Config, grp *group) []string {
|
|||||||
// other metered route — with per-target attribution added separately
|
// other metered route — 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 (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) {
|
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")
|
||||||
@@ -259,7 +292,17 @@ 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
|
||||||
}
|
}
|
||||||
@@ -270,7 +313,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, grp)
|
g.mcpFederatedToolsList(w, r, req, toolsListNames)
|
||||||
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/"):
|
||||||
@@ -433,7 +476,7 @@ 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, maxRequestBytes))
|
respBody, err := io.ReadAll(io.LimitReader(httpResp.Body, mcpBackendResponseMaxBytes))
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, status, respSessionID, err
|
return nil, status, respSessionID, err
|
||||||
}
|
}
|
||||||
@@ -632,43 +675,64 @@ func lastSSEDataLine(body []byte) []byte {
|
|||||||
return last
|
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
|
// mcpFederatedToolsList implements the federated "tools/list": it fans
|
||||||
// out one mcpBackendCall per server allowedMCPServerNames(grp) permits,
|
// out one mcpBackendCall per server in names (handleMCPFederated's own
|
||||||
// concurrently — each under its own toolsListBackendTimeout budget
|
// allowedMCPServerNames(grp) call, made once there so the SAME set that
|
||||||
// (context.WithTimeout off r.Context(), MF1) — then merges every
|
// was weighed into checkAndCountWeighted is the set actually contacted),
|
||||||
// reachable server's own tools into one result, each tool's name prefixed
|
// at most mcpFederatedFanoutConcurrency at a time — each under its own
|
||||||
// "<serverName>_<toolName>", the exact underscore-separator convention the
|
// toolsListBackendTimeout budget (context.WithTimeout off r.Context(),
|
||||||
// old agentgateway this plugin replaces used, and the one pugbot's
|
// MF1) — then merges every reachable server's own tools into one result,
|
||||||
// mcpclient and agentkit have already persisted into their own tool-id
|
// each tool's name prefixed "<serverName>_<toolName>", the exact
|
||||||
// databases (e.g. "brave-search_brave_web_search" — verified live). A
|
// underscore-separator convention the old agentgateway this plugin
|
||||||
// server that errors, times out, or returns an unparsable/invalid result
|
// replaces used, and the one pugbot's mcpclient and agentkit have already
|
||||||
// is logged and skipped, not surfaced as a whole-call failure — UNLESS
|
// persisted into their own tool-id databases (e.g.
|
||||||
// every single attempted server failed, in which case an empty tools list
|
// "brave-search_brave_web_search" — verified live). A server that errors,
|
||||||
// would misleadingly look like "this caller's group has no MCP access" —
|
// times out, or returns an unparsable/invalid result is logged
|
||||||
// see the loud-failure branch below.
|
// 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
|
// countTargetRequests (limits.go) attributes one request to EVERY server
|
||||||
// ATTEMPTED here, in ONE batched call after wg.Wait() — "attempted", not
|
// ATTEMPTED here, in ONE batched call after wg.Wait() — "attempted", not
|
||||||
// "reached" or "succeeded": a server this gateway dialed and got a
|
// "reached" or "succeeded": a server this gateway dialed and got a
|
||||||
// response (or a timeout, or a connection refusal) from still had a real
|
// response (or a timeout, or a connection refusal) from still had
|
||||||
// request sent to it and a real slot of this gateway's outbound capacity
|
// a real request sent to it and a real slot of this gateway's outbound
|
||||||
// spent on it, which is what these counters exist to track. A single
|
// capacity spent on it, which is what these counters exist to track. A
|
||||||
// federated tools/list call can move several targets' own counters, not
|
// single federated tools/list call can move several targets' own
|
||||||
// just one — unlike handleTargetProxy's always-exactly-one-target shape
|
// counters, not just one — unlike handleTargetProxy's always-exactly-
|
||||||
// (mcp_a2a.go).
|
// one-target shape (mcp_a2a.go).
|
||||||
func (g *Gateway) mcpFederatedToolsList(w http.ResponseWriter, r *http.Request, req jsonrpcRequest, grp *group) {
|
func (g *Gateway) mcpFederatedToolsList(w http.ResponseWriter, r *http.Request, req jsonrpcRequest, names []string) {
|
||||||
names := allowedMCPServerNames(g.cfg, grp)
|
|
||||||
|
|
||||||
var (
|
var (
|
||||||
mu sync.Mutex
|
mu sync.Mutex
|
||||||
wg sync.WaitGroup
|
wg sync.WaitGroup
|
||||||
merged = make([]mcpTool, 0, len(names))
|
merged = make([]mcpTool, 0, len(names))
|
||||||
failed []string
|
failed []string
|
||||||
)
|
)
|
||||||
|
sem := make(chan struct{}, mcpFederatedFanoutConcurrency)
|
||||||
for _, name := range names {
|
for _, name := range names {
|
||||||
wg.Add(1)
|
wg.Add(1)
|
||||||
go func(name string) {
|
go func(name string) {
|
||||||
defer wg.Done()
|
defer wg.Done()
|
||||||
|
sem <- struct{}{}
|
||||||
|
defer func() { <-sem }()
|
||||||
|
|
||||||
ctx, cancel := context.WithTimeout(r.Context(), toolsListBackendTimeout)
|
ctx, cancel := context.WithTimeout(r.Context(), toolsListBackendTimeout)
|
||||||
defer cancel()
|
defer cancel()
|
||||||
targetURL := g.cfg.MCPServers[name].URL
|
targetURL := g.cfg.MCPServers[name].URL
|
||||||
|
|||||||
@@ -4,6 +4,7 @@ import (
|
|||||||
"bytes"
|
"bytes"
|
||||||
"context"
|
"context"
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
"io"
|
"io"
|
||||||
"net/http"
|
"net/http"
|
||||||
"net/http/httptest"
|
"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 ---
|
// --- MF3: legacy HTTP+SSE transport exclusion ---
|
||||||
|
|
||||||
func TestIsLegacySSETransportURL(t *testing.T) {
|
func TestIsLegacySSETransportURL(t *testing.T) {
|
||||||
|
|||||||
Reference in New Issue
Block a user