diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index a4d2977a971..1d0583d7ebb 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -989,6 +989,67 @@ def _merge_sso_token_claims( ) +def _id_token_claims_for_unreadable_userinfo( + id_token: str | None, userinfo_endpoint: str | None, error: json.JSONDecodeError +) -> Mapping[str, object]: + claims: Final = _decode_sso_token_claims(id_token) + if not claims: + raise ProxyException( + message=( + f"SSO user info unavailable: userinfo endpoint {userinfo_endpoint} returned a non-JSON " + "response and the token response had no decodable id_token" + ), + type=ProxyErrorTypes.auth_error, + param="userinfo", + code=status.HTTP_401_UNAUTHORIZED, + ) + verbose_proxy_logger.warning( + "Userinfo endpoint %s returned a non-JSON response (%s), falling back to id_token claims", + userinfo_endpoint, + error, + ) + return claims + + +class _UserinfoFallbackCapableSSO(Protocol): + @property + def access_token(self) -> str | None: ... + + @property + def id_token(self) -> str | None: ... + + async def verify_and_process( + self, + request: Request, + *, + params: dict[str, object] | None = None, # mutable-ok: fastapi-sso verify_and_process requires dict params + headers: dict[str, str] | None = None, # mutable-ok: fastapi-sso verify_and_process requires dict headers + ) -> CustomOpenID | None: ... + + +async def _verify_generic_sso_or_fall_back_to_id_token( + generic_sso: _UserinfoFallbackCapableSSO, + request: Request, + token_exchange_params: dict[str, object], # mutable-ok: fastapi-sso verify_and_process requires dict params + additional_headers: dict[str, str], # mutable-ok: fastapi-sso verify_and_process requires dict headers + userinfo_endpoint: str | None, + response_convertor: Callable[[Mapping[str, object], object], CustomOpenID], +) -> CustomOpenID | None: + try: + return await generic_sso.verify_and_process( + request, + params=token_exchange_params, + headers=additional_headers, + ) + except json.JSONDecodeError as userinfo_error: + if generic_sso.access_token is None: + raise + return response_convertor( + _id_token_claims_for_unreadable_userinfo(generic_sso.id_token, userinfo_endpoint, userinfo_error), + generic_sso, + ) + + async def _raise_if_sso_exceeds_free_user_limit(premium_user: bool, prisma_client: PrismaClient | None) -> None: """Free tier allows SSO for up to 5 billable users; beyond that requires an Enterprise license.""" if premium_user is True: @@ -1655,10 +1716,15 @@ async def get_generic_sso_response( # process_sso_jwt_access_token can extract JWT-embedded roles/teams. access_token_str: str | None = combined_response.get("access_token") else: - result = await generic_sso.verify_and_process( - request, - params=token_exchange_params, - headers=additional_generic_sso_headers_dict, + result = await _verify_generic_sso_or_fall_back_to_id_token( + generic_sso=cast( # cast-ok: fastapi-sso builds the provider class dynamically via create_provider + _UserinfoFallbackCapableSSO, generic_sso + ), + request=request, + token_exchange_params=token_exchange_params, + additional_headers=additional_generic_sso_headers_dict, + userinfo_endpoint=generic_userinfo_endpoint, + response_convertor=response_convertor, ) access_token_str = generic_sso.access_token sso_assertion = assertion_from_sso_login(generic_sso.id_token, generic_sso.refresh_token) diff --git a/tests/unit/proxy/management_endpoints/test_ui_sso.py b/tests/unit/proxy/management_endpoints/test_ui_sso.py index 4e128440703..db5ec90335a 100644 --- a/tests/unit/proxy/management_endpoints/test_ui_sso.py +++ b/tests/unit/proxy/management_endpoints/test_ui_sso.py @@ -8914,6 +8914,143 @@ async def test_verify_and_process_arm_captures_sso_assertion(): assert sso_assertion.refresh_token.get_secret_value() == "rt_from_idp" +_GENERIC_SSO_CALLBACK_ENV = { + "GENERIC_CLIENT_SECRET": "test_secret", + "GENERIC_AUTHORIZATION_ENDPOINT": "https://auth.example.com/auth", + "GENERIC_TOKEN_ENDPOINT": "https://auth.example.com/token", + "GENERIC_USERINFO_ENDPOINT": "https://auth.example.com/userinfo", + "GENERIC_CLIENT_USE_PKCE": "", +} + + +def _generic_sso_callback_request() -> Request: + from starlette.requests import Request as StarletteRequest + + return StarletteRequest( + { + "type": "http", + "method": "GET", + "scheme": "https", + "server": ("proxy.example.com", 443), + "path": "/sso/callback", + "query_string": b"code=abc", + "headers": [], + } + ) + + +def _generic_sso_jwt_handler() -> MagicMock: + jwt_handler = MagicMock(spec=JWTHandler) + jwt_handler.get_team_ids_from_jwt.return_value = [] + return jwt_handler + + +@pytest.mark.asyncio +async def test_non_json_userinfo_falls_back_to_id_token_claims(): + """A userinfo endpoint answering with a non-JSON body (ADFS 401, Content-Length 0) + must fall back to id_token claims on the non-PKCE arm instead of surfacing a 500.""" + import jwt as pyjwt + + from litellm.proxy.management_endpoints.ui_sso import get_generic_sso_response + + id_token = pyjwt.encode( + {"sub": "u-adfs", "preferred_username": "u-adfs", "email": "adfs@example.com"}, + "test-idp-signing-key-32-bytes-long-xxxx", + algorithm="HS256", + ) + request = _generic_sso_callback_request() + + with ( + respx.mock, + patch.dict(os.environ, _GENERIC_SSO_CALLBACK_ENV), + ): + respx.post("https://auth.example.com/token").mock( + return_value=httpx.Response(200, json={"access_token": "at", "token_type": "Bearer", "id_token": id_token}) + ) + userinfo_route = respx.get("https://auth.example.com/userinfo").mock( + return_value=httpx.Response(401, content=b"", headers={"content-type": "text/html"}) + ) + result, received_response, _, _ = await get_generic_sso_response( + request=request, + jwt_handler=_generic_sso_jwt_handler(), + generic_client_id="test_client_id", + redirect_url="https://proxy.example.com/sso/callback", + sso_jwt_handler=None, + ) + + assert result.email == "adfs@example.com" + assert result.id == "u-adfs" + assert received_response is not None + assert received_response["email"] == "adfs@example.com" + assert received_response["sub"] == "u-adfs" + assert "id_token" not in received_response + assert "access_token" not in received_response + assert userinfo_route.call_count == 1 + + +@pytest.mark.asyncio +async def test_non_json_userinfo_without_id_token_raises_401(): + """Same unreadable userinfo but no decodable id_token in the token response: + the login must fail closed with a 401 ProxyException naming the userinfo endpoint.""" + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.ui_sso import get_generic_sso_response + + request = _generic_sso_callback_request() + + with ( + respx.mock, + patch.dict(os.environ, _GENERIC_SSO_CALLBACK_ENV), + ): + respx.post("https://auth.example.com/token").mock( + return_value=httpx.Response(200, json={"access_token": "at", "token_type": "Bearer"}) + ) + respx.get("https://auth.example.com/userinfo").mock( + return_value=httpx.Response(401, content=b"", headers={"content-type": "text/html"}) + ) + with pytest.raises(ProxyException) as exc_info: + await get_generic_sso_response( + request=request, + jwt_handler=_generic_sso_jwt_handler(), + generic_client_id="test_client_id", + redirect_url="https://proxy.example.com/sso/callback", + sso_jwt_handler=None, + ) + + assert str(exc_info.value.code) == "401" + assert "https://auth.example.com/userinfo" in exc_info.value.message + + +@pytest.mark.asyncio +async def test_non_json_token_response_is_not_swallowed_by_userinfo_fallback(): + """A non-JSON TOKEN endpoint response must keep the old behaviour (raw JSONDecodeError + through _handle_generic_sso_error): no access token was parsed, so the id_token + fallback must not fire and userinfo must never be called.""" + from litellm.proxy.management_endpoints.ui_sso import get_generic_sso_response + + request = _generic_sso_callback_request() + + with ( + respx.mock, + patch.dict(os.environ, _GENERIC_SSO_CALLBACK_ENV), + ): + respx.post("https://auth.example.com/token").mock( + return_value=httpx.Response(200, content=b"", headers={"content-type": "text/html"}) + ) + userinfo_route = respx.get("https://auth.example.com/userinfo").mock( + return_value=httpx.Response(200, json={"sub": "unused"}) + ) + with pytest.raises(json.JSONDecodeError): + await get_generic_sso_response( + request=request, + jwt_handler=_generic_sso_jwt_handler(), + generic_client_id="test_client_id", + redirect_url="https://proxy.example.com/sso/callback", + sso_jwt_handler=None, + ) + + assert not userinfo_route.called + + @pytest.mark.asyncio async def test_redirect_from_openid_persists_assertion_under_canonical_user_id(): """The browser funnel persists the captured assertion AFTER canonical user