From 0c647d4be9eb2eb2079d02e4ce277f72daa9c7e5 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Fri, 6 Mar 2026 12:49:12 -0800 Subject: [PATCH] address greptile review feedback (greploop iteration 39) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - fix credential leakage: directly assign received_response from combined_response instead of relying on nonlocal mutation; Pyright was flagging the old guard as unreachable, meaning credential stripping might not execute — now it always runs unconditionally - add test for legacy plain-string cache format backward compat branch - add test for HTTP 200 with no error field and no access_token (else branch) - add test for HTTP 200 with JSON null body (new AttributeError guard) --- litellm/proxy/management_endpoints/ui_sso.py | 19 ++-- .../proxy/management_endpoints/test_ui_sso.py | 99 ++++++++++++++++++- 2 files changed, 107 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index f6b20dc6e73..67f09e3a459 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -845,16 +845,15 @@ async def get_generic_sso_response( # Pass the full response so custom response_convertor implementations # can access all fields (including id_token for claim extraction). result = response_convertor(combined_response, generic_sso) - # Strip bearer credentials from received_response after conversion. - # received_response may appear in restricted-group error messages — - # do not expose tokens to callers. - # Note: response_convertor always sets received_response via nonlocal, so it is - # non-None here at runtime. The guard satisfies Pyright's type narrowing since - # it cannot track nonlocal mutations inside closures. - if received_response is not None: - received_response = { - k: v for k, v in received_response.items() if k not in _OAUTH_TOKEN_FIELDS - } + # Strip bearer credentials from combined_response before storing in + # received_response. received_response may appear in restricted-group + # error messages — bearer tokens (access_token, id_token, refresh_token) + # must not be exposed to callers. + # Assign directly rather than relying on nonlocal mutation so that Pyright + # can track that received_response is non-None from this point on. + received_response = { + k: v for k, v in combined_response.items() if k not in _OAUTH_TOKEN_FIELDS + } # In the PKCE path verify_and_process is skipped, so generic_sso.access_token # is never set. Read the token directly from the exchange response instead so # process_sso_jwt_access_token can extract JWT-embedded roles/teams. 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 e9681f6d17c..1c62a70c51d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -3663,7 +3663,7 @@ class TestPKCEFunctionality: ) assert "verifier not found" in exc_info.value.message.lower() or "cache" in exc_info.value.message.lower() - assert exc_info.value.code == 401 + assert str(exc_info.value.code) == "401" @pytest.mark.asyncio @@ -3883,6 +3883,103 @@ class TestPKCEFunctionality: assert "code_verifier" not in result assert "_pkce_cache_key" not in result + @pytest.mark.asyncio + async def test_pkce_legacy_string_cache_format_backward_compat(self): + """Legacy plain-string cache entries (stored before dict format was introduced) + are handled transparently via the backward-compat branch.""" + import os + from unittest.mock import AsyncMock, MagicMock, patch + + from starlette.requests import Request + + from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler + + legacy_verifier = "legacy_plain_string_verifier_abc123" + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=legacy_verifier) + + mock_request = MagicMock(spec=Request) + mock_request.query_params = {"state": "legacy_state_xyz"} + + with patch("litellm.proxy.proxy_server.redis_usage_cache", None), patch( + "litellm.proxy.proxy_server.user_api_key_cache", mock_cache + ), patch.dict(os.environ, {"GENERIC_CLIENT_USE_PKCE": "true"}, clear=False): + result = await SSOAuthenticationHandler.prepare_token_exchange_parameters( + request=mock_request, generic_include_client_id=False + ) + + assert result["code_verifier"] == legacy_verifier + assert result["_pkce_cache_key"] == "pkce_verifier:legacy_state_xyz" + + @pytest.mark.asyncio + async def test_pkce_token_exchange_null_json_body_raises_proxy_exception(self): + """HTTP 200 with JSON body `null` raises a clean ProxyException instead of + AttributeError when .get() is called on the None return value.""" + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler + + with patch("litellm.proxy.management_endpoints.ui_sso.httpx.AsyncClient") as mock_client_cls: + mock_client = AsyncMock() + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=False) + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.json.return_value = None # JSON null response body + mock_resp.text = "null" + mock_client.post = AsyncMock(return_value=mock_resp) + mock_client_cls.return_value = mock_client + + with pytest.raises(ProxyException) as exc_info: + await SSOAuthenticationHandler._pkce_token_exchange( + authorization_code="some_code", + code_verifier="verifier", + client_id="cid", + client_secret="csecret", + token_endpoint="https://example.com/token", + userinfo_endpoint=None, + include_client_id=False, + redirect_url=None, + additional_headers={}, + ) + + assert "unexpected response format" in exc_info.value.message.lower() + assert str(exc_info.value.code) == "401" + + @pytest.mark.asyncio + async def test_pkce_token_exchange_http200_no_error_field_no_access_token(self): + """HTTP 200 with no error field and no access_token raises ProxyException + with a descriptive message showing the actual response keys.""" + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler + + body_without_token = {"token_type": "Bearer", "scope": "openid"} + + with patch("litellm.proxy.management_endpoints.ui_sso.httpx.AsyncClient") as mock_client_cls: + mock_client = AsyncMock() + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=False) + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.json.return_value = body_without_token + mock_client.post = AsyncMock(return_value=mock_resp) + mock_client_cls.return_value = mock_client + + with pytest.raises(ProxyException) as exc_info: + await SSOAuthenticationHandler._pkce_token_exchange( + authorization_code="some_code", + code_verifier="verifier", + client_id="cid", + client_secret="csecret", + token_endpoint="https://example.com/token", + userinfo_endpoint=None, + include_client_id=False, + redirect_url=None, + additional_headers={}, + ) + + assert "no access_token" in exc_info.value.message or "access_token" in exc_info.value.message + assert str(exc_info.value.code) == "401" + # Tests for SSO user team assignment bug (Issue: SSO Users Not Added to Entra-Synced Teams on First Login) class TestAddMissingTeamMember: