From 3cac5e5cd4c12a782e0afe96218aaff986ef3f60 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Wed, 2 Sep 2026 22:44:17 -0700 Subject: [PATCH] fix(sso): resolve multi-valued role claims to the highest privilege role (#39480) * fix(sso): resolve multi-valued role claims to the highest privilege role A role claim carrying several roles used to resolve to whichever one the IdP listed first, so a user holding both proxy_admin_viewer and internal_user lost org-level spend visibility depending on claim ordering alone. get_litellm_user_role now picks the highest privilege role out of a list-valued claim, and the Entra app_roles path shares that same resolution instead of keeping its own copy of the hierarchy. SAML assertions carrying several role values go through the same path rather than taking the first value. * test(sso): lock ranked-over-unranked role resolution for mixed claims org_admin, team and customer sit outside the privilege ladder. Pin the resolution for a claim that mixes one of them with a ranked role so the asymmetry is covered rather than implicit. * fix(sso): label the claim-sequence cast for the type-discipline gate * fix(sso): resolve claim entries without recursing The repo's recursive-function gate rejects self-recursion here, and a role claim is flat anyway. Pull the single-value lookup into its own helper so the list branch maps over it instead of calling back into itself. --- .../management_endpoints/sso/saml_sso.py | 6 +- litellm/proxy/management_endpoints/types.py | 62 +++++++-- litellm/proxy/management_endpoints/ui_sso.py | 19 +-- ruff-strict-budget.json | 4 +- .../management_endpoints/test_saml_sso.py | 33 +++++ .../proxy/management_endpoints/test_ui_sso.py | 129 +++++++++++++++++- type-discipline-budget.json | 4 +- 7 files changed, 218 insertions(+), 39 deletions(-) diff --git a/litellm/proxy/management_endpoints/sso/saml_sso.py b/litellm/proxy/management_endpoints/sso/saml_sso.py index 466b100ea1f..12e1f1a03f3 100644 --- a/litellm/proxy/management_endpoints/sso/saml_sso.py +++ b/litellm/proxy/management_endpoints/sso/saml_sso.py @@ -443,7 +443,9 @@ class SAMLAuthHandler: last_name: Final = SAMLAuthHandler._attribute_value( attributes, "SAML_ATTRIBUTE_LAST_NAME", _LAST_NAME_ATTRIBUTE_CANDIDATES ) - role_value = SAMLAuthHandler._attribute_value(attributes, "SAML_ATTRIBUTE_ROLE", _ROLE_ATTRIBUTE_CANDIDATES) + role_values: Final = SAMLAuthHandler._attribute_values( + attributes, "SAML_ATTRIBUTE_ROLE", _ROLE_ATTRIBUTE_CANDIDATES + ) team_ids: Final = SAMLAuthHandler._attribute_values( attributes, "SAML_ATTRIBUTE_TEAM_IDS", _TEAM_IDS_ATTRIBUTE_CANDIDATES ) @@ -464,7 +466,7 @@ class SAMLAuthHandler: picture=None, provider="saml", team_ids=team_ids, - user_role=get_litellm_user_role(role_value) if role_value else None, + user_role=get_litellm_user_role(role_values), ) except ValidationError as e: raise HTTPException( diff --git a/litellm/proxy/management_endpoints/types.py b/litellm/proxy/management_endpoints/types.py index 070df97d09a..4414eed97b2 100644 --- a/litellm/proxy/management_endpoints/types.py +++ b/litellm/proxy/management_endpoints/types.py @@ -4,12 +4,44 @@ Types for the management endpoints Might include fastapi/proxy requirements.txt related imports """ +from collections.abc import Iterable, Sequence from typing import Any, Final, cast from fastapi_sso.sso.base import OpenID from litellm.proxy._types import LitellmUserRoles +# Ordered highest to lowest privilege +LITELLM_USER_ROLE_HIERARCHY: Final = ( + LitellmUserRoles.PROXY_ADMIN, + LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, + LitellmUserRoles.INTERNAL_USER, + LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, +) + + +def highest_privilege_role(roles: Iterable[LitellmUserRoles]) -> LitellmUserRoles | None: + """ + Pick the highest privilege role out of the roles an IdP asserted for one user. + + IdPs do not guarantee ordering within a multi-valued role claim, so a user holding + several roles resolves to the most privileged one rather than whichever came first. + Roles the hierarchy does not rank (org_admin, team, customer) resolve by name to stay + deterministic. + + Args: + roles: The roles resolved from the claim + + Returns: + The highest privilege role, or None if `roles` is empty + """ + resolved: Final = frozenset(roles) + if not resolved: + return None + + ranked: Final = next((role for role in LITELLM_USER_ROLE_HIERARCHY if role in resolved), None) + return ranked if ranked is not None else min(resolved, key=lambda role: role.value) + def is_valid_litellm_user_role(role_str: str) -> bool: """ @@ -28,12 +60,22 @@ def is_valid_litellm_user_role(role_str: str) -> bool: return False -def get_litellm_user_role(role_str) -> LitellmUserRoles | None: +def _role_from_claim_value(role_str: object) -> LitellmUserRoles | None: + if not isinstance(role_str, str): + return None + # Use _value2member_map_ for O(1) lookup, case-insensitive + result: Final = LitellmUserRoles._value2member_map_.get(role_str.lower()) + return cast(LitellmUserRoles | None, result) + + +def get_litellm_user_role(role_str: object) -> LitellmUserRoles | None: """ Convert a string (or list of strings) to a LitellmUserRoles enum if valid (case-insensitive). Handles list inputs since some SSO providers (e.g., Keycloak) return roles - as arrays like ["proxy_admin"] instead of plain strings. + as arrays like ["proxy_admin"] instead of plain strings. A claim carrying several + roles resolves to the highest privilege one, so a user does not lose access just + because the IdP listed a weaker role first. Args: role_str: String or list to convert (e.g., "proxy_admin", ["proxy_admin"]) @@ -41,16 +83,12 @@ def get_litellm_user_role(role_str) -> LitellmUserRoles | None: Returns: LitellmUserRoles enum if valid, None otherwise """ - try: - if isinstance(role_str, list): - if len(role_str) == 0: - return None - role_str = role_str[0] - # Use _value2member_map_ for O(1) lookup, case-insensitive - result: Final = LitellmUserRoles._value2member_map_.get(role_str.lower()) - return cast(LitellmUserRoles | None, result) - except Exception: - return None + if isinstance(role_str, (list, tuple)): + entries: Final = cast(Sequence[object], role_str) # cast-ok: isinstance narrows the claim, not its elements + return highest_privilege_role( + role for role in (_role_from_claim_value(entry) for entry in entries) if role is not None + ) + return _role_from_claim_value(role_str) class CustomOpenID(OpenID): diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 1feefa5725d..84150ef7935 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -112,6 +112,7 @@ from litellm.proxy.management_endpoints.sso_helper_utils import ( ) from litellm.proxy.management_endpoints.team_endpoints import new_team, team_member_add from litellm.proxy.management_endpoints.types import ( + LITELLM_USER_ROLE_HIERARCHY, CustomOpenID, get_litellm_user_role, is_valid_litellm_user_role, @@ -809,15 +810,6 @@ def normalize_email(email: str | None) -> str | None: return email.lower() if isinstance(email, str) else email -# Ordered highest to lowest privilege -LITELLM_USER_ROLE_HIERARCHY: Final = ( - LitellmUserRoles.PROXY_ADMIN, - LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, - LitellmUserRoles.INTERNAL_USER, - LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, -) - - def determine_role_from_groups( user_groups: list[str], role_mappings: "RoleMappings", @@ -4312,14 +4304,7 @@ class MicrosoftSSOHandler: listed first. Roles the hierarchy does not rank (org_admin, team, customer) resolve by name to stay deterministic """ - resolved: Final = frozenset( - role for role in (get_litellm_user_role(role_str) for role_str in app_roles or ()) if role is not None - ) - if not resolved: - return None - - ranked: Final = next((role for role in LITELLM_USER_ROLE_HIERARCHY if role in resolved), None) - return ranked if ranked is not None else min(resolved, key=lambda role: role.value) + return get_litellm_user_role(tuple(app_roles or ())) @staticmethod def get_app_roles_from_id_token(id_token: str | None) -> list[str]: diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 4fcf650a8bc..be2b30fc189 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,6 +1,6 @@ { "ANN001": { - "limit": 2985 + "limit": 2984 }, "ANN002": { "limit": 71 @@ -57,7 +57,7 @@ "limit": 3 }, "BLE001": { - "limit": 2917 + "limit": 2916 }, "C401": { "limit": 8 diff --git a/tests/test_litellm/proxy/management_endpoints/test_saml_sso.py b/tests/test_litellm/proxy/management_endpoints/test_saml_sso.py index 36d6414ba9f..635e6332958 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_saml_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_saml_sso.py @@ -516,6 +516,39 @@ async def test_team_ids_extracted_from_groups_attribute(saml_env_idp_initiated): assert result.team_ids == ["team-a", "team-b"] +@pytest.mark.asyncio +@pytest.mark.parametrize( + "roles", + [ + ["internal_user", "proxy_admin_viewer"], + ["proxy_admin_viewer", "internal_user"], + ], +) +async def test_multi_valued_role_attribute_resolves_to_highest_privilege(saml_env_idp_initiated, roles): + """An assertion carrying several roles must not depend on the order the IdP emitted them in.""" + key_pem, cert_pem = saml_env_idp_initiated + resp = _build_signed_response( + key_pem, + cert_pem, + attributes={ + "email": ["dave@example.com"], + "role": roles, + }, + ) + + result = await _acs(_b64(resp), _shared_cache()) + assert result.user_role == LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY + + +@pytest.mark.asyncio +async def test_assertion_without_role_attribute_has_no_user_role(saml_env_idp_initiated): + key_pem, cert_pem = saml_env_idp_initiated + resp = _build_signed_response(key_pem, cert_pem, attributes={"email": ["erin@example.com"]}) + + result = await _acs(_b64(resp), _shared_cache()) + assert result.user_role is None + + @pytest.mark.asyncio async def test_build_login_redirect_targets_idp_and_caches_request_id(saml_env): cache = DualCache() 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 dd8c752a868..5dfff53f7c3 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -6598,13 +6598,94 @@ def test_get_litellm_user_role_with_invalid_role(): assert result is None -def test_get_litellm_user_role_with_list_multiple_roles(): - """Test that get_litellm_user_role takes the first element from a multi-element list.""" +@pytest.mark.parametrize( + "role_claim", + [ + ["proxy_admin", "internal_user"], + ["internal_user", "proxy_admin"], + ], +) +def test_get_litellm_user_role_picks_highest_privilege_regardless_of_order(role_claim): + """A multi-valued role claim resolves to the most privileged role, not the first one listed.""" from litellm.proxy._types import LitellmUserRoles from litellm.proxy.management_endpoints.types import get_litellm_user_role - result = get_litellm_user_role(["proxy_admin", "internal_user"]) - assert result == LitellmUserRoles.PROXY_ADMIN + assert get_litellm_user_role(role_claim) == LitellmUserRoles.PROXY_ADMIN + + +@pytest.mark.parametrize( + "role_claim", + [ + ["proxy_admin_viewer", "internal_user"], + ["internal_user", "proxy_admin_viewer"], + ], +) +def test_get_litellm_user_role_keeps_org_spend_visibility_for_mixed_roles(role_claim): + """ + Regression for LIT-6077: a user holding both proxy_admin_viewer and internal_user kept + losing org-level spend visibility whenever the IdP happened to list internal_user first. + """ + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.management_endpoints.types import get_litellm_user_role + + assert get_litellm_user_role(role_claim) == LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY + + +def test_get_litellm_user_role_ignores_unrecognised_entries(): + """Roles LiteLLM does not know about are skipped rather than swallowing the whole claim.""" + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.management_endpoints.types import get_litellm_user_role + + assert get_litellm_user_role(["some_idp_group", "internal_user"]) == LitellmUserRoles.INTERNAL_USER + assert get_litellm_user_role(["some_idp_group", "another_group"]) is None + + +def test_get_litellm_user_role_list_lookup_is_case_insensitive(): + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.management_endpoints.types import get_litellm_user_role + + assert get_litellm_user_role(["INTERNAL_USER", "Proxy_Admin"]) == LitellmUserRoles.PROXY_ADMIN + + +@pytest.mark.parametrize( + "role_claim", + [ + ["org_admin", "team"], + ["team", "org_admin"], + ], +) +def test_get_litellm_user_role_is_deterministic_for_unranked_roles(role_claim): + """Roles outside the privilege hierarchy still resolve the same way in either claim order.""" + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.management_endpoints.types import get_litellm_user_role + + assert get_litellm_user_role(role_claim) == LitellmUserRoles.ORG_ADMIN + + +@pytest.mark.parametrize( + "role_claim", + [ + ["org_admin", "internal_user"], + ["internal_user", "org_admin"], + ], +) +def test_get_litellm_user_role_prefers_a_ranked_role_over_an_unranked_one(role_claim): + """ + org_admin, team and customer sit outside the privilege ladder, so a claim mixing one of + them with a ranked role settles on the ranked role in either order. Same rule the Entra + app_roles and role_mappings paths already follow. + """ + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.management_endpoints.types import get_litellm_user_role + + assert get_litellm_user_role(role_claim) == LitellmUserRoles.INTERNAL_USER + + +def test_get_litellm_user_role_returns_none_for_non_string_claims(): + from litellm.proxy.management_endpoints.types import get_litellm_user_role + + assert get_litellm_user_role(None) is None + assert get_litellm_user_role({"role": "proxy_admin"}) is None # ============================================================================ @@ -6654,6 +6735,46 @@ def test_process_sso_jwt_access_token_extracts_role_from_access_token(): assert result.user_role == LitellmUserRoles.PROXY_ADMIN +@pytest.mark.parametrize( + "role_claim", + [ + ["internal_user", "proxy_admin_viewer"], + ["proxy_admin_viewer", "internal_user"], + ], +) +def test_process_sso_jwt_access_token_resolves_highest_privilege_role(role_claim): + """ + The generic SSO access-token path must land on the same role for a user whose role + claim holds several roles, whichever order the IdP emitted them in. + """ + import jwt as pyjwt + + from litellm.proxy._types import LitellmUserRoles + + access_token_str = pyjwt.encode( + {"sub": "user-123", "email": "mixed@test.com", "litellm_role": role_claim}, + "secret", + algorithm="HS256", + ) + result = CustomOpenID( + id="user-123", + email="mixed@test.com", + display_name="Mixed Role User", + team_ids=[], + user_role=None, + ) + + with patch.dict(os.environ, {"GENERIC_USER_ROLE_ATTRIBUTE": "litellm_role"}): + process_sso_jwt_access_token( + access_token_str=access_token_str, + sso_jwt_handler=None, + result=result, + role_mappings=None, + ) + + assert result.user_role == LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY + + def test_process_sso_jwt_access_token_does_not_override_existing_role(): """ Test that process_sso_jwt_access_token does NOT override a role that was diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 5c25b8722ef..d5bf3883be4 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -27,10 +27,10 @@ "limit": 0 }, "LIT010": { - "limit": 16478 + "limit": 16477 }, "LIT011": { - "limit": 5520 + "limit": 5519 }, "LIT012": { "limit": 4489