diff --git a/litellm/constants.py b/litellm/constants.py index 919122d6e4d..70be1e56341 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1814,6 +1814,7 @@ LITELLM_EXECUTED_BATCH_CONCURRENCY: Final = max(1, get_env_int("LITELLM_EXECUTED ########################### CLI SSO AUTHENTICATION CONSTANTS ########################### LITELLM_CLI_SOURCE_IDENTIFIER: Final = "litellm-cli" LITELLM_CLI_SESSION_TOKEN_PREFIX: Final = "litellm-session-token" +CP_RETURN_TO_COOKIE_PREFIX: Final = "return_to.v1." CLI_SSO_SESSION_CACHE_KEY_PREFIX: Final = "cli_sso_session" CLI_SSO_SESSION_TTL_SECONDS: Final = 600 CLI_SESSION_KEY_PREFIX: Final = "cli-session" diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 4cc849762aa..19bca3d2bf8 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -55,6 +55,7 @@ from litellm.constants import ( CLI_SSO_CLAIM_MAX_SCALAR_LENGTH, CLI_SSO_SESSION_CACHE_KEY_PREFIX, CLI_SSO_SESSION_TTL_SECONDS, + CP_RETURN_TO_COOKIE_PREFIX, LITELLM_CLI_SOURCE_IDENTIFIER, LITELLM_UI_SESSION_DURATION, MAX_SPENDLOG_ROWS_TO_QUERY, @@ -2285,9 +2286,7 @@ async def auth_callback(request: Request, state: str | None = None): sso_assertion=sso_assertion, ) - # Control-plane cross-origin: read return_to from cookie. - # Starlette's cookie_parser already handles RFC 2109 unquoting. - cp_return_to: Final[str | None] = request.cookies.get("litellm_cp_return_to") + cp_return_to: Final[str | None] = read_return_to_cookie(request) return await SSOAuthenticationHandler.get_redirect_response_from_openid( result=result, @@ -3090,10 +3089,11 @@ def _persist_return_to_cookie(response: Response, return_to: str | None, request safe: Final = _is_same_origin_return_path(return_to) or SSOAuthenticationHandler.validate_return_to(return_to) except HTTPException: return # a non-matching absolute return_to is ignored, never blocks sign-in - if safe: + sealed: Final = _seal_return_to(return_to) if safe else None + if sealed is not None: response.set_cookie( key="litellm_cp_return_to", - value=return_to, + value=sealed, max_age=600, httponly=True, samesite="lax", @@ -3101,6 +3101,45 @@ def _persist_return_to_cookie(response: Response, return_to: str | None, request ) +def _seal_return_to(return_to: str) -> str | None: + """``return_to`` sealed under the salt key as an opaque cookie value. The plaintext URL carries the + client's loopback ``redirect_uri`` (``http://127.0.0.1:/...`` for Codex or Claude Code), which + edge firewalls read as an SSRF probe and block before the proxy sees the request. None when there + is no salt or master key yet, so the login entrypoint still renders.""" + from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_bearer_token + + try: + return encrypt_bearer_token(return_to, prefix=CP_RETURN_TO_COOKIE_PREFIX) + except ValueError: + return None + + +def _clear_rejected_return_to_cookie(response: Response, request: Request) -> None: + """Expire a ``litellm_cp_return_to`` that ``read_return_to_cookie`` rejected (plaintext, stale, or sealed + under another key) so the browser stops replaying it until it expires on its own.""" + if "litellm_cp_return_to" in request.cookies: + response.delete_cookie("litellm_cp_return_to") + + +def _dashboard_redirect_response(request: Request, jwt_token: str, litellm_dashboard_ui: str) -> RedirectResponse: + """The post-SSO dashboard landing, used when there is no honored ``return_to``. Sets the session + cookie and expires a ``litellm_cp_return_to`` the reader rejected.""" + verbose_proxy_logger.info("Redirecting to %s", litellm_dashboard_ui) + redirect_response: Final = RedirectResponse(url=litellm_dashboard_ui, status_code=303) + set_session_token_cookie(redirect_response, request, jwt_token) + _clear_rejected_return_to_cookie(redirect_response, request) + return redirect_response + + +def read_return_to_cookie(request: Request) -> str | None: + """The ``return_to`` this proxy sealed into ``litellm_cp_return_to``, or None when the cookie is + missing, was not sealed here, or was tampered with.""" + from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_bearer_token + + sealed: Final = request.cookies.get("litellm_cp_return_to") + return None if sealed is None else decrypt_bearer_token(sealed, prefix=CP_RETURN_TO_COOKIE_PREFIX) + + def _is_blank_sso_user_id(user_id: object) -> bool: return isinstance(user_id, str) and not user_id.strip() @@ -3976,10 +4015,7 @@ class SSOAuthenticationHandler: if user_id is not None and isinstance(user_id, str): litellm_dashboard_ui += "?login=success" - verbose_proxy_logger.info("Redirecting to %s", litellm_dashboard_ui) - redirect_response: Final = RedirectResponse(url=litellm_dashboard_ui, status_code=303) - set_session_token_cookie(redirect_response, request, jwt_token) - return redirect_response + return _dashboard_redirect_response(request, jwt_token, litellm_dashboard_ui) @staticmethod @with_service_target(SSO_SESSIONS_TARGET) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c6d3e12a378..967cae27990 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -18400,6 +18400,7 @@ async def login(request: Request): # _is_same_origin_return_path (strictly relative path) so it can never be an open redirect, and the # one-shot cookie is cleared after use. from litellm.proxy.management_endpoints.ui_sso import ( + read_return_to_cookie, set_session_token_cookie, sso_return_to_redirect, ) @@ -18411,7 +18412,7 @@ async def login(request: Request): # return_to is honored for EVERY sign-in branch, so the write and read sets cannot diverge: it # sets the token cookie on the same-origin arm and hands off via a one-time login code on the # cross-origin arm, and clears the one-shot cookie in both. - cp_return_to: Final = request.cookies.get("litellm_cp_return_to") + cp_return_to: Final = read_return_to_cookie(request) if cp_return_to: try: resumed = await sso_return_to_redirect( # rebind-ok: pre-existing rebinding on a rename-only line @@ -18435,7 +18436,7 @@ async def login(request: Request): # Create redirect response with cookie redirect_response: Final = RedirectResponse(url=litellm_dashboard_ui, status_code=303) set_session_token_cookie(redirect_response, request, jwt_token) - if cp_return_to: + if "litellm_cp_return_to" in request.cookies: redirect_response.delete_cookie(key="litellm_cp_return_to") return redirect_response diff --git a/tests/unit/proxy/management_endpoints/test_ui_sso.py b/tests/unit/proxy/management_endpoints/test_ui_sso.py index 873182a5145..479391e54fe 100644 --- a/tests/unit/proxy/management_endpoints/test_ui_sso.py +++ b/tests/unit/proxy/management_endpoints/test_ui_sso.py @@ -33,6 +33,7 @@ from litellm.proxy.management_endpoints.ui_sso import ( insert_sso_user, normalize_email, process_sso_jwt_access_token, + read_return_to_cookie, sso_lookup_user_ids, ) from litellm.types.proxy.management_endpoints.ui_sso import ( @@ -10603,6 +10604,10 @@ class TestPersistReturnToCookieSharedHelper: def _cookie(resp) -> str: return resp.headers.get("set-cookie", "") + @pytest.fixture(autouse=True) + def _master_key(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "sk-test-master") + def test_sets_cookie_for_same_origin_relative_path(self, monkeypatch): from fastapi import Response @@ -10648,6 +10653,62 @@ class TestPersistReturnToCookieSharedHelper: _persist_return_to_cookie(resp, "https://cp.example.com/ui?page=models", _make_https_request()) assert "litellm_cp_return_to=" in self._cookie(resp) + def test_cookie_is_opaque_and_round_trips_the_loopback_redirect_uri(self, monkeypatch: pytest.MonkeyPatch) -> None: + """The DCR authorize round-trip for a local coding agent stores a return_to carrying + ``redirect_uri=http://127.0.0.1:/...``. An edge firewall (AWS WAF + ``EC2MetaDataSSRF_COOKIE``) reads that loopback URL in a cookie as an SSRF probe and 403s the + callback, so the cookie must carry nothing a middlebox can pattern-match, while the proxy + itself still reads the exact return_to back.""" + from fastapi import Response + + from litellm.proxy.management_endpoints.ui_sso import _persist_return_to_cookie, read_return_to_cookie + + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + return_to: Final = "/mcp/authorize?client_id=llm_dcrc_abc&redirect_uri=http%3A%2F%2F127.0.0.1%3A1455%2Fauth%2Fcallback" + resp: Final = Response() + _persist_return_to_cookie(resp, return_to, _make_https_request()) + cookie: Final = self._cookie(resp) + value: Final = cookie.split("litellm_cp_return_to=", 1)[1].split(";", 1)[0] + assert "127.0.0.1" not in cookie and "authorize" not in cookie, cookie + request: Final = _make_https_request() + request.cookies = {"litellm_cp_return_to": value} + assert read_return_to_cookie(request) == return_to + + def test_a_cookie_this_proxy_did_not_seal_reads_as_nothing(self) -> None: + from litellm.proxy.management_endpoints.ui_sso import read_return_to_cookie + + request: Final = _make_https_request() + request.cookies = {"litellm_cp_return_to": "/mcp/authorize?client_id=llm_dcrc_abc"} + assert read_return_to_cookie(request) is None + + def test_no_master_key_means_no_cookie_and_no_error(self, monkeypatch: pytest.MonkeyPatch) -> None: + from fastapi import Response + + from litellm.proxy.management_endpoints.ui_sso import _persist_return_to_cookie + + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.delenv("LITELLM_SALT_KEY", raising=False) + resp: Final = Response() + _persist_return_to_cookie(resp, "/mcp/authorize", _make_https_request()) + assert "litellm_cp_return_to=" not in self._cookie(resp) + + def test_sso_dashboard_fallback_clears_a_return_to_cookie_it_rejected(self, monkeypatch: pytest.MonkeyPatch) -> None: + """A plaintext (pre-upgrade or client-planted) litellm_cp_return_to is not honored by /sso/callback, + and the dashboard landing must also expire it so the browser stops replaying the URL-bearing cookie + (and tripping the edge firewall) for the rest of its lifetime.""" + from litellm.proxy.management_endpoints.ui_sso import _dashboard_redirect_response + + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + request: Final = _make_https_request() + request.cookies = {"litellm_cp_return_to": "/mcp/authorize?client_id=llm_dcrc_abc"} + assert read_return_to_cookie(request) is None + response: Final = _dashboard_redirect_response(request, "jwt", "/ui?login=success") + set_cookies: Final = response.headers.getlist("set-cookie") + assert response.status_code == 303 and response.headers["location"] == "/ui?login=success" + assert any(c.startswith("token=jwt") for c in set_cookies), set_cookies + assert any(c.startswith('litellm_cp_return_to=""') and "Max-Age=0" in c for c in set_cookies), set_cookies + def test_cookie_is_secure_and_httponly_over_https(self, monkeypatch): from fastapi import Response diff --git a/tests/unit/proxy/proxy_server/test_routes_login_sso.py b/tests/unit/proxy/proxy_server/test_routes_login_sso.py index 75e14a89148..43629b926ad 100644 --- a/tests/unit/proxy/proxy_server/test_routes_login_sso.py +++ b/tests/unit/proxy/proxy_server/test_routes_login_sso.py @@ -11,8 +11,12 @@ Routes covered: from __future__ import annotations from datetime import datetime, timezone +from typing import Final from unittest.mock import AsyncMock, MagicMock +import pytest +from fastapi.testclient import TestClient + from .conftest import normalize # --------------------------------------------------------------------------- @@ -398,6 +402,29 @@ def test_v3_login_exchange_success_returns_token_and_redirect(client, monkeypatc } +def _sealed_return_to(return_to: str) -> str: + """The cookie value /sso/key/generate writes: return_to sealed under the master key, never plaintext.""" + from litellm.constants import CP_RETURN_TO_COOKIE_PREFIX + from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_bearer_token + + return encrypt_bearer_token(return_to, prefix=CP_RETURN_TO_COOKIE_PREFIX) + + +def test_login_form_ignores_a_plaintext_return_to_cookie(client: TestClient, monkeypatch: pytest.MonkeyPatch) -> None: + """A return_to the proxy did not seal (a client-planted or pre-upgrade plaintext cookie) is not + honored: /login lands on the dashboard and clears the cookie.""" + _install_login_mocks(monkeypatch) + response: Final = client.post( + "/login", + data={"username": "admin", "password": "password"}, + cookies={"litellm_cp_return_to": "/mcp/authorize?client_id=llm_dcrc_abc"}, + follow_redirects=False, + ) + assert response.status_code == 303 + assert "/ui" in response.headers.get("location", "") + assert 'litellm_cp_return_to=""' in response.headers.get("set-cookie", "") + + def test_login_form_honors_same_origin_return_to_cookie(client, monkeypatch): """The aggregate DCR connect flow preserves a same-origin return_to in the litellm_cp_return_to cookie; /login must RESUME there after password sign-in instead of dead-ending at the dashboard.""" @@ -406,7 +433,7 @@ def test_login_form_honors_same_origin_return_to_cookie(client, monkeypatch): response = client.post( "/login", data={"username": "admin", "password": "password"}, - cookies={"litellm_cp_return_to": return_to}, + cookies={"litellm_cp_return_to": _sealed_return_to(return_to)}, follow_redirects=False, ) assert response.status_code == 303 @@ -425,7 +452,7 @@ def test_login_form_honors_control_plane_return_to_cookie(client, monkeypatch): response = client.post( "/login", data={"username": "admin", "password": "password"}, - cookies={"litellm_cp_return_to": "https://cp.example.com/console"}, + cookies={"litellm_cp_return_to": _sealed_return_to("https://cp.example.com/console")}, follow_redirects=False, ) location = response.headers.get("location", "") @@ -447,7 +474,7 @@ def test_login_form_survives_stale_control_plane_return_to(client, monkeypatch): response = client.post( "/login", data={"username": "admin", "password": "password"}, - cookies={"litellm_cp_return_to": "https://old-cp.example.com/console"}, + cookies={"litellm_cp_return_to": _sealed_return_to("https://old-cp.example.com/console")}, follow_redirects=False, ) assert response.status_code == 303, "login must not break on a stale return_to cookie" @@ -463,7 +490,7 @@ def test_login_form_ignores_open_redirect_return_to(client, monkeypatch): response = client.post( "/login", data={"username": "admin", "password": "password"}, - cookies={"litellm_cp_return_to": "https://evil.example.com/steal"}, + cookies={"litellm_cp_return_to": _sealed_return_to("https://evil.example.com/steal")}, follow_redirects=False, ) assert response.status_code == 303