mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-19 04:20:26 +00:00
fix(jwt): OIDC discovery URLs, roles array handling, dot-notation error hints (#22336)
* fix(jwt): support OIDC discovery URLs, handle roles array, improve error hints Three fixes for Azure AD JWT auth: 1. OIDC discovery URL support - JWT_PUBLIC_KEY_URL can now be set to .well-known/openid-configuration endpoints. The proxy fetches the discovery doc, extracts jwks_uri, and caches it. 2. Handle roles claim as array - when team_id_jwt_field points to a list (e.g. AAD's "roles": ["team1"]), auto-unwrap the first element instead of crashing with 'unhashable type: list'. 3. Better error hint for dot-notation indexing - when team_id_jwt_field is set to "roles.0" or "roles[0]", the 401 error now explains to use "roles" instead and that LiteLLM auto-unwraps lists. * Add integration demo script for JWT auth fixes (OIDC discovery, array roles, dot-notation hints) Co-authored-by: Ishaan Jaff <ishaan-jaff@users.noreply.github.com> * Add demo_servers.py for manual JWT auth testing with mock JWKS/OIDC endpoints Co-authored-by: Ishaan Jaff <ishaan-jaff@users.noreply.github.com> * Add demo screenshots for PR comment Co-authored-by: Ishaan Jaff <ishaan-jaff@users.noreply.github.com> * Add integration test results with screenshots for PR review Co-authored-by: Ishaan Jaff <ishaan-jaff@users.noreply.github.com> * address greptile review feedback (greploop iteration 1) - fix: add HTTP status code check in _resolve_jwks_url before parsing JSON - fix: remove misleading bracket-notation hint from debug log (get_nested_value does not support it) * Update tests/test_litellm/proxy/auth/test_handle_jwt.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * remove demo scripts and assets --------- Co-authored-by: Cursor Agent <cursoragent@cursor.com> Co-authored-by: Ishaan Jaff <ishaan-jaff@users.noreply.github.com> Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
This commit is contained in:
Ishaan Jaff
co-authored by
Cursor Agent
Ishaan Jaff
greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Cursor Agent
Ishaan Jaff
greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
parent
8ab542875c
commit
ee703cea99
@@ -8,6 +8,7 @@ JWT token must have 'litellm_proxy_admin' in scope.
|
||||
|
||||
import fnmatch
|
||||
import os
|
||||
import re
|
||||
from typing import Any, List, Literal, Optional, Set, Tuple, cast
|
||||
|
||||
from cryptography import x509
|
||||
@@ -235,7 +236,17 @@ class JWTHandler:
|
||||
return self.litellm_jwtauth.team_id_default
|
||||
else:
|
||||
return default_value
|
||||
# At this point, team_id is not the sentinel, so it should be a string
|
||||
# AAD and other IdPs often send roles/groups as a list of strings.
|
||||
# team_id_jwt_field is singular, so take the first element when a list
|
||||
# is returned. This avoids "unhashable type: 'list'" errors downstream.
|
||||
if isinstance(team_id, list):
|
||||
if not team_id:
|
||||
return default_value
|
||||
verbose_proxy_logger.debug(
|
||||
f"JWT Auth: team_id_jwt_field '{self.litellm_jwtauth.team_id_jwt_field}' "
|
||||
f"returned a list {team_id}; using first element '{team_id[0]}' automatically."
|
||||
)
|
||||
team_id = team_id[0]
|
||||
return team_id # type: ignore[return-value]
|
||||
elif self.litellm_jwtauth.team_id_default is not None:
|
||||
team_id = self.litellm_jwtauth.team_id_default
|
||||
@@ -453,6 +464,52 @@ class JWTHandler:
|
||||
scopes = []
|
||||
return scopes
|
||||
|
||||
async def _resolve_jwks_url(self, url: str) -> str:
|
||||
"""
|
||||
If url points to an OIDC discovery document (*.well-known/openid-configuration),
|
||||
fetch it and return the jwks_uri contained within. Otherwise return url unchanged.
|
||||
This lets JWT_PUBLIC_KEY_URL be set to a well-known discovery endpoint instead of
|
||||
requiring operators to manually find the JWKS URL.
|
||||
"""
|
||||
if ".well-known/openid-configuration" not in url:
|
||||
return url
|
||||
|
||||
cache_key = f"litellm_oidc_discovery_{url}"
|
||||
cached_jwks_uri = await self.user_api_key_cache.async_get_cache(cache_key)
|
||||
if cached_jwks_uri is not None:
|
||||
return cached_jwks_uri
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
f"JWT Auth: Fetching OIDC discovery document from {url}"
|
||||
)
|
||||
response = await self.http_handler.get(url)
|
||||
if response.status_code != 200:
|
||||
raise Exception(
|
||||
f"JWT Auth: OIDC discovery endpoint {url} returned status {response.status_code}: {response.text}"
|
||||
)
|
||||
try:
|
||||
discovery = response.json()
|
||||
except Exception as e:
|
||||
raise Exception(
|
||||
f"JWT Auth: Failed to parse OIDC discovery document at {url}: {e}"
|
||||
)
|
||||
|
||||
jwks_uri = discovery.get("jwks_uri")
|
||||
if not jwks_uri:
|
||||
raise Exception(
|
||||
f"JWT Auth: OIDC discovery document at {url} does not contain a 'jwks_uri' field."
|
||||
)
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
f"JWT Auth: Resolved OIDC discovery {url} -> jwks_uri={jwks_uri}"
|
||||
)
|
||||
await self.user_api_key_cache.async_set_cache(
|
||||
key=cache_key,
|
||||
value=jwks_uri,
|
||||
ttl=self.litellm_jwtauth.public_key_ttl,
|
||||
)
|
||||
return jwks_uri
|
||||
|
||||
async def get_public_key(self, kid: Optional[str]) -> dict:
|
||||
keys_url = os.getenv("JWT_PUBLIC_KEY_URL")
|
||||
|
||||
@@ -462,6 +519,7 @@ class JWTHandler:
|
||||
keys_url_list = [url.strip() for url in keys_url.split(",")]
|
||||
|
||||
for key_url in keys_url_list:
|
||||
key_url = await self._resolve_jwks_url(key_url)
|
||||
cache_key = f"litellm_jwt_auth_keys_{key_url}"
|
||||
|
||||
cached_keys = await self.user_api_key_cache.async_get_cache(cache_key)
|
||||
@@ -913,8 +971,30 @@ class JWTAuthManager:
|
||||
if jwt_handler.is_required_team_id() is True:
|
||||
team_id_field = jwt_handler.litellm_jwtauth.team_id_jwt_field
|
||||
team_alias_field = jwt_handler.litellm_jwtauth.team_alias_jwt_field
|
||||
hint = ""
|
||||
if team_id_field:
|
||||
# "roles.0" — dot-notation numeric indexing is not supported
|
||||
if "." in team_id_field:
|
||||
parts = team_id_field.rsplit(".", 1)
|
||||
if parts[-1].isdigit():
|
||||
base_field = parts[0]
|
||||
hint = (
|
||||
f" Hint: dot-notation array indexing (e.g. '{team_id_field}') is not "
|
||||
f"supported. Use '{base_field}' instead — LiteLLM automatically "
|
||||
f"uses the first element when the field value is a list."
|
||||
)
|
||||
# "roles[0]" — bracket-notation indexing is also not supported in get_nested_value
|
||||
elif "[" in team_id_field and team_id_field.endswith("]"):
|
||||
m = re.match(r"^(\w+)\[(\d+)\]$", team_id_field)
|
||||
if m:
|
||||
base_field = m.group(1)
|
||||
hint = (
|
||||
f" Hint: array indexing (e.g. '{team_id_field}') is not supported "
|
||||
f"in team_id_jwt_field. Use '{base_field}' instead — LiteLLM "
|
||||
f"automatically uses the first element when the field value is a list."
|
||||
)
|
||||
raise Exception(
|
||||
f"No team found in token. Checked team_id field '{team_id_field}' and team_alias field '{team_alias_field}'"
|
||||
f"No team found in token. Checked team_id field '{team_id_field}' and team_alias field '{team_alias_field}'.{hint}"
|
||||
)
|
||||
|
||||
return individual_team_id, team_object
|
||||
|
||||
@@ -1485,4 +1485,273 @@ async def test_get_objects_resolves_org_by_name():
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fix 1: OIDC discovery URL resolution
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_jwks_url_passthrough_for_direct_jwks_url():
|
||||
"""Non-discovery URLs are returned unchanged."""
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
|
||||
handler = JWTHandler()
|
||||
handler.update_environment(
|
||||
prisma_client=None,
|
||||
user_api_key_cache=DualCache(),
|
||||
litellm_jwtauth=LiteLLM_JWTAuth(),
|
||||
)
|
||||
url = "https://login.microsoftonline.com/common/discovery/keys"
|
||||
result = await handler._resolve_jwks_url(url)
|
||||
assert result == url
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_jwks_url_resolves_oidc_discovery_document():
|
||||
"""
|
||||
A .well-known/openid-configuration URL should be fetched and its
|
||||
jwks_uri returned.
|
||||
"""
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
|
||||
handler = JWTHandler()
|
||||
cache = DualCache()
|
||||
handler.update_environment(
|
||||
prisma_client=None,
|
||||
user_api_key_cache=cache,
|
||||
litellm_jwtauth=LiteLLM_JWTAuth(),
|
||||
)
|
||||
|
||||
discovery_url = "https://login.microsoftonline.com/tenant/.well-known/openid-configuration"
|
||||
jwks_url = "https://login.microsoftonline.com/tenant/discovery/keys"
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {"jwks_uri": jwks_url, "issuer": "https://..."}
|
||||
|
||||
with patch.object(handler.http_handler, "get", new_callable=AsyncMock, return_value=mock_response) as mock_get:
|
||||
result = await handler._resolve_jwks_url(discovery_url)
|
||||
|
||||
assert result == jwks_url
|
||||
mock_get.assert_called_once_with(discovery_url)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_jwks_url_caches_resolved_jwks_uri():
|
||||
"""Resolved jwks_uri is cached — second call does not hit the network."""
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
|
||||
handler = JWTHandler()
|
||||
cache = DualCache()
|
||||
handler.update_environment(
|
||||
prisma_client=None,
|
||||
user_api_key_cache=cache,
|
||||
litellm_jwtauth=LiteLLM_JWTAuth(),
|
||||
)
|
||||
|
||||
discovery_url = "https://login.microsoftonline.com/tenant/.well-known/openid-configuration"
|
||||
jwks_url = "https://login.microsoftonline.com/tenant/discovery/keys"
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = {"jwks_uri": jwks_url}
|
||||
|
||||
with patch.object(handler.http_handler, "get", new_callable=AsyncMock, return_value=mock_response) as mock_get:
|
||||
first = await handler._resolve_jwks_url(discovery_url)
|
||||
second = await handler._resolve_jwks_url(discovery_url)
|
||||
|
||||
assert first == jwks_url
|
||||
assert second == jwks_url
|
||||
# Network should only be hit once
|
||||
assert mock_get.call_count == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_jwks_url_raises_if_no_jwks_uri_in_discovery_doc():
|
||||
"""Raise a helpful error if the discovery document has no jwks_uri."""
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
|
||||
handler = JWTHandler()
|
||||
handler.update_environment(
|
||||
prisma_client=None,
|
||||
user_api_key_cache=DualCache(),
|
||||
litellm_jwtauth=LiteLLM_JWTAuth(),
|
||||
)
|
||||
|
||||
discovery_url = "https://example.com/.well-known/openid-configuration"
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = {"issuer": "https://example.com"} # no jwks_uri
|
||||
|
||||
with patch.object(handler.http_handler, "get", new_callable=AsyncMock, return_value=mock_response):
|
||||
with pytest.raises(Exception, match="jwks_uri"):
|
||||
await handler._resolve_jwks_url(discovery_url)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fix 2: handle array values in team_id_jwt_field (e.g. AAD "roles" claim)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_jwt_handler(team_id_jwt_field: str) -> JWTHandler:
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
|
||||
handler = JWTHandler()
|
||||
handler.update_environment(
|
||||
prisma_client=None,
|
||||
user_api_key_cache=DualCache(),
|
||||
litellm_jwtauth=LiteLLM_JWTAuth(team_id_jwt_field=team_id_jwt_field),
|
||||
)
|
||||
return handler
|
||||
|
||||
|
||||
def test_get_team_id_returns_first_element_when_roles_is_list():
|
||||
"""
|
||||
AAD sends roles as a list. get_team_id() must return the first string
|
||||
element rather than the raw list (which would later crash with
|
||||
'unhashable type: list').
|
||||
"""
|
||||
handler = _make_jwt_handler("roles")
|
||||
token = {"oid": "user-oid", "roles": ["team1"]}
|
||||
result = handler.get_team_id(token=token, default_value=None)
|
||||
assert result == "team1"
|
||||
|
||||
|
||||
def test_get_team_id_returns_first_element_from_multi_value_roles_list():
|
||||
"""When roles has multiple entries, the first one is used."""
|
||||
handler = _make_jwt_handler("roles")
|
||||
token = {"roles": ["team2", "team1"]}
|
||||
result = handler.get_team_id(token=token, default_value=None)
|
||||
assert result == "team2"
|
||||
|
||||
|
||||
def test_get_team_id_returns_default_when_roles_list_is_empty():
|
||||
"""Empty list should fall back to default_value."""
|
||||
handler = _make_jwt_handler("roles")
|
||||
token = {"roles": []}
|
||||
result = handler.get_team_id(token=token, default_value="fallback")
|
||||
assert result == "fallback"
|
||||
|
||||
|
||||
def test_get_team_id_still_works_with_string_value():
|
||||
"""String values (non-array) continue to work as before."""
|
||||
handler = _make_jwt_handler("appid")
|
||||
token = {"appid": "my-team-id"}
|
||||
result = handler.get_team_id(token=token, default_value=None)
|
||||
assert result == "my-team-id"
|
||||
|
||||
|
||||
def test_get_team_id_list_result_is_hashable():
|
||||
"""
|
||||
The value returned by get_team_id() must be hashable so it can be
|
||||
added to a set (the operation that previously crashed).
|
||||
"""
|
||||
handler = _make_jwt_handler("roles")
|
||||
token = {"roles": ["team1"]}
|
||||
result = handler.get_team_id(token=token, default_value=None)
|
||||
# This must not raise TypeError
|
||||
s: set = set()
|
||||
s.add(result)
|
||||
assert "team1" in s
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fix 3: helpful error message for dot-notation array indexing (roles.0)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_find_and_validate_specific_team_id_hints_bracket_notation():
|
||||
"""
|
||||
When team_id_jwt_field is set to 'roles.0' (unsupported dot-notation for
|
||||
array indexing) and no team is found, the exception message should suggest
|
||||
using 'roles' instead (and explain LiteLLM auto-unwraps list values).
|
||||
"""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
|
||||
handler = _make_jwt_handler("roles.0")
|
||||
# token has roles as a list — dot-notation won't find anything
|
||||
token = {"roles": ["team1"]}
|
||||
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
await JWTAuthManager.find_and_validate_specific_team_id(
|
||||
jwt_handler=handler,
|
||||
jwt_valid_token=token,
|
||||
prisma_client=None,
|
||||
user_api_key_cache=DualCache(),
|
||||
parent_otel_span=None,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
error_msg = str(exc_info.value)
|
||||
# Should mention the bad field name and suggest the fix
|
||||
assert "roles.0" in error_msg, f"Expected field name in: {error_msg}"
|
||||
assert "roles" in error_msg and "list" in error_msg, (
|
||||
f"Expected hint about using 'roles' instead: {error_msg}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_find_and_validate_specific_team_id_hints_bracket_index_notation():
|
||||
"""
|
||||
When team_id_jwt_field is set to 'roles[0]' (bracket indexing, also unsupported
|
||||
in get_nested_value) the error message should suggest using 'roles' instead.
|
||||
"""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
|
||||
handler = _make_jwt_handler("roles[0]")
|
||||
token = {"roles": ["team1"]}
|
||||
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
await JWTAuthManager.find_and_validate_specific_team_id(
|
||||
jwt_handler=handler,
|
||||
jwt_valid_token=token,
|
||||
prisma_client=None,
|
||||
user_api_key_cache=DualCache(),
|
||||
parent_otel_span=None,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
error_msg = str(exc_info.value)
|
||||
assert "roles[0]" in error_msg, f"Expected field name in: {error_msg}"
|
||||
assert "roles" in error_msg and "list" in error_msg, (
|
||||
f"Expected hint about using 'roles' instead: {error_msg}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_find_and_validate_specific_team_id_no_hint_for_valid_field():
|
||||
"""
|
||||
When team_id_jwt_field is a normal field name (no dot-notation) the
|
||||
error message should not contain a spurious bracket-notation hint.
|
||||
"""
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
|
||||
handler = _make_jwt_handler("appid")
|
||||
token = {} # no appid — triggers the "no team found" path
|
||||
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
await JWTAuthManager.find_and_validate_specific_team_id(
|
||||
jwt_handler=handler,
|
||||
jwt_valid_token=token,
|
||||
prisma_client=None,
|
||||
user_api_key_cache=DualCache(),
|
||||
parent_otel_span=None,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
error_msg = str(exc_info.value)
|
||||
assert "Hint" not in error_msg
|
||||
|
||||
|
||||
Reference in New Issue
Block a user