diff --git a/litellm/proxy/auth/ui_session_jwt.py b/litellm/proxy/auth/ui_session_jwt.py new file mode 100644 index 00000000000..71c57727e37 --- /dev/null +++ b/litellm/proxy/auth/ui_session_jwt.py @@ -0,0 +1,52 @@ +"""The one place a UI session JWT is signed. + +The ``token`` cookie the dashboard carries is an HS256 JWT signed with the proxy ``master_key``. +It is minted from six places (SSO callback, ``/login``, ``/v2/login``, ``/v3/login``, and the two +onboarding links), and every one of them used to call :func:`jwt.encode` directly with no ``exp``, +so the cookie never expired: a leaked one stayed valid until the master key rotated. Worse, readers +that (correctly) require a bounded lifetime rejected every real cookie, because none of the mints +stamped the claim they were checking for. + +Routing all six through :func:`encode_ui_session_jwt` makes the lifetime a property of the +credential rather than of whichever endpoint happened to issue it. The window is +``LITELLM_UI_SESSION_DURATION`` (default 24h), the same setting that already bounds the virtual key +carried inside the payload, so the cookie and the key it wraps expire together instead of the +cookie outliving its own contents. + +This module deliberately imports nothing from the proxy package: ``login_utils`` already imports +from ``ui_sso``, so a shared helper living in either one would close an import cycle for the other. +""" + +from __future__ import annotations + +from datetime import datetime, timedelta, timezone + +import jwt + +from litellm.constants import LITELLM_UI_SESSION_DURATION +from litellm.litellm_core_utils.duration_parser import duration_in_seconds +from litellm.types.proxy.ui_sso import ReturnedUITokenObject, UISessionJWTClaims + +_UI_SESSION_JWT_ALGORITHM = "HS256" + + +def ui_session_expires_at(now: datetime) -> datetime: + """When a UI session minted at ``now`` expires.""" + return now + timedelta(seconds=duration_in_seconds(LITELLM_UI_SESSION_DURATION)) + + +def encode_ui_session_jwt( + token_object: ReturnedUITokenObject, + master_key: str, + now: datetime | None = None, +) -> str: + """Sign ``token_object`` as the UI session cookie, stamping a bounded ``exp``. + + ``now`` is injectable so tests can pin the expiry without patching the clock. + """ + issued_at = now if now is not None else datetime.now(timezone.utc) + claims: UISessionJWTClaims = { + **token_object, + "exp": int(ui_session_expires_at(issued_at).timestamp()), + } + return jwt.encode(dict(claims), master_key, algorithm=_UI_SESSION_JWT_ALGORITHM) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index de988a0140f..e5f4e92b4ce 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -88,6 +88,7 @@ from litellm.proxy.auth.auth_utils import ( _has_user_setup_sso, ) from litellm.proxy.auth.handle_jwt import JWTHandler +from litellm.proxy.auth.ui_session_jwt import encode_ui_session_jwt from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.admin_ui_utils import ( admin_ui_disabled, @@ -3055,8 +3056,6 @@ class SSOAuthenticationHandler: return_to: Optional[str] = None, sso_assertion: SSOIdentityAssertion | None = None, ) -> RedirectResponse: - import jwt - from litellm.proxy.proxy_server import ( general_settings, generate_key_helper_fn, @@ -3219,11 +3218,7 @@ class SSOAuthenticationHandler: server_root_path=get_server_root_path(), ) - jwt_token = jwt.encode( - cast(dict, returned_ui_token_object), - master_key or "", - algorithm="HS256", - ) + jwt_token = encode_ui_session_jwt(returned_ui_token_object, master_key or "") # Control-plane cross-origin: store JWT behind a single-use opaque # code (60s TTL) so the token never appears in browser history / logs. diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 32845763f22..d432aab23e8 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -281,6 +281,7 @@ from litellm.proxy.auth.model_checks import ( get_mcp_server_ids, get_team_models, ) +from litellm.proxy.auth.ui_session_jwt import encode_ui_session_jwt from litellm.proxy.auth.user_api_key_auth import ( _fetch_global_spend_with_event_coordination, user_api_key_auth, @@ -13497,11 +13498,7 @@ async def login(request: Request): # Generate JWT token import jwt - jwt_token = jwt.encode( - cast(dict, returned_ui_token_object), - cast(str, master_key), - algorithm="HS256", - ) + jwt_token = encode_ui_session_jwt(returned_ui_token_object, cast(str, master_key)) # Build redirect URL litellm_dashboard_ui = get_custom_url(str(request.base_url)) @@ -13543,11 +13540,7 @@ async def login_v2(request: Request): import jwt - jwt_token = jwt.encode( - cast(dict, returned_ui_token_object), - cast(str, master_key), - algorithm="HS256", - ) + jwt_token = encode_ui_session_jwt(returned_ui_token_object, cast(str, master_key)) litellm_dashboard_ui = get_custom_url(str(request.base_url)) if litellm_dashboard_ui.endswith("/"): @@ -13622,11 +13615,7 @@ async def login_v3(request: Request): import jwt - jwt_token = jwt.encode( - cast(dict, returned_ui_token_object), - cast(str, master_key), - algorithm="HS256", - ) + jwt_token = encode_ui_session_jwt(returned_ui_token_object, cast(str, master_key)) litellm_dashboard_ui = get_custom_url(str(request.base_url)) if litellm_dashboard_ui.endswith("/"): @@ -13813,11 +13802,7 @@ async def onboarding(invite_link: str, request: Request): disabled_non_admin_personal_key_creation=disabled_non_admin_personal_key_creation, server_root_path=get_server_root_path(), ) - jwt_token = jwt.encode( # type: ignore - cast(dict, returned_ui_token_object), - master_key, - algorithm="HS256", - ) + jwt_token = encode_ui_session_jwt(returned_ui_token_object, master_key) litellm_dashboard_ui += "?token={}&user_email={}".format(jwt_token, user_email) return { @@ -13923,11 +13908,7 @@ async def _generate_onboarding_ui_session_token(user_obj: Any) -> str: server_root_path=get_server_root_path(), ) assert master_key is not None - return jwt.encode( # type: ignore - cast(dict, returned_ui_token_object), - master_key, - algorithm="HS256", - ) + return encode_ui_session_jwt(returned_ui_token_object, master_key) @app.post("/onboarding/claim_token", include_in_schema=False) diff --git a/litellm/types/proxy/ui_sso.py b/litellm/types/proxy/ui_sso.py index bef200952ba..00f00b7af29 100644 --- a/litellm/types/proxy/ui_sso.py +++ b/litellm/types/proxy/ui_sso.py @@ -19,6 +19,17 @@ class ReturnedUITokenObject(TypedDict): server_root_path: str # e.g. `/litellm` +class UISessionJWTClaims(ReturnedUITokenObject): + """The signed claim set of the UI session cookie. + + ``exp`` is stamped by the encoder rather than carried on + :class:`ReturnedUITokenObject`, because the lifetime belongs to the credential and not to the + payload that endpoints hand back in a response body. + """ + + exp: int + + class ParsedOpenIDResult(TypedDict, total=False): """ Parsed OpenID result diff --git a/tests/test_litellm/proxy/auth/test_ui_session_jwt.py b/tests/test_litellm/proxy/auth/test_ui_session_jwt.py new file mode 100644 index 00000000000..928f6452aeb --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_ui_session_jwt.py @@ -0,0 +1,138 @@ +"""Tests for the single choke point that signs the UI session cookie.""" + +from datetime import datetime, timedelta, timezone +from pathlib import Path + +import jwt +import pytest + +from litellm.constants import LITELLM_UI_SESSION_DURATION +from litellm.litellm_core_utils.duration_parser import duration_in_seconds +from litellm.proxy.auth.ui_session_jwt import encode_ui_session_jwt, ui_session_expires_at +from litellm.types.proxy.ui_sso import ReturnedUITokenObject + +MASTER_KEY = "sk-master-for-tests" + + +def _token_object(**overrides) -> ReturnedUITokenObject: + base = ReturnedUITokenObject( + user_id="user-1", + key="sk-inner-key", + user_email="user@example.com", + user_role="proxy_admin", + login_method="sso", + premium_user=False, + auth_header_name="Authorization", + disabled_non_admin_personal_key_creation=False, + server_root_path="/", + ) + base.update(overrides) # type: ignore[typeddict-item] # test helper takes arbitrary claim overrides + return base + + +def test_encode_stamps_a_bounded_expiry(): + """The cookie expires LITELLM_UI_SESSION_DURATION after it is minted.""" + now = datetime(2026, 1, 1, tzinfo=timezone.utc) + + decoded = jwt.decode( + encode_ui_session_jwt(_token_object(), MASTER_KEY, now=now), + MASTER_KEY, + algorithms=["HS256"], + # The claim's VALUE is what is under test, so expiry enforcement is off; a fixed mint + # instant keeps the assertion independent of when the suite runs. + options={"require": ["exp"], "verify_exp": False}, + ) + + assert decoded["exp"] == int((now + timedelta(seconds=duration_in_seconds(LITELLM_UI_SESSION_DURATION))).timestamp()) + + +def test_encode_preserves_every_payload_claim(): + """Stamping exp must not drop or rewrite anything the dashboard reads out of the cookie.""" + token_object = _token_object() + + decoded = jwt.decode( + encode_ui_session_jwt(token_object, MASTER_KEY), + MASTER_KEY, + algorithms=["HS256"], + options={"require": ["exp"]}, + ) + + assert {claim: decoded[claim] for claim in token_object} == dict(token_object) + + +def test_a_cookie_past_its_expiry_is_rejected(): + long_ago = datetime.now(timezone.utc) - timedelta(seconds=duration_in_seconds(LITELLM_UI_SESSION_DURATION) + 60) + expired = encode_ui_session_jwt(_token_object(), MASTER_KEY, now=long_ago) + + with pytest.raises(jwt.ExpiredSignatureError): + jwt.decode(expired, MASTER_KEY, algorithms=["HS256"]) + + +def test_a_cookie_signed_with_another_key_is_rejected(): + foreign = encode_ui_session_jwt(_token_object(), "sk-some-other-master-key") + + with pytest.raises(jwt.InvalidSignatureError): + jwt.decode(foreign, MASTER_KEY, algorithms=["HS256"]) + + +def test_ui_session_expires_at_is_the_configured_window(): + now = datetime(2026, 6, 1, 12, 0, tzinfo=timezone.utc) + + assert ui_session_expires_at(now) - now == timedelta(seconds=duration_in_seconds(LITELLM_UI_SESSION_DURATION)) + + +def test_minted_cookie_is_accepted_by_the_session_cookie_reader(monkeypatch): + """The regression this module exists for. + + ``_user_id_from_session_cookie`` requires an ``exp`` claim, so before the mints stamped one it + rejected every genuine login cookie and the MCP SSO interpose was an unconditional + redirect-to-login loop. A cookie in the shape the mints previously produced must still be + rejected, and one from the choke point must resolve the user. + """ + from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import _user_id_from_session_cookie + + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", MASTER_KEY, raising=False) + + def _request_with_cookie(cookie_value: str): + from fastapi import Request + + return Request( + scope={ + "type": "http", + "method": "GET", + "path": "/", + "headers": [(b"cookie", f"token={cookie_value}".encode())], + } + ) + + token_object = _token_object() + without_exp = jwt.encode(dict(token_object), MASTER_KEY, algorithm="HS256") + with_exp = encode_ui_session_jwt(token_object, MASTER_KEY) + + assert _user_id_from_session_cookie(_request_with_cookie(without_exp)) is None + assert _user_id_from_session_cookie(_request_with_cookie(with_exp)) == "user-1" + + +def test_no_mint_site_signs_a_ui_token_object_outside_the_choke_point(): + """Pins the 'every UI session cookie goes through encode_ui_session_jwt' claim. + + A new endpoint that mints a session cookie with a bare ``jwt.encode`` would silently + reintroduce the unbounded-lifetime hole, and no behavioral test would catch it because the + cookie it produces still works everywhere except the readers that require ``exp``. + """ + proxy_root = Path(__file__).resolve().parents[4] / "litellm" / "proxy" + offenders = [ + path + for path in proxy_root.rglob("*.py") + if path.name != "ui_session_jwt.py" and "returned_ui_token_object" in path.read_text() and _signs_it(path) + ] + + assert offenders == [], f"these sign a UI token object directly instead of using encode_ui_session_jwt: {offenders}" + + +def _signs_it(path: Path) -> bool: + source = path.read_text() + return any( + "returned_ui_token_object" in source[match : match + 200] + for match in (i for i in range(len(source)) if source.startswith("jwt.encode(", i)) + )