From 1f38359d95b65638b4b2f7dc41ba628756bd5bee Mon Sep 17 00:00:00 2001 From: Derek Higgins Date: Mon, 20 Oct 2025 20:34:55 +0100 Subject: [PATCH] fix: nested claims mapping in OAuth2 token validation (#3814) fix: nested claims mapping in OAuth2 token validation The get_attributes_from_claims function was only checking for top-level claim keys, causing token validation to fail when using nested claims like "resource_access.llamastack.roles" (common in Keycloak JWT tokens). Updated the function to support dot notation for traversing nested claim structures. Give precedence to dot notation over literal keys with dots in claims mapping. Added test coverage. Closes: #3812 Signed-off-by: Derek Higgins --- llama_stack/core/server/auth_providers.py | 23 ++++++- tests/unit/server/test_auth.py | 76 +++++++++++++++++++++++ 2 files changed, 96 insertions(+), 3 deletions(-) diff --git a/llama_stack/core/server/auth_providers.py b/llama_stack/core/server/auth_providers.py index 05a21c8d4..0fe5f1558 100644 --- a/llama_stack/core/server/auth_providers.py +++ b/llama_stack/core/server/auth_providers.py @@ -72,13 +72,30 @@ class AuthProvider(ABC): def get_attributes_from_claims(claims: dict[str, str], mapping: dict[str, str]) -> dict[str, list[str]]: attributes: dict[str, list[str]] = {} for claim_key, attribute_key in mapping.items(): - if claim_key not in claims: + # First try dot notation for nested traversal (e.g., "resource_access.llamastack.roles") + # Then fall back to literal key with dots (e.g., "my.dotted.key") + claim: object = claims + keys = claim_key.split(".") + for key in keys: + if isinstance(claim, dict) and key in claim: + claim = claim[key] + else: + claim = None + break + + if claim is None and claim_key in claims: + # Fall back to checking if claim_key exists as a literal key + claim = claims[claim_key] + + if claim is None: continue - claim = claims[claim_key] + if isinstance(claim, list): values = claim - else: + elif isinstance(claim, str): values = claim.split() + else: + continue if attribute_key in attributes: attributes[attribute_key].extend(values) diff --git a/tests/unit/server/test_auth.py b/tests/unit/server/test_auth.py index 04ae89db8..75cbf518b 100644 --- a/tests/unit/server/test_auth.py +++ b/tests/unit/server/test_auth.py @@ -516,6 +516,82 @@ def test_get_attributes_from_claims(): assert set(attributes["teams"]) == {"my-team", "group1", "group2"} assert attributes["namespaces"] == ["my-tenant"] + # Test nested claims with dot notation (e.g., Keycloak resource_access structure) + claims = { + "sub": "user123", + "resource_access": {"llamastack": {"roles": ["inference_max", "admin"]}, "other-client": {"roles": ["viewer"]}}, + "realm_access": {"roles": ["offline_access", "uma_authorization"]}, + } + attributes = get_attributes_from_claims( + claims, {"resource_access.llamastack.roles": "roles", "realm_access.roles": "realm_roles"} + ) + assert set(attributes["roles"]) == {"inference_max", "admin"} + assert set(attributes["realm_roles"]) == {"offline_access", "uma_authorization"} + + # Test that dot notation takes precedence over literal keys with dots + claims = { + "my.dotted.key": "literal-value", + "my": {"dotted": {"key": "nested-value"}}, + } + attributes = get_attributes_from_claims(claims, {"my.dotted.key": "test"}) + assert attributes["test"] == ["nested-value"] + + # Test that literal key works when nested traversal doesn't exist + claims = { + "my.dotted.key": "literal-value", + } + attributes = get_attributes_from_claims(claims, {"my.dotted.key": "test"}) + assert attributes["test"] == ["literal-value"] + + # Test missing nested paths are handled gracefully + claims = { + "sub": "user123", + "resource_access": {"other-client": {"roles": ["viewer"]}}, + } + attributes = get_attributes_from_claims( + claims, + { + "resource_access.llamastack.roles": "roles", # Missing nested path + "resource_access.missing.key": "missing_attr", # Missing nested path + "completely.missing.path": "another_missing", # Completely missing + "sub": "username", # Existing path + }, + ) + # Only the existing claim should be in attributes + assert attributes["username"] == ["user123"] + assert "roles" not in attributes + assert "missing_attr" not in attributes + assert "another_missing" not in attributes + + # Test mixture of flat and nested claims paths + claims = { + "sub": "user456", + "flat_key": "flat-value", + "scope": "read write admin", + "resource_access": {"app1": {"roles": ["role1", "role2"]}, "app2": {"roles": ["role3"]}}, + "groups": ["group1", "group2"], + "metadata": {"tenant": "tenant1", "region": "us-west"}, + } + attributes = get_attributes_from_claims( + claims, + { + "sub": "user_id", # Flat string + "scope": "permissions", # Flat string with spaces + "groups": "teams", # Flat list + "resource_access.app1.roles": "app1_roles", # Nested list + "resource_access.app2.roles": "app2_roles", # Nested list + "metadata.tenant": "tenant", # Nested string + "metadata.region": "region", # Nested string + }, + ) + assert attributes["user_id"] == ["user456"] + assert set(attributes["permissions"]) == {"read", "write", "admin"} + assert set(attributes["teams"]) == {"group1", "group2"} + assert set(attributes["app1_roles"]) == {"role1", "role2"} + assert attributes["app2_roles"] == ["role3"] + assert attributes["tenant"] == ["tenant1"] + assert attributes["region"] == ["us-west"] + # TODO: add more tests for oauth2 token provider