scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 455

fix(sso): tighten oauth_state cookie — Secure flag + PKCE-only set

Browse source 
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).
This commit is contained in:
user 2026-05-01 00:36:51 +00:00
parent 00442e653c
commit 2c852ba2b1
2 changed files with 130 additions and 18 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

44
litellm/proxy/management_endpoints/ui_sso.py
View file

@ -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,
)

104
tests/test_litellm/proxy/management_endpoints/test_ui_sso.py
View file

@ -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):