From 2c852ba2b1a6729ad36e922898b670b353d111a1 Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Fri, 1 May 2026 00:36:51 +0000 Subject: [PATCH] =?UTF-8?q?fix(sso):=20tighten=20oauth=5Fstate=20cookie=20?= =?UTF-8?q?=E2=80=94=20Secure=20flag=20+=20PKCE-only=20set?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two Greptile review findings addressed: 1. (P1, security) The ``litellm_oauth_state`` cookie is the sole guard against Login-CSRF in the PKCE flow but was set without the ``Secure`` attribute, so a network observer on plain HTTP could read and replay it — bypassing the protection this PR adds. Thread the originating ``Request`` down through ``get_sso_login_redirect`` and ``get_generic_sso_redirect_response`` and set ``Secure`` based on ``request.url.scheme == "https"``. When no request is supplied (programmatic callers / tests) default to ``Secure=True`` — production-safe. Local HTTP dev still works because the request scheme is observed at runtime. 2. (P2) The cookie was set unconditionally, but the callback only validates it inside the PKCE branch. Two concurrent SSO sessions (one PKCE, one plain) could overwrite each other's state cookie and produce spurious 400s for the plain-flow user. Move the ``set_cookie`` call inside the existing ``if code_verifier and "state" in redirect_params`` block so the cookie is only written when PKCE is active and the validation will actually fire. Tests cover both paths: PKCE-on (cookie set with Secure default), PKCE-off (cookie not set), and HTTP dev request (Secure dropped so the browser will actually attach the cookie on the callback hop). --- litellm/proxy/management_endpoints/ui_sso.py | 44 +++++--- .../proxy/management_endpoints/test_ui_sso.py | 104 +++++++++++++++++- 2 files changed, 130 insertions(+), 18 deletions(-) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index ca9fee1a76f..8225b973206 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -674,6 +674,7 @@ async def google_login( google_client_id=google_client_id, generic_client_id=generic_client_id, state=cli_state, + request=request, ) if return_to is not None and sso_redirect is not None: if SSOAuthenticationHandler._validate_return_to(return_to): @@ -2170,6 +2171,7 @@ class SSOAuthenticationHandler: microsoft_client_id: Optional[str] = None, generic_client_id: Optional[str] = None, state: Optional[str] = None, + request: Optional[Request] = None, ) -> Optional[RedirectResponse]: """ Step 1. Call Get Login Redirect for the SSO provider. Send the redirect response to `redirect_url` @@ -2179,6 +2181,8 @@ class SSOAuthenticationHandler: google_client_id (Optional[str], optional): The Google Client ID. Defaults to None. microsoft_client_id (Optional[str], optional): The Microsoft Client ID. Defaults to None. generic_client_id (Optional[str], optional): The Generic Client ID. Defaults to None. + request: Optional FastAPI request, used to drive the ``Secure`` + attribute on the ``litellm_oauth_state`` CSRF cookie. Returns: RedirectResponse: The redirect response from the SSO provider. @@ -2289,6 +2293,7 @@ class SSOAuthenticationHandler: generic_sso=generic_sso, state=state, generic_authorization_endpoint=generic_authorization_endpoint, + request=request, ) raise ValueError( "Unknown SSO provider. Please setup SSO with client IDs https://docs.litellm.ai/docs/proxy/admin_ui_sso" @@ -2299,6 +2304,7 @@ class SSOAuthenticationHandler: generic_sso: Any, state: Optional[str] = None, generic_authorization_endpoint: Optional[str] = None, + request: Optional[Request] = None, ) -> Optional[RedirectResponse]: """ Get the redirect response for Generic SSO @@ -2382,19 +2388,30 @@ class SSOAuthenticationHandler: # Update the redirect response redirect_response.headers["location"] = new_url - # Bind state to the user's browser session. The /callback handler - # validates the URL ``state`` against this cookie via - # ``secrets.compare_digest`` before exchanging the PKCE - # code_verifier. - state_value = redirect_params.get("state") - if state_value and redirect_response is not None: - redirect_response.set_cookie( - key="litellm_oauth_state", - value=state_value, - max_age=600, - httponly=True, - samesite="lax", - ) + # Bind state to the user's browser session. The /callback + # handler validates the URL ``state`` against this cookie via + # ``secrets.compare_digest`` before exchanging the PKCE + # code_verifier. Only set the cookie when PKCE is in use + # (i.e. inside this ``code_verifier`` branch) so two + # concurrent SSO sessions — one PKCE, one plain — cannot + # overwrite each other's state cookie. + state_value = redirect_params.get("state") + if state_value and redirect_response is not None: + # Production-safe default: require HTTPS for the + # CSRF-protection cookie unless we can prove the + # incoming request is HTTP (local dev). Without + # ``Secure`` the cookie is sent over plain HTTP, + # letting a network observer read and replay the + # state value and bypass this protection. + secure_flag = request is None or request.url.scheme == "https" + redirect_response.set_cookie( + key="litellm_oauth_state", + value=state_value, + max_age=600, + httponly=True, + samesite="lax", + secure=secure_flag, + ) return redirect_response @staticmethod @@ -4012,6 +4029,7 @@ async def debug_sso_login(request: Request): microsoft_client_id=microsoft_client_id, google_client_id=google_client_id, generic_client_id=generic_client_id, + request=request, ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index 019d6218740..231fca36cf6 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -5782,17 +5782,18 @@ class TestPKCEStateCookieBinding: the PKCE token exchange.""" @pytest.mark.asyncio - async def test_redirect_response_sets_oauth_state_cookie(self): + async def test_redirect_response_sets_oauth_state_cookie_when_pkce_enabled(self): """``get_generic_sso_redirect_response`` must set - ``litellm_oauth_state`` on the redirect response so the callback - can verify it later.""" + ``litellm_oauth_state`` on the redirect response when PKCE is on so + the callback can verify it later. The cookie must carry HttpOnly, + SameSite=Lax, and (because no http request was supplied to the + helper) the production-safe ``Secure`` default.""" from fastapi.responses import RedirectResponse from litellm.proxy.management_endpoints.ui_sso import ( SSOAuthenticationHandler, ) - # generic_sso is a context manager + redirect-response factory. mock_redirect = RedirectResponse( url="https://idp.example.com/authorize?state=test-state-xyz" ) @@ -5801,7 +5802,13 @@ class TestPKCEStateCookieBinding: mock_generic_sso.__exit__ = MagicMock(return_value=None) mock_generic_sso.get_login_redirect = AsyncMock(return_value=mock_redirect) - with patch.dict(os.environ, {"GENERIC_CLIENT_STATE": "test-state-xyz"}): + with patch.dict( + os.environ, + { + "GENERIC_CLIENT_STATE": "test-state-xyz", + "GENERIC_CLIENT_USE_PKCE": "true", + }, + ): response = await SSOAuthenticationHandler.get_generic_sso_redirect_response( generic_sso=mock_generic_sso, state=None, @@ -5819,6 +5826,93 @@ class TestPKCEStateCookieBinding: assert "test-state-xyz" in cookie_str assert "HttpOnly" in cookie_str assert "SameSite=lax" in cookie_str + # No incoming Request supplied → ``Secure`` defaults to True so a + # network observer on plain HTTP cannot read the state value. + assert "Secure" in cookie_str + + @pytest.mark.asyncio + async def test_redirect_response_omits_oauth_state_cookie_when_pkce_disabled( + self, + ): + """Non-PKCE flows delegate to fastapi-sso's own session-cookie + binding; we do not set our cookie there because it would never be + validated (and could collide with a concurrent PKCE session in + the same browser).""" + from fastapi.responses import RedirectResponse + + from litellm.proxy.management_endpoints.ui_sso import ( + SSOAuthenticationHandler, + ) + + mock_redirect = RedirectResponse( + url="https://idp.example.com/authorize?state=test-state-xyz" + ) + mock_generic_sso = MagicMock() + mock_generic_sso.__enter__ = MagicMock(return_value=mock_generic_sso) + mock_generic_sso.__exit__ = MagicMock(return_value=None) + mock_generic_sso.get_login_redirect = AsyncMock(return_value=mock_redirect) + + with patch.dict( + os.environ, + { + "GENERIC_CLIENT_STATE": "test-state-xyz", + "GENERIC_CLIENT_USE_PKCE": "false", + }, + ): + response = await SSOAuthenticationHandler.get_generic_sso_redirect_response( + generic_sso=mock_generic_sso, + state=None, + generic_authorization_endpoint="https://idp.example.com/authorize", + ) + + assert response is not None + cookie_headers = response.headers.getlist("set-cookie") + assert not any( + "litellm_oauth_state=" in c for c in cookie_headers + ), f"litellm_oauth_state cookie set on non-PKCE flow; got: {cookie_headers}" + + @pytest.mark.asyncio + async def test_redirect_response_drops_secure_flag_for_http_dev(self): + """When the incoming request is plain HTTP (local dev), ``Secure`` + must be dropped so the browser will actually attach the cookie on + the callback hop.""" + from fastapi.responses import RedirectResponse + + from litellm.proxy.management_endpoints.ui_sso import ( + SSOAuthenticationHandler, + ) + + mock_redirect = RedirectResponse( + url="http://idp.local/authorize?state=local-dev-state" + ) + mock_generic_sso = MagicMock() + mock_generic_sso.__enter__ = MagicMock(return_value=mock_generic_sso) + mock_generic_sso.__exit__ = MagicMock(return_value=None) + mock_generic_sso.get_login_redirect = AsyncMock(return_value=mock_redirect) + + http_request = MagicMock(spec=Request) + http_request.url.scheme = "http" + + with patch.dict( + os.environ, + { + "GENERIC_CLIENT_STATE": "local-dev-state", + "GENERIC_CLIENT_USE_PKCE": "true", + }, + ): + response = await SSOAuthenticationHandler.get_generic_sso_redirect_response( + generic_sso=mock_generic_sso, + state=None, + generic_authorization_endpoint="http://idp.local/authorize", + request=http_request, + ) + + cookie_headers = response.headers.getlist("set-cookie") + cookie_str = next( + (c for c in cookie_headers if "litellm_oauth_state=" in c), None + ) + assert cookie_str is not None + assert "Secure" not in cookie_str @pytest.mark.asyncio async def test_pkce_callback_rejects_missing_cookie(self):