mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
feat(auth): extend JWT routing overrides with scope and wildcard selectors
Support optional scope matching and shell-style wildcard selectors (*, ?) for JWT OAuth2 routing overrides. Space-delimited tokenization applies only to the scope claim; iss/aud/client_id keep full-string matching on unverified claims. Document case-sensitive wildcard semantics. Add parametrized and integration tests for matcher, override composition, and routing behavior. Made-with: Cursor
This commit is contained in:
Milan
parent
850fe595ac
commit
e5a3e9c51c
3 changed files with 376 additions and 4 deletions
|
|
@ -4152,10 +4152,16 @@ class JWTRoutingOverride(BaseModel):
|
|||
|
||||
A rule matches when all provided selectors match token claims.
|
||||
If matched, request is routed to the configured auth path.
|
||||
|
||||
Wildcard selectors use shell-style patterns (* and ?) and are matched with
|
||||
case-sensitive semantics; use the same casing your IdP emits in JWT claims.
|
||||
Space-delimited tokenization applies only to the ``scope`` claim (OAuth/OIDC
|
||||
scope strings), not to ``iss``, ``aud``, or ``client_id``.
|
||||
"""
|
||||
|
||||
iss: Union[str, List[str]]
|
||||
client_id: Optional[Union[str, List[str]]] = None
|
||||
scope: Optional[Union[str, List[str]]] = None
|
||||
aud: Optional[Union[str, List[str]]] = None
|
||||
path: Literal["oauth2"] = "oauth2"
|
||||
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ Returns a UserAPIKeyAuth object if the API key is valid
|
|||
"""
|
||||
|
||||
import asyncio
|
||||
import fnmatch
|
||||
import re
|
||||
import secrets
|
||||
from datetime import datetime, timezone
|
||||
|
|
@ -140,22 +141,49 @@ def _get_bearer_token_or_received_api_key(api_key: str) -> str:
|
|||
|
||||
|
||||
def _routing_selector_matches_claim(
|
||||
selector_value: Optional[Any], claim_value: Optional[Any]
|
||||
selector_value: Optional[Any],
|
||||
claim_value: Optional[Any],
|
||||
*,
|
||||
split_space_delimited: bool = False,
|
||||
) -> bool:
|
||||
if selector_value is None:
|
||||
return True
|
||||
|
||||
selector_list = (
|
||||
selector_list: List[str] = (
|
||||
[str(v) for v in selector_value]
|
||||
if isinstance(selector_value, list)
|
||||
else [str(selector_value)]
|
||||
)
|
||||
|
||||
if claim_value is None:
|
||||
return False
|
||||
|
||||
if isinstance(claim_value, list):
|
||||
claim_list = [str(v) for v in claim_value]
|
||||
return any(v in claim_list for v in selector_list)
|
||||
elif (
|
||||
split_space_delimited
|
||||
and isinstance(claim_value, str)
|
||||
and " " in claim_value.strip()
|
||||
):
|
||||
# OAuth/OIDC often sends scope as a single space-delimited string. Only split
|
||||
# for the scope selector: iss/aud/client_id must stay exact full-string match
|
||||
# on unverified claims (see routing override security review).
|
||||
split_values = [v for v in claim_value.strip().split(" ") if v]
|
||||
claim_list = split_values if len(split_values) > 1 else [claim_value]
|
||||
else:
|
||||
claim_list = [str(claim_value)]
|
||||
|
||||
return str(claim_value) in selector_list if claim_value is not None else False
|
||||
def _selector_matches_claim(selector: str, claim: str) -> bool:
|
||||
# NOTE: wildcard matching is case-sensitive (fnmatch.fnmatchcase).
|
||||
if "*" in selector or "?" in selector:
|
||||
return fnmatch.fnmatchcase(claim, selector)
|
||||
return selector == claim
|
||||
|
||||
return any(
|
||||
_selector_matches_claim(selector=s, claim=c)
|
||||
for s in selector_list
|
||||
for c in claim_list
|
||||
)
|
||||
|
||||
|
||||
def _matches_routing_override(
|
||||
|
|
@ -166,6 +194,11 @@ def _matches_routing_override(
|
|||
and _routing_selector_matches_claim(
|
||||
override.client_id, token_claims.get("client_id")
|
||||
)
|
||||
and _routing_selector_matches_claim(
|
||||
override.scope,
|
||||
token_claims.get("scope"),
|
||||
split_space_delimited=True,
|
||||
)
|
||||
and _routing_selector_matches_claim(override.aud, token_claims.get("aud"))
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -23,6 +23,8 @@ from litellm.proxy._types import (
|
|||
from litellm.proxy.auth.handle_jwt import JWTHandler
|
||||
from litellm.proxy.auth.route_checks import RouteChecks
|
||||
from litellm.proxy.auth.user_api_key_auth import (
|
||||
_matches_routing_override,
|
||||
_routing_selector_matches_claim,
|
||||
_run_post_custom_auth_checks,
|
||||
get_api_key,
|
||||
user_api_key_auth,
|
||||
|
|
@ -431,6 +433,140 @@ def test_get_api_key_with_custom_litellm_key_header(
|
|||
) == (api_key, passed_in_key)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"selector_value, claim_value, expected, split_space_delimited",
|
||||
[
|
||||
(None, "any-value", True, False),
|
||||
("issuer.example.com", "issuer.example.com", True, False),
|
||||
("issuer.example.com", "other-issuer.example.com", False, False),
|
||||
# iss (and other non-scope claims) must not match via space-split injection
|
||||
(
|
||||
"trusted.example.com",
|
||||
"trusted.example.com attacker.example.com",
|
||||
False,
|
||||
False,
|
||||
),
|
||||
(
|
||||
["issuer-a.example.com", "issuer-b.example.com"],
|
||||
"issuer-b.example.com",
|
||||
True,
|
||||
False,
|
||||
),
|
||||
("*MID_LITELLM", "STREAM_MID_LITELLM", True, False),
|
||||
("*MID_LITELLM", "REDIS_LITELLM", False, False),
|
||||
("machine-??", "machine-01", True, False),
|
||||
("machine-??", "machine-001", False, False),
|
||||
# Wildcard matching is case-sensitive (fnmatch.fnmatchcase)
|
||||
("*litellm", "BATCH_LITELLM", False, False),
|
||||
("*LITELLM", "BATCH_LITELLM", True, False),
|
||||
("App:LiteLLM", "App:LiteLLM openid", True, True),
|
||||
("App:*", "App:LiteLLM openid", True, True),
|
||||
(["openid", "App:LiteLLM"], "openid profile", True, True),
|
||||
(["service-*", "batch-*"], "batch-123", True, False),
|
||||
(["service-*", "batch-*"], "other-123", False, False),
|
||||
("App:LiteLLM", ["openid", "App:LiteLLM"], True, False),
|
||||
("App:LiteLLM", None, False, False),
|
||||
],
|
||||
)
|
||||
def test_routing_selector_matches_claim_parametrized(
|
||||
selector_value, claim_value, expected, split_space_delimited
|
||||
):
|
||||
assert (
|
||||
_routing_selector_matches_claim(
|
||||
selector_value=selector_value,
|
||||
claim_value=claim_value,
|
||||
split_space_delimited=split_space_delimited,
|
||||
)
|
||||
is expected
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"override, token_claims, expected",
|
||||
[
|
||||
# Only iss selector is required and should match.
|
||||
(
|
||||
JWTRoutingOverride(iss="oauth-issuer.example.com", path="oauth2"),
|
||||
{"iss": "oauth-issuer.example.com"},
|
||||
True,
|
||||
),
|
||||
# Scope selector narrows the match.
|
||||
(
|
||||
JWTRoutingOverride(
|
||||
iss="oauth-issuer.example.com",
|
||||
scope="App:LiteLLM",
|
||||
path="oauth2",
|
||||
),
|
||||
{"iss": "oauth-issuer.example.com", "scope": "App:LiteLLM openid"},
|
||||
True,
|
||||
),
|
||||
# client_id wildcard selector narrows the match.
|
||||
(
|
||||
JWTRoutingOverride(
|
||||
iss="oauth-issuer.example.com",
|
||||
client_id="*MID_LITELLM",
|
||||
path="oauth2",
|
||||
),
|
||||
{"iss": "oauth-issuer.example.com", "client_id": "BATCH_MID_LITELLM"},
|
||||
True,
|
||||
),
|
||||
(
|
||||
JWTRoutingOverride(
|
||||
iss="oauth-issuer.example.com",
|
||||
client_id="*MID_LITELLM",
|
||||
path="oauth2",
|
||||
),
|
||||
{"iss": "oauth-issuer.example.com", "client_id": "BATCH_PORTAL"},
|
||||
False,
|
||||
),
|
||||
# aud selector still works with list claims.
|
||||
(
|
||||
JWTRoutingOverride(
|
||||
iss="oauth-issuer.example.com",
|
||||
aud=["api://litellm", "api://fallback"],
|
||||
path="oauth2",
|
||||
),
|
||||
{"iss": "oauth-issuer.example.com", "aud": ["api://other", "api://litellm"]},
|
||||
True,
|
||||
),
|
||||
# All provided selectors are AND-ed.
|
||||
(
|
||||
JWTRoutingOverride(
|
||||
iss="oauth-issuer.example.com",
|
||||
scope="App:LiteLLM",
|
||||
client_id="*MID_LITELLM",
|
||||
path="oauth2",
|
||||
),
|
||||
{
|
||||
"iss": "oauth-issuer.example.com",
|
||||
"scope": "App:LiteLLM openid",
|
||||
"client_id": "BATCH_MID_LITELLM",
|
||||
},
|
||||
True,
|
||||
),
|
||||
(
|
||||
JWTRoutingOverride(
|
||||
iss="oauth-issuer.example.com",
|
||||
scope="App:LiteLLM",
|
||||
client_id="*MID_LITELLM",
|
||||
path="oauth2",
|
||||
),
|
||||
{
|
||||
"iss": "oauth-issuer.example.com",
|
||||
"scope": "App:Other openid",
|
||||
"client_id": "BATCH_MID_LITELLM",
|
||||
},
|
||||
False,
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_matches_routing_override_parametrized(override, token_claims, expected):
|
||||
assert (
|
||||
_matches_routing_override(token_claims=token_claims, override=override)
|
||||
is expected
|
||||
)
|
||||
|
||||
|
||||
def test_team_metadata_with_tags_flows_through_jwt_auth():
|
||||
"""
|
||||
Test that team_metadata (specifically tags) flows through JWT authentication.
|
||||
|
|
@ -1293,6 +1429,203 @@ class TestJWTOAuth2Coexistence:
|
|||
mock_jwt_auth.assert_not_called()
|
||||
assert result.user_id == "machine-client-aud-list"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_routing_override_matches_scope_claim(self):
|
||||
"""
|
||||
Match routing override when scope selector is configured and scope claim matches.
|
||||
"""
|
||||
jwt_token = (
|
||||
"eyJhbGciOiJSUzI1NiJ9."
|
||||
"eyJpc3MiOiJvYXV0aC1pc3N1ZXIuZXhhbXBsZS5jb20iLCJzY29wZSI6IkFwcDpMaXRlTExNIiwiY2xpZW50X2lkIjoiTUFDSElORV9NSURfTElURUxMTSJ9."
|
||||
"c2ln"
|
||||
)
|
||||
general_settings = {
|
||||
"enable_oauth2_auth": False,
|
||||
"enable_jwt_auth": True,
|
||||
}
|
||||
mock_oauth2_response = UserAPIKeyAuth(
|
||||
api_key=jwt_token,
|
||||
user_id="machine-client-scope-match",
|
||||
)
|
||||
|
||||
mock_request = MagicMock()
|
||||
mock_request.url.path = "/v1/chat/completions"
|
||||
mock_request.headers = {"authorization": f"Bearer {jwt_token}"}
|
||||
mock_request.query_params = {}
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.proxy_server.general_settings", general_settings
|
||||
), patch("litellm.proxy.proxy_server.premium_user", True), patch(
|
||||
"litellm.proxy.proxy_server.master_key", "sk-master"
|
||||
), patch(
|
||||
"litellm.proxy.proxy_server.prisma_client", None
|
||||
), patch(
|
||||
"litellm.proxy.auth.user_api_key_auth.Oauth2Handler.check_oauth2_token",
|
||||
new_callable=AsyncMock,
|
||||
return_value=mock_oauth2_response,
|
||||
) as mock_oauth2, patch(
|
||||
"litellm.proxy.auth.user_api_key_auth.JWTAuthManager.auth_builder",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_jwt_auth:
|
||||
litellm.proxy.proxy_server.jwt_handler.update_environment(
|
||||
prisma_client=None,
|
||||
user_api_key_cache=DualCache(),
|
||||
litellm_jwtauth=LiteLLM_JWTAuth(
|
||||
routing_overrides=[
|
||||
JWTRoutingOverride(
|
||||
iss="oauth-issuer.example.com",
|
||||
scope="App:LiteLLM",
|
||||
path="oauth2",
|
||||
)
|
||||
]
|
||||
),
|
||||
)
|
||||
|
||||
result = await user_api_key_auth(
|
||||
request=mock_request,
|
||||
api_key=f"Bearer {jwt_token}",
|
||||
)
|
||||
|
||||
mock_oauth2.assert_called_once_with(token=jwt_token)
|
||||
mock_jwt_auth.assert_not_called()
|
||||
assert result.user_id == "machine-client-scope-match"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_routing_override_scope_mismatch_falls_back_to_jwt(self):
|
||||
"""
|
||||
If scope selector does not match, continue default JWT flow.
|
||||
"""
|
||||
jwt_token = (
|
||||
"eyJhbGciOiJSUzI1NiJ9."
|
||||
"eyJpc3MiOiJvYXV0aC1pc3N1ZXIuZXhhbXBsZS5jb20iLCJzY29wZSI6IkFwcDpPdGhlciIsImNsaWVudF9pZCI6IlBPUlRBTF9NSURfTElURUxMTSJ9."
|
||||
"c2ln"
|
||||
)
|
||||
general_settings = {
|
||||
"enable_oauth2_auth": False,
|
||||
"enable_jwt_auth": True,
|
||||
}
|
||||
mock_jwt_result = {
|
||||
"is_proxy_admin": True,
|
||||
"team_object": None,
|
||||
"user_object": None,
|
||||
"end_user_object": None,
|
||||
"org_object": None,
|
||||
"token": jwt_token,
|
||||
"team_id": "jwt-team",
|
||||
"user_id": "jwt-user-scope-mismatch",
|
||||
"end_user_id": None,
|
||||
"org_id": None,
|
||||
"team_membership": None,
|
||||
"jwt_claims": {"sub": "user1"},
|
||||
}
|
||||
|
||||
mock_request = MagicMock()
|
||||
mock_request.url.path = "/v1/chat/completions"
|
||||
mock_request.headers = {"authorization": f"Bearer {jwt_token}"}
|
||||
mock_request.query_params = {}
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.proxy_server.general_settings", general_settings
|
||||
), patch("litellm.proxy.proxy_server.premium_user", True), patch(
|
||||
"litellm.proxy.proxy_server.master_key", "sk-master"
|
||||
), patch(
|
||||
"litellm.proxy.proxy_server.prisma_client", None
|
||||
), patch(
|
||||
"litellm.proxy.auth.user_api_key_auth.Oauth2Handler.check_oauth2_token",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_oauth2, patch(
|
||||
"litellm.proxy.auth.user_api_key_auth.JWTAuthManager.auth_builder",
|
||||
new_callable=AsyncMock,
|
||||
return_value=mock_jwt_result,
|
||||
) as mock_jwt_auth:
|
||||
litellm.proxy.proxy_server.jwt_handler.update_environment(
|
||||
prisma_client=None,
|
||||
user_api_key_cache=DualCache(),
|
||||
litellm_jwtauth=LiteLLM_JWTAuth(
|
||||
routing_overrides=[
|
||||
JWTRoutingOverride(
|
||||
iss="oauth-issuer.example.com",
|
||||
scope="App:LiteLLM",
|
||||
path="oauth2",
|
||||
)
|
||||
]
|
||||
),
|
||||
)
|
||||
|
||||
result = await user_api_key_auth(
|
||||
request=mock_request,
|
||||
api_key=f"Bearer {jwt_token}",
|
||||
)
|
||||
|
||||
mock_oauth2.assert_not_called()
|
||||
mock_jwt_auth.assert_called_once()
|
||||
assert result.user_id == "jwt-user-scope-mismatch"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_routing_override_matches_scope_and_client_wildcard_when_scope_claim_is_space_delimited(
|
||||
self,
|
||||
):
|
||||
"""
|
||||
Integration check: combined scope + wildcard selectors match on OAuth2 path
|
||||
when scope claim is a space-delimited string.
|
||||
"""
|
||||
jwt_token = (
|
||||
"eyJhbGciOiJSUzI1NiJ9."
|
||||
"eyJpc3MiOiJvYXV0aC1pc3N1ZXIuZXhhbXBsZS5jb20iLCJzY29wZSI6IkFwcDpMaXRlTExNIG9wZW5pZCIsImNsaWVudF9pZCI6IkJBVENIX01JRF9MSVRFTExNIn0."
|
||||
"c2ln"
|
||||
)
|
||||
general_settings = {
|
||||
"enable_oauth2_auth": False,
|
||||
"enable_jwt_auth": True,
|
||||
}
|
||||
mock_oauth2_response = UserAPIKeyAuth(
|
||||
api_key=jwt_token,
|
||||
user_id="machine-client-space-delimited-scope-match",
|
||||
)
|
||||
|
||||
mock_request = MagicMock()
|
||||
mock_request.url.path = "/v1/chat/completions"
|
||||
mock_request.headers = {"authorization": f"Bearer {jwt_token}"}
|
||||
mock_request.query_params = {}
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.proxy_server.general_settings", general_settings
|
||||
), patch("litellm.proxy.proxy_server.premium_user", True), patch(
|
||||
"litellm.proxy.proxy_server.master_key", "sk-master"
|
||||
), patch(
|
||||
"litellm.proxy.proxy_server.prisma_client", None
|
||||
), patch(
|
||||
"litellm.proxy.auth.user_api_key_auth.Oauth2Handler.check_oauth2_token",
|
||||
new_callable=AsyncMock,
|
||||
return_value=mock_oauth2_response,
|
||||
) as mock_oauth2, patch(
|
||||
"litellm.proxy.auth.user_api_key_auth.JWTAuthManager.auth_builder",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_jwt_auth:
|
||||
litellm.proxy.proxy_server.jwt_handler.update_environment(
|
||||
prisma_client=None,
|
||||
user_api_key_cache=DualCache(),
|
||||
litellm_jwtauth=LiteLLM_JWTAuth(
|
||||
routing_overrides=[
|
||||
JWTRoutingOverride(
|
||||
iss="oauth-issuer.example.com",
|
||||
scope="App:LiteLLM",
|
||||
client_id="*MID_LITELLM",
|
||||
path="oauth2",
|
||||
)
|
||||
]
|
||||
),
|
||||
)
|
||||
|
||||
result = await user_api_key_auth(
|
||||
request=mock_request,
|
||||
api_key=f"Bearer {jwt_token}",
|
||||
)
|
||||
|
||||
mock_oauth2.assert_called_once_with(token=jwt_token)
|
||||
mock_jwt_auth.assert_not_called()
|
||||
assert result.user_id == "machine-client-space-delimited-scope-match"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_routing_override_routes_jwt_to_oauth2_when_oauth2_globally_disabled(
|
||||
self,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue