From 6b166e046dc4700c315b7d65f6c7dec8bc574876 Mon Sep 17 00:00:00 2001 From: berri-teddy Date: Wed, 15 Oct 2025 17:17:55 -0700 Subject: [PATCH 1/2] fix: correct EntraID app roles JWT claim extraction - Fix get_app_roles_from_id_token to use 'app_roles' claim instead of 'roles' - Add comprehensive unit tests for EntraID app roles functionality - Prevent regressions in Microsoft EntraID SSO authentication Resolves issue where EntraID app roles were not being extracted correctly from JWT tokens, causing authentication failures for users with assigned app roles in Microsoft EntraID. --- litellm/proxy/management_endpoints/ui_sso.py | 30 +-- .../test_entraid_app_roles.py | 215 ++++++++++++++++++ 2 files changed, 230 insertions(+), 15 deletions(-) create mode 100644 tests/test_litellm/proxy/management_endpoints/test_entraid_app_roles.py diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index d8f426e1ae6..54909b9712b 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -875,9 +875,9 @@ async def insert_sso_user( if user_defined_values.get("max_budget") is None: user_defined_values["max_budget"] = litellm.max_internal_user_budget if user_defined_values.get("budget_duration") is None: - user_defined_values["budget_duration"] = ( - litellm.internal_user_budget_duration - ) + user_defined_values[ + "budget_duration" + ] = litellm.internal_user_budget_duration if user_defined_values["user_role"] is None: user_defined_values["user_role"] = LitellmUserRoles.INTERNAL_USER_VIEW_ONLY @@ -1118,9 +1118,9 @@ class SSOAuthenticationHandler: generic_authorization_endpoint and "okta" in generic_authorization_endpoint ): - redirect_params["state"] = ( - uuid.uuid4().hex - ) # set state param for okta - required + redirect_params[ + "state" + ] = uuid.uuid4().hex # set state param for okta - required return redirect_params @@ -1699,9 +1699,9 @@ class MicrosoftSSOHandler: # if user is trying to get the raw sso response for debugging, return the raw sso response if return_raw_sso_response: - original_msft_result[MicrosoftSSOHandler.GRAPH_API_RESPONSE_KEY] = ( - user_team_ids - ) + original_msft_result[ + MicrosoftSSOHandler.GRAPH_API_RESPONSE_KEY + ] = user_team_ids original_msft_result["app_roles"] = app_roles return original_msft_result or {} @@ -1739,7 +1739,7 @@ class MicrosoftSSOHandler: Extract app roles from the Microsoft Entra ID (Azure AD) id_token JWT. App roles are assigned in the Azure AD Enterprise Application and appear - in the 'roles' claim of the id_token. + in the 'app_roles' claim of the id_token. Args: id_token (Optional[str]): The JWT id_token from Microsoft SSO @@ -1758,8 +1758,8 @@ class MicrosoftSSOHandler: # (signature is already verified by fastapi_sso) decoded_token = jwt.decode(id_token, options={"verify_signature": False}) - # Extract roles claim from the token - roles = decoded_token.get("roles", []) + # Extract app_roles claim from the token + roles = decoded_token.get("app_roles", []) if roles and isinstance(roles, list): verbose_proxy_logger.debug( @@ -1817,9 +1817,9 @@ class MicrosoftSSOHandler: # Fetch user membership from Microsoft Graph API all_group_ids = [] - next_link: Optional[str] = ( - MicrosoftSSOHandler.graph_api_user_groups_endpoint - ) + next_link: Optional[ + str + ] = MicrosoftSSOHandler.graph_api_user_groups_endpoint auth_headers = {"Authorization": f"Bearer {access_token}"} page_count = 0 diff --git a/tests/test_litellm/proxy/management_endpoints/test_entraid_app_roles.py b/tests/test_litellm/proxy/management_endpoints/test_entraid_app_roles.py new file mode 100644 index 00000000000..de984496824 --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/test_entraid_app_roles.py @@ -0,0 +1,215 @@ +""" +Unit tests for EntraID app roles JWT claim extraction. + +This module tests the get_app_roles_from_id_token method to ensure it correctly +extracts app roles from Microsoft EntraID JWT tokens and prevents regressions. +""" + +import pytest +from unittest.mock import patch +import jwt + +from litellm.proxy.management_endpoints.ui_sso import MicrosoftSSOHandler + + +class TestEntraIDAppRoles: + """Test EntraID app roles extraction from JWT tokens""" + + @pytest.fixture + def sample_jwt_token(self): + """Create a sample JWT token with app_roles claim""" + payload = { + "sub": "user123", + "email": "user@company.com", + "app_roles": ["proxy_admin"], + "aud": "litellm-app", + "iss": "https://login.microsoftonline.com/tenant-id/v2.0", + "exp": 9999999999, + } + return jwt.encode(payload, "secret", algorithm="HS256") + + @pytest.fixture + def sample_jwt_token_single_role(self): + """Create a sample JWT token with single app role""" + payload = { + "sub": "user456", + "email": "admin@company.com", + "app_roles": ["proxy_admin_viewer"], + "aud": "litellm-app", + "iss": "https://login.microsoftonline.com/tenant-id/v2.0", + "exp": 9999999999, + } + return jwt.encode(payload, "secret", algorithm="HS256") + + @pytest.fixture + def sample_jwt_token_no_roles(self): + """Create a sample JWT token without app_roles claim""" + payload = { + "sub": "user789", + "email": "user@company.com", + "aud": "litellm-app", + "iss": "https://login.microsoftonline.com/tenant-id/v2.0", + "exp": 9999999999, + } + return jwt.encode(payload, "secret", algorithm="HS256") + + @pytest.fixture + def sample_jwt_token_empty_roles(self): + """Create a sample JWT token with empty app_roles array""" + payload = { + "sub": "user000", + "email": "user@company.com", + "app_roles": [], + "aud": "litellm-app", + "iss": "https://login.microsoftonline.com/tenant-id/v2.0", + "exp": 9999999999, + } + return jwt.encode(payload, "secret", algorithm="HS256") + + def test_get_app_roles_from_id_token_single_role( + self, sample_jwt_token_single_role + ): + """Test extracting single app role from JWT token""" + # Act + result = MicrosoftSSOHandler.get_app_roles_from_id_token( + sample_jwt_token_single_role + ) + + # Assert + assert result == ["proxy_admin_viewer"] + assert len(result) == 1 + + def test_get_app_roles_from_id_token_no_roles_claim( + self, sample_jwt_token_no_roles + ): + """Test handling JWT token without app_roles claim""" + # Act + result = MicrosoftSSOHandler.get_app_roles_from_id_token( + sample_jwt_token_no_roles + ) + + # Assert + assert result == [] + assert len(result) == 0 + + def test_get_app_roles_from_id_token_empty_roles( + self, sample_jwt_token_empty_roles + ): + """Test handling JWT token with empty app_roles array""" + # Act + result = MicrosoftSSOHandler.get_app_roles_from_id_token( + sample_jwt_token_empty_roles + ) + + # Assert + assert result == [] + assert len(result) == 0 + + def test_get_app_roles_from_id_token_none_input(self): + """Test handling None input""" + # Act + result = MicrosoftSSOHandler.get_app_roles_from_id_token(None) + + # Assert + assert result == [] + assert len(result) == 0 + + def test_get_app_roles_from_id_token_empty_string(self): + """Test handling empty string input""" + # Act + result = MicrosoftSSOHandler.get_app_roles_from_id_token("") + + # Assert + assert result == [] + assert len(result) == 0 + + def test_get_app_roles_from_id_token_invalid_jwt(self): + """Test handling invalid JWT token""" + # Act + result = MicrosoftSSOHandler.get_app_roles_from_id_token("invalid.jwt.token") + + # Assert + assert result == [] + assert len(result) == 0 + + def test_get_app_roles_from_id_token_malformed_roles(self): + """Test handling JWT with malformed app_roles (not a list)""" + # Arrange + payload = { + "sub": "user123", + "app_roles": "not_a_list", # Should be a list + "exp": 9999999999, + } + malformed_token = jwt.encode(payload, "secret", algorithm="HS256") + + # Act + result = MicrosoftSSOHandler.get_app_roles_from_id_token(malformed_token) + + # Assert + assert result == [] + assert len(result) == 0 + + def test_get_app_roles_from_id_token_jwt_decode_exception(self): + """Test handling JWT decode exceptions gracefully""" + # Arrange + invalid_token = "completely.invalid.token" + + # Act + result = MicrosoftSSOHandler.get_app_roles_from_id_token(invalid_token) + + # Assert + assert result == [] + assert len(result) == 0 + + def test_get_app_roles_from_id_token_import_error(self): + """Test handling import error for jwt library""" + # Arrange + with patch( + "builtins.__import__", side_effect=ImportError("No module named 'jwt'") + ): + # Act + result = MicrosoftSSOHandler.get_app_roles_from_id_token("any.token") + + # Assert + assert result == [] + assert len(result) == 0 + + def test_get_app_roles_from_id_token_uses_correct_claim_name( + self, sample_jwt_token + ): + """Test that the method uses 'app_roles' claim, not 'roles' claim""" + # This test ensures we don't regress to the old bug where 'roles' was used + + # Arrange - Create a token with both claims to verify correct one is used + payload = { + "sub": "user123", + "roles": ["old_roles_claim"], # This should be ignored + "app_roles": ["proxy_admin"], # This should be used + "exp": 9999999999, + } + token_with_both_claims = jwt.encode(payload, "secret", algorithm="HS256") + + # Act + result = MicrosoftSSOHandler.get_app_roles_from_id_token(token_with_both_claims) + + # Assert + assert result == ["proxy_admin"] # Should use app_roles, not roles + assert "old_roles_claim" not in result + + def test_get_app_roles_from_id_token_case_sensitivity(self): + """Test that app roles are extracted as-is (case sensitive)""" + # Arrange + payload = { + "sub": "user123", + "app_roles": ["PROXY_ADMIN", "Internal_User"], # Mixed case + "exp": 9999999999, + } + mixed_case_token = jwt.encode(payload, "secret", algorithm="HS256") + + # Act + result = MicrosoftSSOHandler.get_app_roles_from_id_token(mixed_case_token) + + # Assert + assert result == ["PROXY_ADMIN", "Internal_User"] + assert "PROXY_ADMIN" in result + assert "Internal_User" in result From 1b55803c76fc74771c89b168a3f9e6699c6c2d00 Mon Sep 17 00:00:00 2001 From: berri-teddy Date: Wed, 15 Oct 2025 18:15:04 -0700 Subject: [PATCH 2/2] test: reduce EntraID app roles tests to essential scenarios - Keep only 2 focused tests: single role extraction and no roles claim - Remove complex fixtures and edge case tests - Maintain clean, maintainable test coverage --- .../test_entraid_app_roles.py | 214 +++--------------- 1 file changed, 27 insertions(+), 187 deletions(-) diff --git a/tests/test_litellm/proxy/management_endpoints/test_entraid_app_roles.py b/tests/test_litellm/proxy/management_endpoints/test_entraid_app_roles.py index de984496824..f6248d36628 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_entraid_app_roles.py +++ b/tests/test_litellm/proxy/management_endpoints/test_entraid_app_roles.py @@ -6,7 +6,6 @@ extracts app roles from Microsoft EntraID JWT tokens and prevents regressions. """ import pytest -from unittest.mock import patch import jwt from litellm.proxy.management_endpoints.ui_sso import MicrosoftSSOHandler @@ -15,201 +14,42 @@ from litellm.proxy.management_endpoints.ui_sso import MicrosoftSSOHandler class TestEntraIDAppRoles: """Test EntraID app roles extraction from JWT tokens""" - @pytest.fixture - def sample_jwt_token(self): - """Create a sample JWT token with app_roles claim""" + def test_get_app_roles_from_id_token_works_without_roles(self): + """Test that JWT token works fine without app_roles claim""" + # Arrange - Token without app_roles (normal user) payload = { "sub": "user123", "email": "user@company.com", + "aud": "litellm-app", + "iss": "https://login.microsoftonline.com/tenant-id/v2.0", + "exp": 9999999999, + } + no_roles_token = jwt.encode(payload, "secret", algorithm="HS256") + + # Act + result = MicrosoftSSOHandler.get_app_roles_from_id_token(no_roles_token) + + # Assert - Should return empty list, not error + assert result == [] + assert len(result) == 0 + + def test_get_app_roles_from_id_token_assigns_roles_when_present(self): + """Test that valid app roles are properly assigned when present""" + # Arrange - Token with valid roles + payload = { + "sub": "user123", + "email": "admin@company.com", "app_roles": ["proxy_admin"], "aud": "litellm-app", "iss": "https://login.microsoftonline.com/tenant-id/v2.0", "exp": 9999999999, } - return jwt.encode(payload, "secret", algorithm="HS256") + valid_roles_token = jwt.encode(payload, "secret", algorithm="HS256") - @pytest.fixture - def sample_jwt_token_single_role(self): - """Create a sample JWT token with single app role""" - payload = { - "sub": "user456", - "email": "admin@company.com", - "app_roles": ["proxy_admin_viewer"], - "aud": "litellm-app", - "iss": "https://login.microsoftonline.com/tenant-id/v2.0", - "exp": 9999999999, - } - return jwt.encode(payload, "secret", algorithm="HS256") - - @pytest.fixture - def sample_jwt_token_no_roles(self): - """Create a sample JWT token without app_roles claim""" - payload = { - "sub": "user789", - "email": "user@company.com", - "aud": "litellm-app", - "iss": "https://login.microsoftonline.com/tenant-id/v2.0", - "exp": 9999999999, - } - return jwt.encode(payload, "secret", algorithm="HS256") - - @pytest.fixture - def sample_jwt_token_empty_roles(self): - """Create a sample JWT token with empty app_roles array""" - payload = { - "sub": "user000", - "email": "user@company.com", - "app_roles": [], - "aud": "litellm-app", - "iss": "https://login.microsoftonline.com/tenant-id/v2.0", - "exp": 9999999999, - } - return jwt.encode(payload, "secret", algorithm="HS256") - - def test_get_app_roles_from_id_token_single_role( - self, sample_jwt_token_single_role - ): - """Test extracting single app role from JWT token""" # Act - result = MicrosoftSSOHandler.get_app_roles_from_id_token( - sample_jwt_token_single_role - ) + result = MicrosoftSSOHandler.get_app_roles_from_id_token(valid_roles_token) - # Assert - assert result == ["proxy_admin_viewer"] + # Assert - Should extract the role + assert result == ["proxy_admin"] assert len(result) == 1 - - def test_get_app_roles_from_id_token_no_roles_claim( - self, sample_jwt_token_no_roles - ): - """Test handling JWT token without app_roles claim""" - # Act - result = MicrosoftSSOHandler.get_app_roles_from_id_token( - sample_jwt_token_no_roles - ) - - # Assert - assert result == [] - assert len(result) == 0 - - def test_get_app_roles_from_id_token_empty_roles( - self, sample_jwt_token_empty_roles - ): - """Test handling JWT token with empty app_roles array""" - # Act - result = MicrosoftSSOHandler.get_app_roles_from_id_token( - sample_jwt_token_empty_roles - ) - - # Assert - assert result == [] - assert len(result) == 0 - - def test_get_app_roles_from_id_token_none_input(self): - """Test handling None input""" - # Act - result = MicrosoftSSOHandler.get_app_roles_from_id_token(None) - - # Assert - assert result == [] - assert len(result) == 0 - - def test_get_app_roles_from_id_token_empty_string(self): - """Test handling empty string input""" - # Act - result = MicrosoftSSOHandler.get_app_roles_from_id_token("") - - # Assert - assert result == [] - assert len(result) == 0 - - def test_get_app_roles_from_id_token_invalid_jwt(self): - """Test handling invalid JWT token""" - # Act - result = MicrosoftSSOHandler.get_app_roles_from_id_token("invalid.jwt.token") - - # Assert - assert result == [] - assert len(result) == 0 - - def test_get_app_roles_from_id_token_malformed_roles(self): - """Test handling JWT with malformed app_roles (not a list)""" - # Arrange - payload = { - "sub": "user123", - "app_roles": "not_a_list", # Should be a list - "exp": 9999999999, - } - malformed_token = jwt.encode(payload, "secret", algorithm="HS256") - - # Act - result = MicrosoftSSOHandler.get_app_roles_from_id_token(malformed_token) - - # Assert - assert result == [] - assert len(result) == 0 - - def test_get_app_roles_from_id_token_jwt_decode_exception(self): - """Test handling JWT decode exceptions gracefully""" - # Arrange - invalid_token = "completely.invalid.token" - - # Act - result = MicrosoftSSOHandler.get_app_roles_from_id_token(invalid_token) - - # Assert - assert result == [] - assert len(result) == 0 - - def test_get_app_roles_from_id_token_import_error(self): - """Test handling import error for jwt library""" - # Arrange - with patch( - "builtins.__import__", side_effect=ImportError("No module named 'jwt'") - ): - # Act - result = MicrosoftSSOHandler.get_app_roles_from_id_token("any.token") - - # Assert - assert result == [] - assert len(result) == 0 - - def test_get_app_roles_from_id_token_uses_correct_claim_name( - self, sample_jwt_token - ): - """Test that the method uses 'app_roles' claim, not 'roles' claim""" - # This test ensures we don't regress to the old bug where 'roles' was used - - # Arrange - Create a token with both claims to verify correct one is used - payload = { - "sub": "user123", - "roles": ["old_roles_claim"], # This should be ignored - "app_roles": ["proxy_admin"], # This should be used - "exp": 9999999999, - } - token_with_both_claims = jwt.encode(payload, "secret", algorithm="HS256") - - # Act - result = MicrosoftSSOHandler.get_app_roles_from_id_token(token_with_both_claims) - - # Assert - assert result == ["proxy_admin"] # Should use app_roles, not roles - assert "old_roles_claim" not in result - - def test_get_app_roles_from_id_token_case_sensitivity(self): - """Test that app roles are extracted as-is (case sensitive)""" - # Arrange - payload = { - "sub": "user123", - "app_roles": ["PROXY_ADMIN", "Internal_User"], # Mixed case - "exp": 9999999999, - } - mixed_case_token = jwt.encode(payload, "secret", algorithm="HS256") - - # Act - result = MicrosoftSSOHandler.get_app_roles_from_id_token(mixed_case_token) - - # Assert - assert result == ["PROXY_ADMIN", "Internal_User"] - assert "PROXY_ADMIN" in result - assert "Internal_User" in result + assert "proxy_admin" in result