From 7fa438df06817eacf298f96cb4636b9c826de79a Mon Sep 17 00:00:00 2001 From: mrinal-berri Date: Tue, 3 Nov 2026 16:24:51 -0800 Subject: [PATCH] feat(auth): add link-only JWT mapping and remove route eligibility filtering (#49956) * feat(auth): allow skipping JWT key route eligibility check * feat(auth): support link-only JWT key mapping * fix(auth): model JWT mapping rejection as a failure value * feat(auth): select existing JWT keys without route filtering --- litellm/proxy/_types.py | 16 +- litellm/proxy/auth/user_api_key_auth.py | 75 ++-- .../verification_token_repository.py | 10 +- ...test_jwt_auto_register_map_existing_key.py | 349 +++++++++++++++++- 4 files changed, 406 insertions(+), 44 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index fc1e7128e10..2ef4c0089f9 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -4796,11 +4796,14 @@ class UnregisteredJWTClientBehavior(str, enum.Enum): - auto_register: Automatically create a new virtual key and mapping on first encounter. The new key has no budget/model restrictions; admins can tighten it later via /jwt_client/update. + - link_only: Map an eligible existing user key or return HTTP 403, without + creating a key. Proxy admins authenticate without a mapped key. """ FALLBACK_TEAM_MAPPING = "fallback_team_mapping" REJECT = "reject" AUTO_REGISTER = "auto_register" + LINK_ONLY = "link_only" class JWTIssuerConfig(LiteLLMBaseModel): @@ -5038,7 +5041,12 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase): "What to do when virtual_key_claim_field is set but the JWT claim value " "has no registered mapping. 'fallback_team_mapping' (default): fall through " "to team-based JWT auth. 'reject': return HTTP 403. " - "'auto_register': auto-create a virtual key and mapping on first encounter." + "'auto_register': auto-create a virtual key and mapping on first encounter. " + "'link_only': map an eligible existing user key or return HTTP 403, without creating a key. " + "link_only ignores auto_register_map_existing_key; user provisioning still follows user_id_upsert. " + "Existing-key selection ignores allowed_routes; JWTs inherit the selected key's permissions, " + "including management routes, with request-time key restrictions still enforced. " + "Unmapped proxy admins authenticate directly in auto_register and link_only modes." ), ) auto_register_map_existing_key: bool = Field( @@ -5048,8 +5056,10 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase): "field is the user_id_jwt_field or user_email_jwt_field, the JWT claim is mapped to a virtual key the " "JWT-resolved user already owns instead of minting a new one. If the user owns several, the most recently created key in the " "JWT-resolved team (or with no team when the JWT resolves none) is chosen among keys that never " - "expire, are not blocked, are not Admin UI session keys, were not minted by auto_register, and " - "have no allowed_routes or include llm_api_routes. Otherwise a new key is minted as usual." + "expire, are not blocked, are not Admin UI session keys, and were not minted by auto_register. " + "allowed_routes does not affect key selection. JWTs inherit the selected key's permissions, " + "including management routes; request-time key restrictions still apply. " + "Otherwise a new key is minted as usual." ), ) routing_overrides: list[JWTRoutingOverride] | None = Field( diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 1dd0647d25f..82135cb1b74 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -13,6 +13,7 @@ import re import secrets from collections.abc import Mapping, Sequence from datetime import datetime, timezone +from enum import Enum from types import MappingProxyType from typing import Final, NamedTuple, Protocol, Union, cast @@ -1576,10 +1577,14 @@ _JWT_AUTH_DISABLED_HINT = ( ) +class _JWTKeyMappingFailure(Enum): + NO_ELIGIBLE_KEY = "JWT Key Mapping: No eligible existing virtual key for this identity. Access denied." + + class _PendingAutoRegister(NamedTuple): """ Signal returned by ``_resolve_jwt_to_virtual_key`` when the JWT's claim is - unmapped and ``unregistered_jwt_client_behavior`` is AUTO_REGISTER. + unmapped and ``unregistered_jwt_client_behavior`` is AUTO_REGISTER or LINK_ONLY. The caller MUST run standard ``JWTAuthManager.auth_builder`` to apply RBAC, scope mappings, ``custom_validate``, and ``user_allowed_email_domain`` @@ -1593,23 +1598,22 @@ class _PendingAutoRegister(NamedTuple): claim_value: str cache_key: str jwt_issuer: str | None = None + behavior: UnregisteredJWTClientBehavior = UnregisteredJWTClientBehavior.AUTO_REGISTER def _claim_identifies_user(jwt_handler: JWTHandler, claim_field: str, jwt_issuer: str | None) -> bool: - if not jwt_handler.litellm_jwtauth.auto_register_map_existing_key: - return False if jwt_handler.litellm_jwtauth.is_user_identity_claim(claim_field, jwt_issuer): return True verbose_proxy_logger.warning( - "JWT Key Mapping (auto_register_map_existing_key): claim '%s' is not the user_id or user_email JWT field " - "and may be shared by several users, so a new key is minted instead of reusing one the user owns.", + "JWT Key Mapping: claim '%s' is not the user_id or user_email JWT field " + "and may be shared by several users, so existing user keys cannot be linked automatically.", claim_field, ) return False async def _reusable_key_hash_for_user(prisma_client: PrismaClient, user_id: str, team_id: str | None) -> str | None: - key: Final = await VerificationTokenRepository(prisma_client).find_newest_reusable_llm_api_key(user_id, team_id) + key: Final = await VerificationTokenRepository(prisma_client).find_newest_reusable_key(user_id, team_id) return None if key is None else key.token @@ -1628,11 +1632,13 @@ async def _auto_register_jwt_mapping( org_id: str | None = None, end_user_id: str | None = None, agent_id: str | None = None, -) -> UserAPIKeyAuth | None: + behavior: UnregisteredJWTClientBehavior = UnregisteredJWTClientBehavior.AUTO_REGISTER, +) -> UserAPIKeyAuth | _JWTKeyMappingFailure | None: """ Auto-register: create a new virtual key + mapping for an unrecognised JWT claim value, or point the mapping at a key the resolved user already owns - when ``auto_register_map_existing_key`` is set. ``team_id`` and ``user_id`` + when ``auto_register_map_existing_key`` is set. LINK_ONLY always searches + for an existing key and rejects when none is eligible. ``team_id`` and ``user_id`` must come from a successful ``JWTAuthManager.auth_builder`` run — they encode the JWT identity AFTER RBAC/scope/custom_validate/email-domain policy has been enforced. The key @@ -1650,11 +1656,21 @@ async def _auto_register_jwt_mapping( generate_key_helper_fn, ) + link_only: Final = behavior == UnregisteredJWTClientBehavior.LINK_ONLY + reuse_existing: Final = link_only or jwt_handler.litellm_jwtauth.auto_register_map_existing_key existing_token_hash: Final = ( - await _reusable_key_hash_for_user(prisma_client, user_id, team_id) - if user_id is not None and _claim_identifies_user(jwt_handler, virtual_key_claim_field, jwt_issuer) + await _reusable_key_hash_for_user( + prisma_client, + user_id, + team_id, + ) + if reuse_existing + and user_id is not None + and _claim_identifies_user(jwt_handler, virtual_key_claim_field, jwt_issuer) else None ) + if link_only and existing_token_hash is None: + return _JWTKeyMappingFailure.NO_ELIGIBLE_KEY minted: Final = existing_token_hash is None if existing_token_hash is not None: token_hash = existing_token_hash @@ -1691,8 +1707,8 @@ async def _auto_register_jwt_mapping( "jwt_claim_name": virtual_key_claim_field, "jwt_claim_value": claim_value, "token": token_hash, - "created_by": "auto_register", - "updated_by": "auto_register", + "created_by": behavior.value, + "updated_by": behavior.value, } ) except Exception as e: @@ -1704,8 +1720,8 @@ async def _auto_register_jwt_mapping( # the DB and the cleartext is in memory on this request. Delete it # so orphans don't accumulate under sustained concurrency. verbose_proxy_logger.debug( - "JWT Key Mapping (auto_register): unique conflict on create — " - "deleting orphaned virtual key and fetching winner's mapping for %s='%s'.", + "JWT Key Mapping (%s): unique conflict on create, fetching winner's mapping for %s='%s'.", + behavior.value, virtual_key_claim_field, claim_value, ) @@ -1735,7 +1751,7 @@ async def _auto_register_jwt_mapping( raise HTTPException( status_code=503, detail=( - "JWT Key Mapping: AUTO_REGISTER race resolution failed — " + f"JWT Key Mapping: {behavior.value.upper()} race resolution failed — " "winner's mapping was concurrently removed. Retry the request." ), ) @@ -1749,7 +1765,8 @@ async def _auto_register_jwt_mapping( ) verbose_proxy_logger.info( - "JWT Key Mapping (auto_register): %s virtual key for %s='%s'.", + "JWT Key Mapping (%s): %s virtual key for %s='%s'.", + behavior.value, "created new" if minted else "mapped existing", virtual_key_claim_field, claim_value, @@ -1821,7 +1838,7 @@ async def _resolve_jwt_to_virtual_key( - ``UserAPIKeyAuth``: a resolved virtual key (cache hit or DB hit). The caller may use this directly; JWT policy has been enforced previously (at key-creation time or, for cached results, before caching). - - ``_PendingAutoRegister``: claim is unmapped and behavior is AUTO_REGISTER. + - ``_PendingAutoRegister``: claim is unmapped and behavior is AUTO_REGISTER or LINK_ONLY. The caller MUST run ``JWTAuthManager.auth_builder`` to enforce JWT policy (RBAC, scope, custom_validate, email-domain), then invoke ``_auto_register_jwt_mapping`` with the validated team_id/user_id. @@ -1829,7 +1846,7 @@ async def _resolve_jwt_to_virtual_key( The caller falls through to standard team-based JWT auth (which itself enforces full JWT policy via auth_builder). - Raises HTTPException: REJECT policy hit, missing claim under - REJECT/AUTO_REGISTER, or other policy violations. + REJECT/AUTO_REGISTER/LINK_ONLY, or other policy violations. """ raw_issuer: Final = jwt_claims.get(JWTHandler.LITELLM_JWT_ISSUER_CLAIM) normalized_issuer: Final = raw_issuer if isinstance(raw_issuer, str) else None @@ -1856,6 +1873,7 @@ async def _resolve_jwt_to_virtual_key( if behavior in ( UnregisteredJWTClientBehavior.REJECT, UnregisteredJWTClientBehavior.AUTO_REGISTER, + UnregisteredJWTClientBehavior.LINK_ONLY, ): raise HTTPException( status_code=403, @@ -1887,7 +1905,7 @@ async def _resolve_jwt_to_virtual_key( status_code=403, detail=f"JWT Key Mapping: No registered mapping for {virtual_key_claim_field}='{claim_value}'. Access denied.", ) - if behavior == UnregisteredJWTClientBehavior.AUTO_REGISTER: + if behavior in (UnregisteredJWTClientBehavior.AUTO_REGISTER, UnregisteredJWTClientBehavior.LINK_ONLY): # Stale sentinel written under a prior fallback_team_mapping config — # evict it and defer auto-register to after auth_builder runs. Raise # the same 500 as the fresh-path AUTO_REGISTER branch when there is @@ -1897,7 +1915,7 @@ async def _resolve_jwt_to_virtual_key( raise HTTPException( status_code=500, detail=( - "JWT Key Mapping: AUTO_REGISTER requires a database connection. " + f"JWT Key Mapping: {behavior.value.upper()} requires a database connection. " "Configure a database or change unregistered_jwt_client_behavior." ), ) @@ -1907,6 +1925,7 @@ async def _resolve_jwt_to_virtual_key( claim_value=str(claim_value), cache_key=cache_key, jwt_issuer=normalized_issuer, + behavior=behavior, ) return None elif cached_mapping is not None: @@ -1967,12 +1986,12 @@ async def _resolve_jwt_to_virtual_key( detail=f"JWT Key Mapping: No registered mapping for {virtual_key_claim_field}='{claim_value}'. Access denied.", ) - if behavior == UnregisteredJWTClientBehavior.AUTO_REGISTER: + if behavior in (UnregisteredJWTClientBehavior.AUTO_REGISTER, UnregisteredJWTClientBehavior.LINK_ONLY): if prisma_client is None: raise HTTPException( status_code=500, detail=( - "JWT Key Mapping: AUTO_REGISTER requires a database connection. " + f"JWT Key Mapping: {behavior.value.upper()} requires a database connection. " "Configure a database or change unregistered_jwt_client_behavior." ), ) @@ -1985,6 +2004,7 @@ async def _resolve_jwt_to_virtual_key( claim_value=str(claim_value), cache_key=cache_key, jwt_issuer=normalized_issuer, + behavior=behavior, ) # FALLBACK_TEAM_MAPPING (default): cache the miss and return None so the @@ -2477,10 +2497,13 @@ async def user_api_key_auth_builder( org_id=org_id, end_user_id=end_user_id, agent_id=agent_id, + behavior=pending_auto_register.behavior, ) if pending_auto_register is not None and prisma_client is not None else None ) + if isinstance(auto_registered, _JWTKeyMappingFailure): + raise HTTPException(status_code=403, detail=auto_registered.value) if auto_registered is not None: auto_registered.jwt_claims = jwt_claims auto_registered.user_email = user_email @@ -2496,7 +2519,13 @@ async def user_api_key_auth_builder( falls_through_to_key_checks: Final = ( auto_registered is not None - and jwt_handler.litellm_jwtauth.auto_register_map_existing_key + and ( + jwt_handler.litellm_jwtauth.auto_register_map_existing_key + or ( + pending_auto_register is not None + and pending_auto_register.behavior == UnregisteredJWTClientBehavior.LINK_ONLY + ) + ) and master_key is not None ) if not falls_through_to_key_checks: diff --git a/litellm/repositories/verification_token_repository.py b/litellm/repositories/verification_token_repository.py index b20fb47306e..fb1f1c89e69 100644 --- a/litellm/repositories/verification_token_repository.py +++ b/litellm/repositories/verification_token_repository.py @@ -124,9 +124,7 @@ class VerificationTokenRepository(BaseRepository[LiteLLM_VerificationToken]): records: Final[Sequence[PrismaVerificationToken]] = await self.table.find_many(where={"user_id": user_id}) return self._to_model_list(records) - async def find_newest_reusable_llm_api_key( - self, user_id: str, team_id: str | None - ) -> LiteLLM_VerificationToken | None: + async def find_newest_reusable_key(self, user_id: str, team_id: str | None) -> LiteLLM_VerificationToken | None: records: Final[Sequence[PrismaVerificationToken]] = await self.table.find_many( where={ "user_id": user_id, @@ -140,12 +138,6 @@ class VerificationTokenRepository(BaseRepository[LiteLLM_VerificationToken]): {"team_id": {"not": UI_SESSION_TOKEN_TEAM_ID}}, ] }, - { - "OR": [ - {"allowed_routes": {"is_empty": True}}, - {"allowed_routes": {"has": "llm_api_routes"}}, - ] - }, ], }, order={"created_at": "desc"}, diff --git a/tests/integration/authorization/test_jwt_auto_register_map_existing_key.py b/tests/integration/authorization/test_jwt_auto_register_map_existing_key.py index 62eda977438..42fb55a5c4b 100644 --- a/tests/integration/authorization/test_jwt_auto_register_map_existing_key.py +++ b/tests/integration/authorization/test_jwt_auto_register_map_existing_key.py @@ -7,13 +7,15 @@ from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager from hashlib import sha256 from pathlib import Path -from typing import Final +from typing import Final, Literal import httpx import jwt import pytest import yaml from cryptography.hazmat.primitives.asymmetric import rsa +from pydantic import JsonValue +from redis import Redis from tests.integration._support.client import Gateway, eventually, string_value from tests.integration._support.database import read_rows @@ -31,7 +33,15 @@ def _hash(key: str) -> str: return sha256(key.encode()).hexdigest() -def _config(directory: Path, claim_field: str) -> Path: +def _config( + directory: Path, + claim_field: str, + *, + behavior: Literal["auto_register", "link_only", "reject"] = "auto_register", + map_existing_key: bool = True, + user_id_upsert: bool = False, + issuers: tuple[dict[str, JsonValue], ...] = (), +) -> Path: config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) config["general_settings"] = { **config["general_settings"], @@ -40,8 +50,10 @@ def _config(directory: Path, claim_field: str) -> Path: "user_id_jwt_field": "sub", "user_email_jwt_field": "email", "virtual_key_claim_field": claim_field, - "unregistered_jwt_client_behavior": "auto_register", - "auto_register_map_existing_key": True, + "unregistered_jwt_client_behavior": behavior, + "auto_register_map_existing_key": map_existing_key, + "user_id_upsert": user_id_upsert, + "issuers": list(issuers), }, } path: Final = directory / f"jwt_map_existing_key_{claim_field}.yaml" @@ -107,7 +119,7 @@ def _billed_key(response: httpx.Response) -> str: return string_value(rows[0]["api_key"]) -def test_first_jwt_call_reuses_the_newest_durable_llm_key_and_skips_every_ineligible_newer_key( +def test_first_jwt_call_reuses_the_newest_durable_key_and_skips_every_ineligible_newer_key( gateway: Gateway, tmp_path: Path ) -> None: with _issuer() as (private_key, jwks_url), gateway.scenario() as scenario: @@ -118,7 +130,6 @@ def test_first_jwt_call_reuses_the_newest_durable_llm_key_and_skips_every_inelig skipped: Final = { "older_durable": older_durable, "expiring": scenario.key(user_id=user, duration="1h"), - "management_only": scenario.key(user_id=user, allowed_routes=["management_routes"]), "auto_registered_look_alike": scenario.key(user_id=user, metadata={"auto_registered": True}), "other_team": scenario.key(user_id=user, team_id=scenario.team()), "blocked": scenario.key(user_id=user), @@ -140,6 +151,84 @@ def test_first_jwt_call_reuses_the_newest_durable_llm_key_and_skips_every_inelig assert _billed_key(response) == _hash(durable) +@pytest.mark.parametrize( + ("behavior", "map_existing_key"), [("auto_register", True), ("link_only", False), ("link_only", True)] +) +def test_route_restricted_keys_are_linked_without_bypassing_other_key_checks( + gateway: Gateway, + tmp_path: Path, + behavior: Literal["auto_register", "link_only"], + map_existing_key: bool, +) -> None: + with _issuer() as (private_key, jwks_url), gateway.scenario() as scenario: + model: Final = scenario.model() + user: Final = scenario.user(user_role="internal_user") + scenario.key(user_id=user) + restricted: Final = scenario.key(user_id=user, allowed_routes=["/v1/chat/completions"]) + scenario.key(user_id=user, duration="1h") + scenario.key(user_id=user, metadata={"auto_registered": True}) + scenario.key(user_id=user, team_id=scenario.team()) + blocked: Final = scenario.key(user_id=user) + gateway.post("/key/block", {"key": blocked}) + keys_before: Final = _user_key_hashes(user) + token: Final = _token(private_key, user) + + with owned_proxy( + gateway, + tmp_path, + {"JWT_PUBLIC_KEY_URL": jwks_url}, + config=_config( + tmp_path, + "sub", + behavior=behavior, + map_existing_key=map_existing_key, + ), + ) as candidate: + response: Final = _chat(candidate, model, token) + repeated: Final = _chat(candidate, model, token) + restricted_route: Final = candidate.request("GET", "/v1/models", key=token) + + assert response.status_code == 200, response.text + assert repeated.status_code == 200, repeated.text + assert restricted_route.status_code == 403, restricted_route.text + assert _mapped_token("sub", user) == _hash(restricted) + assert _user_key_hashes(user) == keys_before, "a new key was minted instead of linking an existing key" + assert _billed_key(response) == _hash(restricted) + assert _billed_key(repeated) == _hash(restricted) + + +@pytest.mark.parametrize("behavior", ["auto_register", "link_only"]) +def test_management_route_keys_are_linked_by_default_and_keep_their_route_permissions( + gateway: Gateway, tmp_path: Path, behavior: Literal["auto_register", "link_only"] +) -> None: + with _issuer() as (private_key, jwks_url), gateway.scenario() as scenario: + model: Final = scenario.model() + user: Final = scenario.user(user_role="internal_user") + target: Final = f"integration-{uuid.uuid4().hex}" + scenario.key(user_id=user) + management_key: Final = scenario.key(user_id=user, allowed_routes=["/user/new"]) + keys_before: Final = _user_key_hashes(user) + token: Final = _token(private_key, user) + with owned_proxy( + gateway, tmp_path, {"JWT_PUBLIC_KEY_URL": jwks_url}, config=_config(tmp_path, "sub", behavior=behavior) + ) as candidate: + denied: Final = _chat(candidate, model, token) + allowed: Final = candidate.request( + "POST", + "/user/new", + {"user_id": target, "auto_create_key": False, "user_role": "internal_user"}, + key=token, + ) + assert allowed.status_code == 200, allowed.text + scenario.cleanups.callback(scenario.delete_user, target) + repeated: Final = _chat(candidate, model, token) + assert denied.status_code == 403, denied.text + assert repeated.status_code == 403, repeated.text + assert allowed.json()["user_id"] == target + assert _mapped_token("sub", user) == _hash(management_key) + assert _user_key_hashes(user) == keys_before + + def test_user_matched_by_email_instead_of_sub_still_reuses_their_existing_key(gateway: Gateway, tmp_path: Path) -> None: with _issuer() as (private_key, jwks_url), gateway.scenario() as scenario: model: Final = scenario.model() @@ -183,8 +272,9 @@ def test_shared_client_claim_never_maps_a_second_user_onto_the_first_users_perso assert read_rows('SELECT request_id FROM "LiteLLM_SpendLogs" WHERE api_key = %s', (_hash(personal),)) == [] -def test_concurrent_first_jwt_calls_of_a_keyless_user_both_succeed_on_one_surviving_mapped_key( - gateway: Gateway, tmp_path: Path +@pytest.mark.parametrize("behavior", ["auto_register", "link_only"]) +def test_concurrent_first_jwt_calls_both_succeed_on_one_surviving_mapped_key( + gateway: Gateway, tmp_path: Path, behavior: Literal["auto_register", "link_only"] ) -> None: writer_url: Final = os.environ.get("INTEGRATION_PROXY_DATABASE_URL") or os.environ["DATABASE_URL"] with ( @@ -194,6 +284,7 @@ def test_concurrent_first_jwt_calls_of_a_keyless_user_both_succeed_on_one_surviv ): model: Final = scenario.model() user: Final = scenario.user(user_role="internal_user") + existing: Final = scenario.key(user_id=user) if behavior == "link_only" else None token: Final = _token(private_key, user) overrides: Final = { "JWT_PUBLIC_KEY_URL": jwks_url, @@ -202,7 +293,12 @@ def test_concurrent_first_jwt_calls_of_a_keyless_user_both_succeed_on_one_surviv } with ( - owned_proxy(gateway, tmp_path, overrides, config=_config(tmp_path, "sub")) as candidate, + owned_proxy( + gateway, + tmp_path, + overrides, + config=_config(tmp_path, "sub", behavior=behavior, map_existing_key=False), + ) as candidate, ThreadPoolExecutor(max_workers=1) as pool, ): held_call: Final = pool.submit(_chat, candidate, model, token) @@ -215,5 +311,240 @@ def test_concurrent_first_jwt_calls_of_a_keyless_user_both_succeed_on_one_surviv assert held.status_code == 200, held.text keys: Final = _user_key_hashes(user) assert len(keys) == 1, keys + if existing is not None: + assert keys == frozenset({_hash(existing)}) assert _mapped_token("sub", user) in keys assert (_billed_key(held), _billed_key(racing)) == (_mapped_token("sub", user),) * 2 + + +@pytest.mark.parametrize("ineligible_keys", [False, True]) +def test_link_only_rejects_without_writes_then_links_a_newly_available_key( + gateway: Gateway, tmp_path: Path, ineligible_keys: bool +) -> None: + with _issuer() as (private_key, jwks_url), gateway.scenario() as scenario: + model: Final = scenario.model() + user: Final = scenario.user(user_role="internal_user") + if ineligible_keys: + scenario.key(user_id=user, duration="1h") + scenario.key(user_id=user, metadata={"auto_registered": True}) + scenario.key(user_id=user, team_id=scenario.team()) + blocked: Final = scenario.key(user_id=user) + gateway.post("/key/block", {"key": blocked}) + keys_before: Final = _user_key_hashes(user) + token: Final = _token(private_key, user) + with owned_proxy( + gateway, + tmp_path, + {"JWT_PUBLIC_KEY_URL": jwks_url}, + config=_config(tmp_path, "sub", behavior="link_only", map_existing_key=False), + ) as candidate: + first: Final = _chat(candidate, model, token) + repeated: Final = _chat(candidate, model, token) + assert first.status_code == 403, first.text + assert repeated.status_code == 403, repeated.text + assert _user_key_hashes(user) == keys_before + assert read_rows('SELECT token FROM "LiteLLM_JWTKeyMapping" WHERE jwt_claim_value=%s', (user,)) == [] + available: Final = scenario.key(user_id=user) + accepted: Final = _chat(candidate, model, token) + assert accepted.status_code == 200, accepted.text + assert _mapped_token("sub", user) == _hash(available) + assert _user_key_hashes(user) == keys_before | {_hash(available)} + assert _billed_key(accepted) == _hash(available) + + +@pytest.mark.parametrize("claim_field", ["client_id", "email"]) +def test_link_only_rejects_shared_or_missing_claims_without_linking_a_personal_key( + gateway: Gateway, tmp_path: Path, claim_field: str +) -> None: + with _issuer() as (private_key, jwks_url), gateway.scenario() as scenario: + model: Final = scenario.model() + user: Final = scenario.user(user_role="internal_user") + key: Final = scenario.key(user_id=user) + shared_claim: Final = f"integration-client-{uuid.uuid4().hex}" + with owned_proxy( + gateway, + tmp_path, + {"JWT_PUBLIC_KEY_URL": jwks_url}, + config=_config(tmp_path, claim_field, behavior="link_only", map_existing_key=False), + ) as candidate: + response: Final = _chat(candidate, model, _token(private_key, user, client_id=shared_claim)) + assert response.status_code == 403, response.text + assert _user_key_hashes(user) == frozenset({_hash(key)}) + assert read_rows('SELECT token FROM "LiteLLM_JWTKeyMapping" WHERE token=%s', (_hash(key),)) == [] + assert read_rows('SELECT token FROM "LiteLLM_JWTKeyMapping" WHERE jwt_claim_value=%s', (shared_claim,)) == [] + + +@pytest.mark.parametrize("restriction", ["models", "allowed_routes"]) +def test_link_only_enforces_key_permissions_on_the_first_and_cached_requests( + gateway: Gateway, tmp_path: Path, restriction: str +) -> None: + with _issuer() as (private_key, jwks_url), gateway.scenario() as scenario: + allowed: Final = scenario.model() + denied: Final = scenario.model() + user: Final = scenario.user(user_role="internal_user") + key: Final = scenario.key( + user_id=user, + **({"models": [allowed]} if restriction == "models" else {"allowed_routes": ["/v1/models"]}), + ) + token: Final = _token(private_key, user) + with owned_proxy( + gateway, + tmp_path, + {"JWT_PUBLIC_KEY_URL": jwks_url}, + config=_config(tmp_path, "sub", behavior="link_only", map_existing_key=False), + ) as candidate: + first: Final = _chat(candidate, denied, token) + repeated: Final = _chat(candidate, denied, token) + control: Final = ( + _chat(candidate, allowed, token) + if restriction == "models" + else candidate.request("GET", "/v1/models", key=token) + ) + assert first.status_code == 403, first.text + assert repeated.status_code == 403, repeated.text + assert control.status_code == 200, control.text + assert _mapped_token("sub", user) == _hash(key) + assert _user_key_hashes(user) == frozenset({_hash(key)}) + + +def test_link_only_does_not_replace_an_existing_mapping_to_a_blocked_key(gateway: Gateway, tmp_path: Path) -> None: + with _issuer() as (private_key, jwks_url), gateway.scenario() as scenario: + model: Final = scenario.model() + user: Final = scenario.user(user_role="internal_user") + blocked: Final = scenario.key(user_id=user) + available: Final = scenario.key(user_id=user) + gateway.post("/jwt/key/mapping/new", {"jwt_claim_name": "sub", "jwt_claim_value": user, "key": blocked}) + gateway.post("/key/block", {"key": blocked}) + token: Final = _token(private_key, user) + with owned_proxy( + gateway, + tmp_path, + {"JWT_PUBLIC_KEY_URL": jwks_url}, + config=_config(tmp_path, "sub", behavior="link_only", map_existing_key=False), + ) as candidate: + first: Final = _chat(candidate, model, token) + repeated: Final = _chat(candidate, model, token) + assert first.status_code == 401, first.text + assert repeated.status_code == 401, repeated.text + assert _mapped_token("sub", user) == _hash(blocked) + assert _user_key_hashes(user) == frozenset({_hash(blocked), _hash(available)}) + + +def test_link_only_exempts_proxy_admins_but_rechecks_the_role_after_a_cached_exemption( + gateway: Gateway, tmp_path: Path +) -> None: + with _issuer() as (private_key, jwks_url), gateway.scenario() as scenario: + model: Final = scenario.model() + user: Final = scenario.user(user_role="internal_user") + admin: Final = _token(private_key, user, scope="litellm_proxy_admin") + ordinary: Final = _token(private_key, user) + with owned_proxy( + gateway, + tmp_path, + {"JWT_PUBLIC_KEY_URL": jwks_url}, + config=_config(tmp_path, "sub", behavior="link_only", map_existing_key=False), + ) as candidate: + first: Final = candidate.request("GET", "/v1/models", key=admin) + repeated: Final = candidate.request("GET", "/v1/models", key=admin) + denied: Final = candidate.request("GET", "/v1/models", key=ordinary) + forbidden_route: Final = _chat(candidate, model, admin) + assert first.status_code == 200, first.text + assert repeated.status_code == 200, repeated.text + assert denied.status_code == 403, denied.text + assert forbidden_route.status_code == 401, forbidden_route.text + assert _user_key_hashes(user) == frozenset() + assert read_rows('SELECT token FROM "LiteLLM_JWTKeyMapping" WHERE jwt_claim_value=%s', (user,)) == [] + + +@pytest.mark.parametrize("user_id_upsert", [False, True]) +def test_link_only_keeps_user_provisioning_independent_of_key_creation( + gateway: Gateway, tmp_path: Path, user_id_upsert: bool +) -> None: + with _issuer() as (private_key, jwks_url), gateway.scenario() as scenario: + model: Final = scenario.model() + subject: Final = f"integration-link-only-{uuid.uuid4().hex}" + try: + with owned_proxy( + gateway, + tmp_path, + {"JWT_PUBLIC_KEY_URL": jwks_url}, + config=_config( + tmp_path, "sub", behavior="link_only", map_existing_key=False, user_id_upsert=user_id_upsert + ), + ) as candidate: + response: Final = _chat(candidate, model, _token(private_key, subject)) + assert response.status_code == (403 if user_id_upsert else 401), response.text + users: Final = read_rows('SELECT user_id FROM "LiteLLM_UserTable" WHERE user_id=%s', (subject,)) + assert users == ([{"user_id": subject}] if user_id_upsert else []) + assert _user_key_hashes(subject) == frozenset() + assert read_rows('SELECT token FROM "LiteLLM_JWTKeyMapping" WHERE jwt_claim_value=%s', (subject,)) == [] + finally: + if read_rows('SELECT user_id FROM "LiteLLM_UserTable" WHERE user_id=%s', (subject,)): + scenario.delete_user(subject) + + +@pytest.mark.parametrize("issuer_behavior", ["auto_register", "link_only"]) +def test_issuer_specific_registration_mode_overrides_the_global_mode( + gateway: Gateway, tmp_path: Path, issuer_behavior: Literal["auto_register", "link_only"] +) -> None: + with _issuer() as (private_key, jwks_url), gateway.scenario() as scenario: + model: Final = scenario.model() + user: Final = scenario.user(user_role="internal_user") + issuer: Final = f"https://integration.example.com/{uuid.uuid4().hex}" + token: Final = _token(private_key, user, iss=issuer, aud="integration") + path: Final = _config( + tmp_path, + "sub", + behavior="auto_register" if issuer_behavior == "link_only" else "link_only", + map_existing_key=False, + issuers=( + { + "issuer": issuer, + "jwks_url": jwks_url, + "audience": "integration", + "user_id_jwt_field": "sub", + "virtual_key_claim_field": "sub", + "unregistered_jwt_client_behavior": issuer_behavior, + }, + ), + ) + with owned_proxy(gateway, tmp_path, {"JWT_PUBLIC_KEY_URL": jwks_url}, config=path) as candidate: + response: Final = _chat(candidate, model, token) + assert response.status_code == (403 if issuer_behavior == "link_only" else 200), response.text + if issuer_behavior == "link_only": + assert _user_key_hashes(user) == frozenset() + assert read_rows('SELECT token FROM "LiteLLM_JWTKeyMapping" WHERE jwt_claim_value=%s', (user,)) == [] + else: + assert len(_user_key_hashes(user)) == 1 + assert _mapped_token("sub", user) in _user_key_hashes(user) + assert _billed_key(response) == _mapped_token("sub", user) + + +def test_link_only_links_after_a_cached_rejection_from_the_previous_mode(gateway: Gateway, tmp_path: Path) -> None: + with _issuer() as (private_key, jwks_url), gateway.scenario() as scenario: + model: Final = scenario.model() + user: Final = scenario.user(user_role="internal_user") + key: Final = scenario.key(user_id=user) + token: Final = _token(private_key, user) + with owned_proxy( + gateway, + tmp_path, + {"JWT_PUBLIC_KEY_URL": jwks_url}, + config=_config(tmp_path, "sub", behavior="reject", map_existing_key=False), + ) as rejecting: + denied: Final = _chat(rejecting, model, token) + assert denied.status_code == 403, denied.text + with Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) as cache: + cached: Final = cache.get(f"jwt_key_mapping:sub:{user}") + assert cached is not None and b"__NO_MAPPING__" in cached, cached + with owned_proxy( + gateway, + tmp_path, + {"JWT_PUBLIC_KEY_URL": jwks_url}, + config=_config(tmp_path, "sub", behavior="link_only", map_existing_key=False), + ) as linking: + accepted: Final = _chat(linking, model, token) + assert accepted.status_code == 200, accepted.text + assert _mapped_token("sub", user) == _hash(key) + assert _user_key_hashes(user) == frozenset({_hash(key)}) + assert _billed_key(accepted) == _hash(key)