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