From ca8062e506135081e1d7805c23780be40a6b3f6b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 18:07:21 -0700 Subject: [PATCH] fix(claude_code_gateway): keep the device secret out of the browser URL and validate the login before claiming it --- .../anthropic_endpoints/gateway_endpoints.py | 117 +++++++++++------- .../test_gateway_endpoints.py | 97 ++++++++++++--- 2 files changed, 146 insertions(+), 68 deletions(-) diff --git a/litellm/proxy/anthropic_endpoints/gateway_endpoints.py b/litellm/proxy/anthropic_endpoints/gateway_endpoints.py index 08579186f5e..259d5202db6 100644 --- a/litellm/proxy/anthropic_endpoints/gateway_endpoints.py +++ b/litellm/proxy/anthropic_endpoints/gateway_endpoints.py @@ -18,6 +18,7 @@ import hashlib import json import secrets from collections.abc import Mapping +from dataclasses import dataclass from types import MappingProxyType from typing import Final @@ -32,13 +33,16 @@ from litellm.constants import ( CLI_SSO_SESSION_TTL_SECONDS, LITELLM_CLI_SOURCE_IDENTIFIER, ) +from litellm.proxy._types import LiteLLM_UserTable, LitellmUserRoles from litellm.proxy.anthropic_endpoints.endpoints import anthropic_response, count_tokens from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.http_parsing_utils import _safe_set_request_parsed_body +from litellm.proxy.management_endpoints.ui_sso import CliSsoTeamDetail GATEWAY_PREFIX: Final = "/claude_code_gateway" _DEVICE_CODE_GRANT: Final = "urn:ietf:params:oauth:grant-type:device_code" _REFRESH_TOKEN_GRANT: Final = "refresh_token" +_DEVICE_CODE_SEPARATOR: Final = "." _DEVICE_POLL_INTERVAL_SECONDS: Final = 5 _SECONDS_PER_HOUR: Final = 3600 _MANAGED_SETTINGS_ADAPTER: Final = TypeAdapter(dict[str, object]) @@ -48,12 +52,19 @@ _POST_ONLY: Final = ["POST"] # mutable-ok: FastAPI's add_api_route only accepts class _GatewaySessionData(BaseModel): user_id: str - user_role: str | None + user_role: LitellmUserRoles models: list[str] = Field(default_factory=list) teams: tuple[str, ...] = () team_details: object | None = None +@dataclass(frozen=True, slots=True) +class _GatewayLogin: + user_info: LiteLLM_UserTable + team_id: str | None + team: CliSsoTeamDetail + + class _OAuthErrorBody(BaseModel): error: str error_description: str | None = None @@ -70,7 +81,7 @@ class _DeviceAuthorizationBody(BaseModel): device_code: str user_code: str verification_uri: str - verification_uri_complete: str + verification_uri_complete: str | None = None expires_in: int interval: int @@ -111,15 +122,11 @@ def _managed_settings() -> dict[str, object] | None: return _MANAGED_SETTINGS_ADAPTER.validate_python(settings) -def _oauth_error(*, status_code: int, error: str, description: str | None = None) -> "_OAuthError": - return _OAuthError(status_code=status_code, error=error, description=description) - - -class _OAuthError(Exception): - def __init__(self, *, status_code: int, error: str, description: str | None) -> None: - self.status_code = status_code - self.error = error - self.description = description +@dataclass(frozen=True, slots=True) +class _OAuthError: + status_code: int + error: str + description: str | None = None def _oauth_error_response(err: _OAuthError) -> JSONResponse: @@ -153,7 +160,7 @@ router.add_api_route( @router.get("/.well-known/oauth-authorization-server", include_in_schema=False) async def oauth_authorization_server(request: Request) -> JSONResponse: if not _is_gateway_enabled(): - return _oauth_error_response(_oauth_error(status_code=404, error="not_found")) + return _oauth_error_response(_OAuthError(status_code=404, error="not_found")) from litellm.proxy.utils import get_custom_url @@ -175,6 +182,7 @@ async def device_authorization(request: Request) -> JSONResponse: from litellm.proxy.management_endpoints.ui_sso import ( _check_cli_sso_start_rate_limit, # pyright: ignore[reportPrivateUsage] # shared device-flow helper + _cli_sso_verification_uri_complete_enabled, # pyright: ignore[reportPrivateUsage] # shared device-flow helper _generate_cli_sso_user_code, # pyright: ignore[reportPrivateUsage] # shared device-flow helper _hash_cli_sso_secret, # pyright: ignore[reportPrivateUsage] # shared device-flow helper _normalize_cli_sso_user_code, # pyright: ignore[reportPrivateUsage] # shared device-flow helper @@ -184,7 +192,7 @@ async def device_authorization(request: Request) -> JSONResponse: from litellm.proxy.utils import get_custom_url if not _is_gateway_enabled(): - return _oauth_error_response(_oauth_error(status_code=404, error="not_found")) + return _oauth_error_response(_OAuthError(status_code=404, error="not_found")) _check_cli_sso_start_rate_limit( request=request, @@ -192,50 +200,51 @@ async def device_authorization(request: Request) -> JSONResponse: use_x_forwarded_for=bool(_general_settings().get("use_x_forwarded_for", False)), ) - device_code: Final = f"cli-{secrets.token_urlsafe(24)}" + login_id: Final = f"cli-{secrets.token_urlsafe(24)}" + poll_secret: Final = secrets.token_urlsafe(32) user_code: Final = _generate_cli_sso_user_code() flow: Final = { # mutable-ok: the shared CLI SSO cache entry is a dict the browser leg mutates - "poll_secret_hash": _hash_cli_sso_secret(device_code), + "poll_secret_hash": _hash_cli_sso_secret(poll_secret), "user_code_hash": _hash_cli_sso_secret(_normalize_cli_sso_user_code(user_code)), "sso_complete": False, "user_code_verified": False, "session_data": None, } - _set_cli_sso_flow(login_id=device_code, cache=cli_sso_session_cache, flow=flow) + _set_cli_sso_flow(login_id=login_id, cache=cli_sso_session_cache, flow=flow) request_base_url: Final = str(request.base_url) verification_uri: Final = get_custom_url(request_base_url=request_base_url, route="sso/key/generate") - query: Final = MappingProxyType({"source": LITELLM_CLI_SOURCE_IDENTIFIER, "key": device_code}) + query: Final = MappingProxyType({"source": LITELLM_CLI_SOURCE_IDENTIFIER, "key": login_id}) body: Final = _DeviceAuthorizationBody( - device_code=device_code, + device_code=f"{login_id}{_DEVICE_CODE_SEPARATOR}{poll_secret}", user_code=user_code, verification_uri=f"{verification_uri}?{urlencode(query)}", verification_uri_complete=( f"{verification_uri}?{urlencode(MappingProxyType({**query, 'user_code': user_code}))}" + if _cli_sso_verification_uri_complete_enabled() + else None ), expires_in=CLI_SSO_SESSION_TTL_SECONDS, interval=_DEVICE_POLL_INTERVAL_SECONDS, ) - return JSONResponse(content=body.model_dump()) + return JSONResponse(content=body.model_dump(exclude_none=True)) -def _mint_access_token_from_flow(flow: Mapping[str, object]) -> str: - from litellm.proxy._types import LiteLLM_UserTable - from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken +def _validate_login(flow: Mapping[str, object]) -> _GatewayLogin | _OAuthError: from litellm.proxy.management_endpoints.ui_sso import selected_cli_sso_team_detail try: session_data: Final = _GatewaySessionData.model_validate(flow.get("session_data")) except ValidationError as err: verbose_proxy_logger.warning("Claude Code gateway login session is malformed: %s", err) - raise _oauth_error( + return _OAuthError( status_code=400, error="invalid_grant", description="The login session is malformed; sign in again" - ) from err + ) team_id: Final = session_data.teams[0] if session_data.teams else None selected_team: Final = selected_cli_sso_team_detail(team_details=session_data.team_details, team_id=team_id) if selected_team is None: - raise _oauth_error( + return _OAuthError( status_code=400, error="invalid_grant", description=f"Could not resolve the model grants for team {team_id}; sign in again", @@ -243,26 +252,32 @@ def _mint_access_token_from_flow(flow: Mapping[str, object]) -> str: user_info: Final = LiteLLM_UserTable( user_id=session_data.user_id, - user_role=session_data.user_role, + user_role=session_data.user_role.value, models=session_data.models, ) + return _GatewayLogin(user_info=user_info, team_id=team_id, team=selected_team) + + +def _mint_access_token(login: _GatewayLogin) -> str: + from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken + return ExperimentalUIJWTToken.get_cli_jwt_auth_token( - user_info=user_info, - team_id=team_id, - team_alias=selected_team.team_alias, - team_models=selected_team.team_models, - team_model_aliases=selected_team.team_model_aliases, + user_info=login.user_info, + team_id=login.team_id, + team_alias=login.team.team_alias, + team_models=login.team.team_models, + team_model_aliases=login.team.team_model_aliases, max_budget=None, ) -async def _claim_device_code(device_code: str, cache: DualCache) -> bool: +async def _claim_device_code(login_id: str, cache: DualCache) -> bool: from litellm.proxy.management_endpoints.ui_sso import ( _get_cli_sso_flow_cache_key, # pyright: ignore[reportPrivateUsage] # shared device-flow helper ) claims: Final = await cache.async_increment_cache( - key=f"{_get_cli_sso_flow_cache_key(device_code)}:claimed", + key=f"{_get_cli_sso_flow_cache_key(login_id)}:claimed", value=1, ttl=CLI_SSO_SESSION_TTL_SECONDS, ) @@ -275,39 +290,45 @@ async def _handle_device_code_grant(device_code: str | None) -> JSONResponse: from litellm.proxy.management_endpoints.ui_sso import ( _get_cli_sso_flow_cache_key, # pyright: ignore[reportPrivateUsage] # shared device-flow helper _get_cli_sso_flow_or_raise, # pyright: ignore[reportPrivateUsage] # shared device-flow helper + _verify_cli_sso_poll_secret, # pyright: ignore[reportPrivateUsage] # shared device-flow helper ) from litellm.proxy.proxy_server import cli_sso_session_cache if not device_code: return _oauth_error_response( - _oauth_error(status_code=400, error="invalid_request", description="device_code is required") + _OAuthError(status_code=400, error="invalid_request", description="device_code is required") ) + login_id, _, poll_secret = device_code.partition(_DEVICE_CODE_SEPARATOR) try: - flow: Final = _get_cli_sso_flow_or_raise(login_id=device_code, cache=cli_sso_session_cache) + flow: Final = _get_cli_sso_flow_or_raise(login_id=login_id, cache=cli_sso_session_cache) except HTTPException: - return _oauth_error_response(_oauth_error(status_code=400, error="expired_token")) + return _oauth_error_response(_OAuthError(status_code=400, error="expired_token")) + + if not _verify_cli_sso_poll_secret(flow, poll_secret): + return _oauth_error_response(_OAuthError(status_code=400, error="expired_token")) if not flow.get("sso_complete") or not flow.get("user_code_verified"): - return _oauth_error_response(_oauth_error(status_code=400, error="authorization_pending")) + return _oauth_error_response(_OAuthError(status_code=400, error="authorization_pending")) - if not await _claim_device_code(device_code, cli_sso_session_cache): - return _oauth_error_response(_oauth_error(status_code=400, error="expired_token")) + login: Final = _validate_login(flow) + if isinstance(login, _OAuthError): + return _oauth_error_response(login) - await cli_sso_session_cache.async_delete_cache(key=_get_cli_sso_flow_cache_key(device_code)) - try: - access_token: Final = _mint_access_token_from_flow(flow) - except _OAuthError as err: - return _oauth_error_response(err) + if not await _claim_device_code(login_id, cli_sso_session_cache): + return _oauth_error_response(_OAuthError(status_code=400, error="expired_token")) - body: Final = _AccessTokenBody(access_token=access_token, expires_in=CLI_JWT_EXPIRATION_HOURS * _SECONDS_PER_HOUR) + await cli_sso_session_cache.async_delete_cache(key=_get_cli_sso_flow_cache_key(login_id)) + body: Final = _AccessTokenBody( + access_token=_mint_access_token(login), expires_in=CLI_JWT_EXPIRATION_HOURS * _SECONDS_PER_HOUR + ) return JSONResponse(content=body.model_dump()) @router.post("/oauth/token", include_in_schema=False) async def oauth_token(request: Request) -> JSONResponse: if not _is_gateway_enabled(): - return _oauth_error_response(_oauth_error(status_code=404, error="not_found")) + return _oauth_error_response(_OAuthError(status_code=404, error="not_found")) form: Final = await request.form() grant_type: Final = form.get("grant_type") @@ -318,7 +339,7 @@ async def oauth_token(request: Request) -> JSONResponse: if grant_type == _REFRESH_TOKEN_GRANT: return _oauth_error_response( - _oauth_error( + _OAuthError( status_code=401, error="invalid_grant", description="This gateway does not issue refresh tokens; sign in again", @@ -326,7 +347,7 @@ async def oauth_token(request: Request) -> JSONResponse: ) return _oauth_error_response( - _oauth_error( + _OAuthError( status_code=400, error="unsupported_grant_type", description=f"Unsupported grant_type: {grant_type}" ) ) diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_gateway_endpoints.py b/tests/test_litellm/proxy/anthropic_endpoints/test_gateway_endpoints.py index e49047634bc..d442ac21307 100644 --- a/tests/test_litellm/proxy/anthropic_endpoints/test_gateway_endpoints.py +++ b/tests/test_litellm/proxy/anthropic_endpoints/test_gateway_endpoints.py @@ -20,11 +20,18 @@ from fastapi.testclient import TestClient from litellm.caching.dual_cache import DualCache from litellm.proxy._types import ProxyException from litellm.proxy.anthropic_endpoints import gateway_endpoints -from litellm.proxy.management_endpoints.ui_sso import _get_cli_sso_flow_cache_key, _set_cli_sso_flow +from litellm.proxy.management_endpoints.ui_sso import ( + _get_cli_sso_flow_cache_key, + _hash_cli_sso_secret, + _set_cli_sso_flow, +) from litellm.proxy.middleware.prometheus_auth_middleware import PrometheusAuthMiddleware _DEVICE_CODE_GRANT: Final = "urn:ietf:params:oauth:grant-type:device_code" _MASTER_KEY: Final = "sk-master-key" +_SHARED_LOGIN_ID: Final = "cli-shared-login-code" +_SHARED_POLL_SECRET: Final = "shared-poll-secret" +_SHARED_DEVICE_CODE: Final = f"{_SHARED_LOGIN_ID}.{_SHARED_POLL_SECRET}" _MINT: Final = "litellm.proxy.auth.auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token" _PROTOBUF_BODY: Final = b"\x0a\x05hello\x12\x03{{{" _COMPLETED_SESSION: Final = MappingProxyType( @@ -100,10 +107,12 @@ def _gateway_env( managed_settings: Mapping[str, object] | None = None, cache: DualCache | None = None, real_auth: bool = False, + extra_settings: Mapping[str, object] = MappingProxyType({}), ) -> Iterator[tuple[TestClient, DualCache]]: general_settings: Final = { "enable_claude_code_gateway": enabled, **({} if managed_settings is None else {"claude_code_gateway_managed_settings": dict(managed_settings)}), + **extra_settings, } session_cache: Final = cache or DualCache(default_in_memory_ttl=600) @@ -147,7 +156,7 @@ def _request_token(client: TestClient, device_code: str) -> httpx.Response: def _completed_flow(session_data: Mapping[str, object] = _COMPLETED_SESSION) -> dict[str, object]: return { - "poll_secret_hash": "unused", + "poll_secret_hash": _hash_cli_sso_secret(_SHARED_POLL_SECRET), "user_code_hash": "unused", "sso_complete": True, "user_code_verified": True, @@ -155,13 +164,18 @@ def _completed_flow(session_data: Mapping[str, object] = _COMPLETED_SESSION) -> } +def _login_id(device_code: str) -> str: + return device_code.partition(".")[0] + + def _complete_flow( cache: DualCache, device_code: str, session_data: Mapping[str, object] = _COMPLETED_SESSION ) -> None: - key: Final = _get_cli_sso_flow_cache_key(device_code) + key: Final = _get_cli_sso_flow_cache_key(_login_id(device_code)) flow: Final = cache.get_cache(key=key) assert isinstance(flow, dict) - cache.set_cache(key=key, value={**flow, **_completed_flow(session_data)}, ttl=600) + completed: Final = {**flow, **_completed_flow(session_data), "poll_secret_hash": flow["poll_secret_hash"]} + cache.set_cache(key=key, value=completed, ttl=600) def test_discovery_shape(): @@ -194,18 +208,35 @@ def test_device_authorization_returns_rfc8628_shape_and_persists_flow(): assert resp.status_code == 200 body = resp.json() device_code = body["device_code"] - assert device_code.startswith("cli-") + login_id, separator, poll_secret = device_code.partition(".") + assert login_id.startswith("cli-") + assert separator == "." + assert len(poll_secret) >= 32 assert body["user_code"] assert body["expires_in"] == 600 assert body["interval"] == 5 - # verification_uri_complete carries the user_code; the short uri does not. - assert f"user_code={body['user_code']}" in body["verification_uri_complete"] - assert "user_code=" not in body["verification_uri"] - assert f"key={device_code}" in body["verification_uri"] - # The device flow is stored under the device_code so the browser SSO leg can complete it. - stored = cache.get_cache(key=_get_cli_sso_flow_cache_key(device_code)) + assert "verification_uri_complete" not in body + assert body["verification_uri"].endswith(f"/sso/key/generate?source=litellm-cli&key={login_id}") + assert poll_secret not in body["verification_uri"] + stored = cache.get_cache(key=_get_cli_sso_flow_cache_key(login_id)) assert isinstance(stored, dict) assert stored["sso_complete"] is False + assert stored["poll_secret_hash"] == _hash_cli_sso_secret(poll_secret) + assert cache.get_cache(key=_get_cli_sso_flow_cache_key(device_code)) is None + + +@pytest.mark.parametrize("opted_in", [True, False]) +def test_verification_uri_complete_carries_the_user_code_only_when_the_operator_opts_in(opted_in: bool): + with _gateway_env(extra_settings={"allow_cli_sso_verification_uri_complete": opted_in}) as (client, _): + body = client.post("/claude_code_gateway/oauth/device_authorization").json() + login_id = _login_id(body["device_code"]) + if not opted_in: + assert "verification_uri_complete" not in body + return + assert body["verification_uri_complete"].endswith( + f"/sso/key/generate?source=litellm-cli&key={login_id}&user_code={body['user_code']}" + ) + assert "user_code=" not in body["verification_uri"] def test_token_authorization_pending_before_browser_completes(): @@ -215,6 +246,22 @@ def test_token_authorization_pending_before_browser_completes(): assert resp.json()["error"] == "authorization_pending" +@pytest.mark.parametrize("tamper", ["login_id_only", "wrong_secret"]) +def test_token_refuses_the_browser_login_id_without_the_client_secret(tamper: str): + with _gateway_env() as (client, cache): + device_code = _start_device_flow(client) + _complete_flow(cache, device_code) + login_id = _login_id(device_code) + presented = login_id if tamper == "login_id_only" else f"{login_id}.not-the-secret" + with patch(_MINT, return_value="sk-session") as mint: + resp = _request_token(client, presented) + assert resp.status_code == 400 + assert resp.json()["error"] == "expired_token" + mint.assert_not_called() + with_secret = _request_token(client, device_code) + assert with_secret.status_code == 200 + + def test_token_success_mints_bearer_and_is_single_use(): with _gateway_env() as (client, cache): device_code = _start_device_flow(client) @@ -251,14 +298,25 @@ def test_token_teamless_user_mints_without_a_team(): assert mint.call_args.kwargs["team_models"] == () -def test_token_malformed_session_is_invalid_grant(): +@pytest.mark.parametrize( + "session_data", + [ + {"user_role": "internal_user"}, + {**_COMPLETED_SESSION, "user_role": None}, + {**_COMPLETED_SESSION, "user_role": "not-a-role"}, + ], + ids=["missing_user_id", "no_role", "unknown_role"], +) +def test_token_malformed_session_is_invalid_grant_and_does_not_consume_the_login(session_data: Mapping[str, object]): with _gateway_env() as (client, cache): device_code = _start_device_flow(client) - _complete_flow(cache, device_code, session_data={"user_role": "internal_user"}) + _complete_flow(cache, device_code, session_data=session_data) with patch(_MINT) as mint: resp = _request_token(client, device_code) + again = _request_token(client, device_code) assert resp.status_code == 400 assert resp.json()["error"] == "invalid_grant" + assert again.json()["error"] == "invalid_grant" mint.assert_not_called() @@ -275,25 +333,24 @@ def test_token_unknown_team_grants_is_invalid_grant(): def test_token_mints_on_a_replica_that_did_not_start_the_login(): redis: Final = _SharedRedisFake() - device_code: Final = "cli-shared-login-code" - _set_cli_sso_flow(login_id=device_code, cache=_replica(redis), flow=_completed_flow()) + _set_cli_sso_flow(login_id=_SHARED_LOGIN_ID, cache=_replica(redis), flow=_completed_flow()) with _gateway_env(cache=_replica(redis)) as (client, _), patch(_MINT, return_value="sk-session") as mint: - resp = _request_token(client, device_code) + resp = _request_token(client, _SHARED_DEVICE_CODE) assert resp.status_code == 200 assert resp.json()["access_token"] == "sk-session" assert mint.call_args.kwargs["team_id"] == "team-a" + assert mint.call_args.kwargs["user_info"].user_role == "internal_user" def test_token_refuses_a_device_code_another_replica_already_claimed(): redis: Final = _SharedRedisFake() replica_a: Final = _replica(redis) - device_code: Final = "cli-shared-login-code" - _set_cli_sso_flow(login_id=device_code, cache=replica_a, flow=_completed_flow()) - assert asyncio.run(gateway_endpoints._claim_device_code(device_code, replica_a)) is True + _set_cli_sso_flow(login_id=_SHARED_LOGIN_ID, cache=replica_a, flow=_completed_flow()) + assert asyncio.run(gateway_endpoints._claim_device_code(_SHARED_LOGIN_ID, replica_a)) is True with _gateway_env(cache=_replica(redis)) as (client, _), patch(_MINT) as mint: - resp = _request_token(client, device_code) + resp = _request_token(client, _SHARED_DEVICE_CODE) assert resp.status_code == 400 assert resp.json()["error"] == "expired_token" mint.assert_not_called()