From 36097f2c21742c698d736a2bfe60a7fe567cb26e Mon Sep 17 00:00:00 2001 From: Artur Sidorenko Date: Mon, 2 Nov 2026 13:03:51 +0100 Subject: [PATCH] fix(mcp): keep the BYOK pending-code cap with the shared store and bound the submitted key size With redis_usage_cache the pending codes no longer live in the in-memory dict, so the 1000-code cap checked against that dict no longer applied. Count issued codes per TTL window in the shared store with an atomic increment (fail closed when it cannot count) and reject API keys longer than 8192 characters before a code is issued. --- .../proxy/mcp_server/byok_oauth_endpoints.py | 24 +++++++++- .../mcp_server/test_byok_oauth_endpoints.py | 45 ++++++++++++++++++- 2 files changed, 67 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/mcp_server/byok_oauth_endpoints.py b/litellm/proxy/mcp_server/byok_oauth_endpoints.py index 9b4ecaf5fac..07bdb81a0f2 100644 --- a/litellm/proxy/mcp_server/byok_oauth_endpoints.py +++ b/litellm/proxy/mcp_server/byok_oauth_endpoints.py @@ -50,11 +50,16 @@ from litellm.proxy.middleware.per_request_root_path_middleware import get_server _byok_auth_codes: Final[dict[str, dict]] = {} _AUTH_CODE_CACHE_PREFIX: Final = "mcp_byok_auth_code:" _AUTH_CODE_CLAIM_PREFIX: Final = "mcp_byok_auth_code_claim:" +# Codes issued in the current TTL window, counted in the shared store; bounds pending codes +# there the way len(_byok_auth_codes) bounds them in memory. +_AUTH_CODE_WINDOW_KEY: Final = "mcp_byok_auth_code_window" # Authorization codes expire after 5 minutes. _AUTH_CODE_TTL_SECONDS: Final = 300 # Hard cap to prevent memory exhaustion from incomplete OAuth flows. _AUTH_CODES_MAX_SIZE: Final = 1000 +# Upper bound for a submitted API key, so a pending code cannot carry an arbitrarily large value. +_MAX_API_KEY_LENGTH: Final = 8192 router: Final = APIRouter(tags=["mcp"]) @@ -133,6 +138,21 @@ async def _discard_auth_code(code: str) -> None: await shared.async_delete_cache(_AUTH_CODE_CACHE_PREFIX + code) +async def _reserve_auth_code_slot() -> bool: + """Whether another pending code may be issued. In memory this is the dict size; with the + shared store it is an atomic per-window counter, since the dict stays empty there. A shared + store that cannot count refuses (fail closed).""" + shared: Final = _shared_code_store() + if shared is None: + return len(_byok_auth_codes) < _AUTH_CODES_MAX_SIZE + try: + issued: Final = await shared.async_increment(_AUTH_CODE_WINDOW_KEY, 1, ttl=_AUTH_CODE_TTL_SECONDS) + except Exception as e: # noqa: BLE001 # any shared-store fault refuses the new code + verbose_proxy_logger.warning("byok authorize: pending-code counter unavailable: %s", e) + return False + return issued <= _AUTH_CODES_MAX_SIZE + + async def _claim_auth_code(code: str) -> bool: """Atomic one-time use across workers (RFC 6749 section 4.1.2): of two concurrent redemptions of the same code only the first claim succeeds, and a claim that cannot be @@ -904,7 +924,9 @@ async def byok_authorize_post( # Reject new codes if the store is at capacity (prevents memory exhaustion # from a burst of abandoned OAuth flows). - if len(_byok_auth_codes) >= _AUTH_CODES_MAX_SIZE: + if len(api_key) > _MAX_API_KEY_LENGTH: + raise HTTPException(status_code=400, detail="api_key is too long") + if not await _reserve_auth_code_slot(): raise HTTPException(status_code=503, detail="Too many pending authorization flows") if code_challenge_method != "S256": diff --git a/tests/unit/proxy/mcp_server/test_byok_oauth_endpoints.py b/tests/unit/proxy/mcp_server/test_byok_oauth_endpoints.py index eed5d9fbfa1..f1073c5ed6b 100644 --- a/tests/unit/proxy/mcp_server/test_byok_oauth_endpoints.py +++ b/tests/unit/proxy/mcp_server/test_byok_oauth_endpoints.py @@ -2019,7 +2019,7 @@ class _SharedStore: self.values.pop(key, None) async def async_increment(self, key: str, value: int, **_: Any) -> int: - if self.fail_increment: + if self.fail_increment and key.startswith("mcp_byok_auth_code_claim:"): raise ConnectionError("redis down") self.counts[key] = self.counts.get(key, 0) + value return self.counts[key] @@ -2114,3 +2114,46 @@ async def test_a_shared_code_is_not_redeemed_when_the_single_use_claim_cannot_be result = await _redeem(code, verifier, store_credential) assert result.status_code == 400 store_credential.assert_not_awaited() + + +def test_the_pending_code_cap_holds_with_the_shared_store(client, monkeypatch): + from litellm.proxy.mcp_server.byok_oauth_endpoints import _AUTH_CODE_WINDOW_KEY, _AUTH_CODES_MAX_SIZE + + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-test-salt-key") + shared = _SharedStore() + shared.counts[_AUTH_CODE_WINDOW_KEY] = _AUTH_CODES_MAX_SIZE + with patch("litellm.proxy.proxy_server.redis_usage_cache", shared): + resp = client.post( + "/v1/mcp/oauth/authorize", + data={ + "client_id": "user-123", + "redirect_uri": "http://127.0.0.1:3000/callback", + "code_challenge": _make_challenge("a_code_verifier_that_is_long_enough_for_pkce_43"), + "code_challenge_method": "S256", + "state": "st", + "server_id": "server-xyz", + "api_key": "sk-supersecretkey", + }, + follow_redirects=False, + ) + assert resp.status_code == 503 + assert shared.values == {} + + +def test_an_oversized_api_key_is_rejected_before_a_code_is_issued(client): + from litellm.proxy.mcp_server.byok_oauth_endpoints import _MAX_API_KEY_LENGTH + + resp = client.post( + "/v1/mcp/oauth/authorize", + data={ + "client_id": "user-123", + "redirect_uri": "http://127.0.0.1:3000/callback", + "code_challenge": _make_challenge("a_code_verifier_that_is_long_enough_for_pkce_43"), + "code_challenge_method": "S256", + "state": "st", + "server_id": "server-xyz", + "api_key": "k" * (_MAX_API_KEY_LENGTH + 1), + }, + follow_redirects=False, + ) + assert resp.status_code == 400