mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-18 05:47:38 +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
|
||||
// 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)
|
||||
|
||||
Reference in New Issue
Block a user