mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-04 07:27:51 +00:00
Merge e98a5367fe into 0a1dbc0fb4
This commit is contained in:
Vivek Raj
GitHub
commit
714f123e08
2 changed files with 131 additions and 53 deletions
|
|
@ -896,6 +896,29 @@ def determine_role_from_groups(
|
|||
return role_mappings.default_role
|
||||
|
||||
|
||||
def _groups_from_claim(value: object) -> list[str]:
|
||||
match value:
|
||||
case None:
|
||||
return []
|
||||
case list():
|
||||
groups: Final[Sequence[object]] = value
|
||||
return [str(group) for group in groups]
|
||||
case str():
|
||||
return [group.strip() for group in value.split(",") if group.strip()]
|
||||
case _:
|
||||
return [str(value)]
|
||||
|
||||
|
||||
def role_from_role_mappings(claims: Mapping[str, object], role_mappings: RoleMappings) -> LitellmUserRoles | None:
|
||||
user_groups: Final = _groups_from_claim(get_nested_value(claims, role_mappings.group_claim))
|
||||
if not user_groups:
|
||||
verbose_proxy_logger.debug(
|
||||
"No groups found in '%s', using default_role: %s", role_mappings.group_claim, role_mappings.default_role
|
||||
)
|
||||
return role_mappings.default_role
|
||||
return determine_role_from_groups(user_groups, role_mappings)
|
||||
|
||||
|
||||
def process_sso_jwt_access_token(
|
||||
access_token_str: str | None,
|
||||
sso_jwt_handler: JWTHandler | None,
|
||||
|
|
@ -953,25 +976,7 @@ def process_sso_jwt_access_token(
|
|||
|
||||
# Try role_mappings first (group-based role determination)
|
||||
if role_mappings is not None and role_mappings.roles:
|
||||
group_claim: Final = role_mappings.group_claim
|
||||
user_groups_raw: Final[object] = get_nested_value(access_token_payload, group_claim)
|
||||
|
||||
user_groups: list[str] = []
|
||||
if isinstance(user_groups_raw, list):
|
||||
raw_groups: Final[Sequence[object]] = user_groups_raw
|
||||
user_groups = [str(g) for g in raw_groups]
|
||||
elif isinstance(user_groups_raw, str):
|
||||
user_groups = [g.strip() for g in user_groups_raw.split(",") if g.strip()]
|
||||
elif user_groups_raw is not None:
|
||||
user_groups = [str(user_groups_raw)]
|
||||
|
||||
if user_groups:
|
||||
user_role = determine_role_from_groups(user_groups, role_mappings)
|
||||
verbose_proxy_logger.debug(
|
||||
"Determined role '%s' from access token groups '%s' using role_mappings", user_role, user_groups
|
||||
)
|
||||
elif role_mappings.default_role:
|
||||
user_role = role_mappings.default_role
|
||||
user_role = role_from_role_mappings(access_token_payload, role_mappings)
|
||||
|
||||
# Fallback: try GENERIC_USER_ROLE_ATTRIBUTE on the access token payload
|
||||
if user_role is None:
|
||||
|
|
@ -1268,35 +1273,7 @@ def generic_response_convertor(
|
|||
"generic",
|
||||
"okta",
|
||||
]:
|
||||
# Use role_mappings to determine role from groups
|
||||
group_claim: Final = role_mappings.group_claim
|
||||
user_groups_raw: Final[object] = get_nested_value(response, group_claim)
|
||||
|
||||
# Handle different formats: could be a list, string (comma-separated), or single value
|
||||
user_groups: list[str] = []
|
||||
if isinstance(user_groups_raw, list):
|
||||
raw_groups: Final[Sequence[object]] = user_groups_raw
|
||||
user_groups = [str(g) for g in raw_groups]
|
||||
elif isinstance(user_groups_raw, str):
|
||||
# Handle comma-separated string
|
||||
user_groups = [g.strip() for g in user_groups_raw.split(",") if g.strip()]
|
||||
elif user_groups_raw is not None:
|
||||
# Single value
|
||||
user_groups = [str(user_groups_raw)]
|
||||
|
||||
if user_groups:
|
||||
user_role = determine_role_from_groups(user_groups, role_mappings)
|
||||
verbose_proxy_logger.debug(
|
||||
"Determined role '%s' from groups '%s' using role_mappings",
|
||||
user_role.value if user_role else None,
|
||||
user_groups,
|
||||
)
|
||||
else:
|
||||
# No groups found, use default_role
|
||||
user_role = role_mappings.default_role
|
||||
verbose_proxy_logger.debug(
|
||||
"No groups found in '%s', using default_role: %s", group_claim, role_mappings.default_role
|
||||
)
|
||||
user_role = role_from_role_mappings(response, role_mappings)
|
||||
|
||||
# Fallback to existing logic if role_mappings not used
|
||||
if user_role is None:
|
||||
|
|
@ -2165,6 +2142,7 @@ async def _sso_provider_callback(
|
|||
microsoft_client_id=microsoft_client_id,
|
||||
redirect_url=redirect_url,
|
||||
return_raw_sso_response=return_raw_sso_response,
|
||||
role_mappings=await _setup_role_mappings(),
|
||||
)
|
||||
return microsoft_result, None, None, microsoft_assertion
|
||||
if generic_client_id is not None:
|
||||
|
|
@ -4575,6 +4553,7 @@ class MicrosoftSSOHandler:
|
|||
microsoft_client_id: str,
|
||||
redirect_url: str,
|
||||
return_raw_sso_response: bool = False,
|
||||
role_mappings: RoleMappings | None = None,
|
||||
) -> tuple[CustomOpenID | OpenID | dict[str, object], SSOIdentityAssertion | None]:
|
||||
"""
|
||||
Get the Microsoft SSO callback response, with the IdP identity assertion the token
|
||||
|
|
@ -4630,8 +4609,12 @@ class MicrosoftSSOHandler:
|
|||
app_roles: Final = MicrosoftSSOHandler.get_app_roles_from_id_token(id_token=microsoft_sso.id_token)
|
||||
verbose_proxy_logger.debug("Extracted app roles from id_token: %s", app_roles)
|
||||
|
||||
# Combine groups and app roles
|
||||
user_role: Final = MicrosoftSSOHandler.get_user_role_from_app_roles(app_roles)
|
||||
mapped_role: Final = (
|
||||
role_from_role_mappings(_decode_sso_token_claims(microsoft_tokens.id_token), role_mappings)
|
||||
if role_mappings is not None and role_mappings.provider.lower() == "microsoft"
|
||||
else None
|
||||
)
|
||||
user_role: Final = mapped_role or MicrosoftSSOHandler.get_user_role_from_app_roles(app_roles)
|
||||
|
||||
verbose_proxy_logger.debug("Combined team_ids (groups + app roles): %s", user_team_ids)
|
||||
|
||||
|
|
|
|||
|
|
@ -17,7 +17,7 @@ from pydantic import SecretStr
|
|||
import litellm
|
||||
from litellm._uuid import uuid
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
from litellm.proxy._types import LiteLLM_UserTable, NewUserResponse, ProxyException
|
||||
from litellm.proxy._types import LiteLLM_UserTable, LitellmUserRoles, NewUserResponse, ProxyException
|
||||
from litellm.proxy.auth.handle_jwt import JWTHandler
|
||||
from litellm.proxy.management_endpoints.sso import CustomMicrosoftSSO
|
||||
from litellm.proxy.management_endpoints.types import CustomOpenID
|
||||
|
|
@ -34,6 +34,7 @@ from litellm.proxy.management_endpoints.ui_sso import (
|
|||
normalize_email,
|
||||
process_sso_jwt_access_token,
|
||||
read_return_to_cookie,
|
||||
role_from_role_mappings,
|
||||
sso_lookup_user_ids,
|
||||
)
|
||||
from litellm.types.proxy.management_endpoints.ui_sso import (
|
||||
|
|
@ -41,6 +42,7 @@ from litellm.types.proxy.management_endpoints.ui_sso import (
|
|||
MicrosoftGraphAPIUserGroupDirectoryObject,
|
||||
MicrosoftGraphAPIUserGroupResponse,
|
||||
MicrosoftServicePrincipalTeam,
|
||||
RoleMappings,
|
||||
TeamMappings,
|
||||
)
|
||||
from tests._master_key import MASTER_KEY
|
||||
|
|
@ -410,6 +412,99 @@ def test_get_microsoft_callback_response_raw_sso_response(stubbed_graph_api):
|
|||
assert result["surname"] == "User"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("role_mappings", "id_token_claims", "expected_role"),
|
||||
[
|
||||
pytest.param(
|
||||
RoleMappings(provider="microsoft", group_claim="groups", roles={LitellmUserRoles.PROXY_ADMIN: ["grp-a"]}),
|
||||
{"groups": ["grp-b", "grp-a"]},
|
||||
LitellmUserRoles.PROXY_ADMIN,
|
||||
id="entra-group-mapped",
|
||||
),
|
||||
pytest.param(
|
||||
RoleMappings(
|
||||
provider="microsoft", group_claim="roles", roles={LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY: ["Auditor"]}
|
||||
),
|
||||
{"roles": ["Auditor"]},
|
||||
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
|
||||
id="entra-app-role-mapped",
|
||||
),
|
||||
pytest.param(
|
||||
RoleMappings(
|
||||
provider="microsoft",
|
||||
group_claim="groups",
|
||||
default_role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
|
||||
roles={LitellmUserRoles.PROXY_ADMIN: ["grp-a"]},
|
||||
),
|
||||
{},
|
||||
LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
|
||||
id="no-groups-uses-default-role",
|
||||
),
|
||||
pytest.param(
|
||||
RoleMappings(provider="microsoft", group_claim="groups", roles={LitellmUserRoles.PROXY_ADMIN: ["grp-a"]}),
|
||||
{"groups": ["grp-b"], "roles": ["internal_user"]},
|
||||
LitellmUserRoles.INTERNAL_USER,
|
||||
id="unmatched-falls-back-to-app-role",
|
||||
),
|
||||
pytest.param(
|
||||
RoleMappings(provider="generic", group_claim="groups", roles={LitellmUserRoles.PROXY_ADMIN: ["grp-a"]}),
|
||||
{"groups": ["grp-a"]},
|
||||
None,
|
||||
id="other-provider-mappings-ignored",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_get_microsoft_callback_response_applies_role_mappings(
|
||||
stubbed_graph_api, role_mappings, id_token_claims, expected_role
|
||||
):
|
||||
mock_request = MagicMock(spec=Request)
|
||||
mock_request.scope = {}
|
||||
id_token = _ema_id_token(sub="msft123", **id_token_claims)
|
||||
|
||||
async def fake_verify(self, *args, **kwargs):
|
||||
self._id_token = id_token
|
||||
return {"mail": "microsoft_user@example.com", "id": "msft123"}
|
||||
|
||||
with patch.dict(os.environ, {"MICROSOFT_CLIENT_SECRET": "mock_secret", "MICROSOFT_TENANT": "mock_tenant"}):
|
||||
with patch( # test-quality-ok: stubs the upstream Microsoft IdP boundary, where the token response is read
|
||||
"fastapi_sso.sso.microsoft.MicrosoftSSO.verify_and_process",
|
||||
autospec=True,
|
||||
side_effect=fake_verify,
|
||||
):
|
||||
result, _ = asyncio.run(
|
||||
MicrosoftSSOHandler.get_microsoft_callback_response(
|
||||
request=mock_request,
|
||||
microsoft_client_id="mock_client_id",
|
||||
redirect_url="http://mock_redirect_url",
|
||||
role_mappings=role_mappings,
|
||||
)
|
||||
)
|
||||
|
||||
assert isinstance(result, CustomOpenID)
|
||||
assert result.user_role == expected_role
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("claims", "expected_role"),
|
||||
[
|
||||
pytest.param({"groups": ["grp-b", "grp-a"]}, LitellmUserRoles.PROXY_ADMIN, id="list"),
|
||||
pytest.param({"groups": "grp-b, grp-a"}, LitellmUserRoles.PROXY_ADMIN, id="comma-separated-string"),
|
||||
pytest.param({"groups": 42}, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, id="single-non-string-value"),
|
||||
pytest.param({"groups": " , "}, LitellmUserRoles.INTERNAL_USER, id="blank-string-uses-default-role"),
|
||||
pytest.param({"groups": "grp-z"}, LitellmUserRoles.INTERNAL_USER, id="no-match-uses-default-role"),
|
||||
],
|
||||
)
|
||||
def test_role_from_role_mappings_reads_every_group_claim_shape(claims, expected_role):
|
||||
role_mappings = RoleMappings(
|
||||
provider="microsoft",
|
||||
group_claim="groups",
|
||||
default_role=LitellmUserRoles.INTERNAL_USER,
|
||||
roles={LitellmUserRoles.PROXY_ADMIN: ["grp-a"], LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY: ["42"]},
|
||||
)
|
||||
|
||||
assert role_from_role_mappings(claims, role_mappings) == expected_role
|
||||
|
||||
|
||||
def test_get_google_callback_response():
|
||||
# Arrange
|
||||
mock_request = MagicMock(spec=Request)
|
||||
|
|
@ -9311,13 +9406,13 @@ async def test_auth_callback_without_oauth_error_proceeds_to_normal_flow():
|
|||
# ── SSO identity assertion capture + persist wiring (EMA) ─────────────────────
|
||||
|
||||
|
||||
def _ema_id_token(sub: str = "u1") -> str:
|
||||
def _ema_id_token(sub: str = "u1", **claims: object) -> str:
|
||||
import time as _time
|
||||
|
||||
import jwt as _pyjwt
|
||||
|
||||
return _pyjwt.encode(
|
||||
{"iss": "https://idp.example.com", "sub": sub, "exp": int(_time.time()) + 3600},
|
||||
{"iss": "https://idp.example.com", "sub": sub, "exp": int(_time.time()) + 3600, **claims},
|
||||
"test-idp-signing-key-32-bytes-long-xxxx",
|
||||
algorithm="HS256",
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue