From e54b93017ba37e49948e7d01f4f4d794bb913a3f Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Tue, 15 Sep 2026 13:16:50 -0700 Subject: [PATCH] fix(jwt-auth): scope JWT key mappings by issuer to prevent cross-issuer collisions --- .../migration.sql | 18 ++ .../litellm_proxy_extras/schema.prisma | 8 +- litellm/proxy/_lazy_openapi_snapshot.json | 33 ++++ litellm/proxy/_types.py | 3 + litellm/proxy/auth/auth_checks.py | 26 ++- litellm/proxy/auth/user_api_key_auth.py | 55 +++++- .../jwt_key_mapping_endpoints.py | 25 ++- litellm/proxy/schema.prisma | 8 +- schema.prisma | 8 +- .../proxy_unit_tests/test_jwt_key_mapping.py | 180 ++++++++++++++++++ .../proxy/auth/test_user_api_key_auth.py | 179 ++++++++++++++++- .../test_key_management_endpoints.py | 15 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 6 + 13 files changed, 527 insertions(+), 37 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260915000000_scope_jwt_key_mapping_by_issuer/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260915000000_scope_jwt_key_mapping_by_issuer/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260915000000_scope_jwt_key_mapping_by_issuer/migration.sql new file mode 100644 index 00000000000..c9572066ab6 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260915000000_scope_jwt_key_mapping_by_issuer/migration.sql @@ -0,0 +1,18 @@ +-- DropIndex +DROP INDEX IF EXISTS "LiteLLM_JWTKeyMapping_jwt_claim_name_jwt_claim_value_is_act_idx"; + +-- DropIndex +DROP INDEX IF EXISTS "LiteLLM_JWTKeyMapping_jwt_claim_name_jwt_claim_value_key"; + +-- AlterTable +-- NOT NULL DEFAULT '' (not nullable): Postgres unique constraints treat every +-- NULL as distinct, so a nullable column would let multiple unscoped mappings +-- collide on the same claim without a constraint violation. The constant +-- default is a fast, metadata-only backfill for existing rows, not a rewrite. +ALTER TABLE "LiteLLM_JWTKeyMapping" ADD COLUMN IF NOT EXISTS "jwt_issuer" TEXT NOT NULL DEFAULT ''; + +-- CreateIndex +CREATE INDEX IF NOT EXISTS "LiteLLM_JWTKeyMapping_jwt_issuer_jwt_claim_name_jwt_claim_v_idx" ON "LiteLLM_JWTKeyMapping"("jwt_issuer", "jwt_claim_name", "jwt_claim_value", "is_active"); + +-- CreateIndex +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_JWTKeyMapping_jwt_issuer_jwt_claim_name_jwt_claim_v_key" ON "LiteLLM_JWTKeyMapping"("jwt_issuer", "jwt_claim_name", "jwt_claim_value"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 8072df5aa5b..62853d8e4b8 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -487,6 +487,10 @@ model LiteLLM_VerificationToken { model LiteLLM_JWTKeyMapping { id String @id @default(uuid()) + jwt_issuer String @default("") // Scopes the mapping to one configured issuer; "" matches any issuer. + // Not nullable: Postgres unique constraints treat every NULL as + // distinct, so a nullable column would let multiple unscoped + // mappings collide on the same claim without a constraint violation. jwt_claim_name String // e.g. "sub", "email" jwt_claim_value String // The claim value to match token String // Hashed virtual key (FK) @@ -499,8 +503,8 @@ model LiteLLM_JWTKeyMapping { litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token], onDelete: Cascade) - @@unique([jwt_claim_name, jwt_claim_value]) - @@index([jwt_claim_name, jwt_claim_value, is_active]) + @@unique([jwt_issuer, jwt_claim_name, jwt_claim_value]) + @@index([jwt_issuer, jwt_claim_name, jwt_claim_value, is_active]) } // Deprecated keys during grace period - allows old key to work until revoke_at diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index f3b579d22c7..c5d1e7e8ece 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -15226,6 +15226,17 @@ "title": "Jwt Claim Value", "type": "string" }, + "jwt_issuer": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Jwt Issuer" + }, "key": { "title": "Key", "type": "string" @@ -15310,6 +15321,17 @@ "title": "Jwt Claim Value", "type": "string" }, + "jwt_issuer": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Jwt Issuer" + }, "updated_at": { "format": "date-time", "title": "Updated At", @@ -15366,6 +15388,17 @@ ], "title": "Is Active" }, + "jwt_issuer": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Jwt Issuer" + }, "key": { "anyOf": [ { diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index afc2150934e..9e71a54c5ba 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -4485,12 +4485,14 @@ class CreateJWTKeyMappingRequest(LiteLLMPydanticObjectBase): jwt_claim_name: str jwt_claim_value: str key: str + jwt_issuer: str | None = None description: str | None = None class UpdateJWTKeyMappingRequest(LiteLLMPydanticObjectBase): id: str key: str | None = None + jwt_issuer: str | None = None description: str | None = None is_active: bool | None = None @@ -4501,6 +4503,7 @@ class DeleteJWTKeyMappingRequest(LiteLLMPydanticObjectBase): class JWTKeyMappingResponse(LiteLLMPydanticObjectBase): id: str + jwt_issuer: str | None = None jwt_claim_name: str jwt_claim_value: str description: str | None = None diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 355fc3f6a21..e3783c94dc7 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -148,6 +148,7 @@ class _PrismaDictableRow(Protocol): class _PrismaJWTKeyMappingRow(Protocol): token: str + jwt_issuer: str jwt_claim_name: str jwt_claim_value: str @@ -3601,9 +3602,18 @@ async def _fetch_key_object_from_db_with_reconnect( raise -def jwt_key_mapping_cache_key(jwt_claim_name: str, jwt_claim_value: str) -> str: - """Cache key under which ``_resolve_jwt_to_virtual_key`` stores a JWT-claim-to-key mapping.""" - return f"jwt_key_mapping:{jwt_claim_name}:{jwt_claim_value}" +def jwt_key_mapping_cache_key(jwt_claim_name: str, jwt_claim_value: str, jwt_issuer: str | None = None) -> str: + """Cache key under which a JWT-claim-to-key mapping is stored, scoped to one + issuer (or the issuer-agnostic/global scope when ``jwt_issuer`` is falsy). + + Scoped by issuer (when one is configured) so a cached hit or ``__NO_MAPPING__`` miss + for one issuer's claim value can never be served to a different issuer whose claim + value happens to collide. Unchanged for the global scope, keeping the single-issuer + (no ``litellm_jwtauth.issuers`` configured) cache key format stable across this fix. + """ + if not jwt_issuer: + return f"jwt_key_mapping:{jwt_claim_name}:{jwt_claim_value}" + return f"jwt_key_mapping:{jwt_issuer}:{jwt_claim_name}:{jwt_claim_value}" @log_db_metrics @@ -3615,7 +3625,7 @@ async def get_jwt_key_mapping_cache_keys_for_token( mappings: Final = await _jwt_key_mapping_table(JWTKeyMappingRepository(prisma_client)).find_many( where={"token": hashed_token} ) - return tuple(jwt_key_mapping_cache_key(m.jwt_claim_name, m.jwt_claim_value) for m in mappings) + return tuple(jwt_key_mapping_cache_key(m.jwt_claim_name, m.jwt_claim_value, m.jwt_issuer) for m in mappings) @log_db_metrics @@ -3623,9 +3633,14 @@ async def get_jwt_key_mapping_object( jwt_claim_name: str, jwt_claim_value: str, prisma_client: PrismaClient, + jwt_issuer: str | None = None, ) -> str | None: """ - Lookup a JWT-to-virtual-key mapping from the database. + Lookup a JWT-to-virtual-key mapping from the database for one exact scope: + ``jwt_issuer`` (or the global/issuer-agnostic scope when falsy). Does not fall + back to the global scope itself -- a caller that wants "issuer-scoped mapping, + else the global one" queries both scopes itself, so each result can be cached + under its own scope's key (see ``_resolve_jwt_to_virtual_key``). Returns the hashed token (str) if a matching active mapping is found, else None. """ @@ -3633,6 +3648,7 @@ async def get_jwt_key_mapping_object( where={ "jwt_claim_name": jwt_claim_name, "jwt_claim_value": jwt_claim_value, + "jwt_issuer": jwt_issuer or "", "is_active": True, } ) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 7d62baf39a8..1beacae819f 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -269,6 +269,17 @@ class _TokenTeamModels(Protocol): def team_models(self) -> list[str]: ... +class _RawCacheRead(Protocol): + async def async_get_cache(self, *, key: str) -> object: ... + + +def _raw_cache(cache: _RawCacheRead) -> _RawCacheRead: + """View an untyped cache object's ``async_get_cache`` as returning ``object`` + instead of ``Any``, so a caller can ``isinstance``-narrow it without paying + the ``reportAny`` cost of the underlying (unannotated) cache implementation.""" + return cache + + def _token_team_models(valid_token: _TokenTeamModels) -> list[str]: return valid_token.team_models @@ -842,6 +853,7 @@ class _PendingAutoRegister(NamedTuple): claim_field: str claim_value: str cache_key: str + jwt_issuer: str | None = None async def _auto_register_jwt_mapping( @@ -853,6 +865,7 @@ async def _auto_register_jwt_mapping( parent_otel_span: Span | None, proxy_logging_obj: ProxyLogging, cache_key: str, + jwt_issuer: str | None = None, team_id: str | None = None, user_id: str | None = None, org_id: str | None = None, @@ -905,6 +918,7 @@ async def _auto_register_jwt_mapping( try: await prisma_client.db.litellm_jwtkeymapping.create( data={ + "jwt_issuer": jwt_issuer or "", "jwt_claim_name": virtual_key_claim_field, "jwt_claim_value": claim_value, "token": token_hash, @@ -939,6 +953,7 @@ async def _auto_register_jwt_mapping( jwt_claim_name=virtual_key_claim_field, jwt_claim_value=claim_value, prisma_client=prisma_client, + jwt_issuer=jwt_issuer, ) if token_hash is None: # The winner's mapping vanished between the unique-constraint @@ -1041,7 +1056,7 @@ async def _resolve_jwt_to_virtual_key( ) return None - cache_key: Final = jwt_key_mapping_cache_key(virtual_key_claim_field, str(claim_value)) + cache_key: Final = jwt_key_mapping_cache_key(virtual_key_claim_field, str(claim_value), normalized_issuer) raw_cached_mapping: Final = await user_api_key_cache.async_get_cache(cache_key) sentinel_written_by_this_policy: Final = behavior == UnregisteredJWTClientBehavior.AUTO_REGISTER cached_mapping: Final = ( @@ -1081,6 +1096,7 @@ async def _resolve_jwt_to_virtual_key( claim_field=virtual_key_claim_field, claim_value=str(claim_value), cache_key=cache_key, + jwt_issuer=normalized_issuer, ) return None elif cached_mapping is not None: @@ -1094,21 +1110,44 @@ async def _resolve_jwt_to_virtual_key( ) # Resolve the mapping from DB, or treat prisma_client=None as a definitive - # miss (no DB → no mapping can exist → apply no-match policy below). + # miss (no DB → no mapping can exist → apply no-match policy below). An + # issuer-scoped row wins; falling back to the global (no-issuer) row keeps + # mappings created before issuer scoping existed working for every issuer. + # Each tier is cached under ITS OWN key (the global tier under the + # issuer-less cache key, not under `cache_key`/this issuer's key) so that + # updating or deleting either row invalidates exactly the cache entries it + # can affect. Caching a global-row hit under the requesting issuer's key + # would leave every OTHER issuer that had fallen back to that same global + # mapping serving its stale token until TTL after the row changes. + ttl: Final = jwt_handler.litellm_jwtauth.virtual_key_mapping_cache_ttl token_hash: str | None = None if prisma_client is not None: token_hash = await get_jwt_key_mapping_object( jwt_claim_name=virtual_key_claim_field, jwt_claim_value=str(claim_value), prisma_client=prisma_client, + jwt_issuer=normalized_issuer, ) + if token_hash is not None: + await user_api_key_cache.async_set_cache(key=cache_key, value=token_hash, ttl=ttl) + elif normalized_issuer is not None: + # Another issuer may have already resolved (and cached) this same + # global mapping -- check its cache entry before re-querying the DB. + global_cache_key: Final = jwt_key_mapping_cache_key(virtual_key_claim_field, str(claim_value)) + cached_global: Final = await _raw_cache(user_api_key_cache).async_get_cache(key=global_cache_key) + if isinstance(cached_global, str) and cached_global != "__NO_MAPPING__": + token_hash = cached_global + else: + token_hash = await get_jwt_key_mapping_object( + jwt_claim_name=virtual_key_claim_field, + jwt_claim_value=str(claim_value), + prisma_client=prisma_client, + jwt_issuer=None, + ) + if token_hash is not None: + await user_api_key_cache.async_set_cache(key=global_cache_key, value=token_hash, ttl=ttl) if token_hash is not None: - await user_api_key_cache.async_set_cache( - key=cache_key, - value=token_hash, - ttl=jwt_handler.litellm_jwtauth.virtual_key_mapping_cache_ttl, - ) return IdentityStore.key_from_principal( await IdentityStore( prisma_client, @@ -1149,6 +1188,7 @@ async def _resolve_jwt_to_virtual_key( claim_field=virtual_key_claim_field, claim_value=str(claim_value), cache_key=cache_key, + jwt_issuer=normalized_issuer, ) # FALLBACK_TEAM_MAPPING (default): cache the miss and return None so the @@ -1641,6 +1681,7 @@ async def _user_api_key_auth_builder( parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, cache_key=pending_auto_register.cache_key, + jwt_issuer=pending_auto_register.jwt_issuer, team_id=team_id, user_id=user_id, org_id=org_id, diff --git a/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py index 694930a543c..07234883062 100644 --- a/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py +++ b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py @@ -28,6 +28,9 @@ class _JWTKeyMappingRecord(Protocol): @property def id(self) -> str: ... + @property + def jwt_issuer(self) -> str: ... + @property def jwt_claim_name(self) -> str: ... @@ -78,6 +81,7 @@ def _to_response(mapping: _JWTKeyMappingRecord) -> JWTKeyMappingResponse: """Convert a Prisma mapping object to a safe response (no hashed token).""" return JWTKeyMappingResponse( id=mapping.id, + jwt_issuer=mapping.jwt_issuer or None, jwt_claim_name=mapping.jwt_claim_name, jwt_claim_value=mapping.jwt_claim_value, description=mapping.description, @@ -109,6 +113,7 @@ async def create_jwt_key_mapping( try: hashed_key: Final = hash_token(data.key) create_data: Final = { + "jwt_issuer": data.jwt_issuer or "", "jwt_claim_name": data.jwt_claim_name, "jwt_claim_value": data.jwt_claim_value, "token": hashed_key, @@ -120,7 +125,7 @@ async def create_jwt_key_mapping( new_mapping: Final = await _mapping_table(prisma_client).create(data=create_data) - cache_key: Final = jwt_key_mapping_cache_key(data.jwt_claim_name, data.jwt_claim_value) + cache_key: Final = jwt_key_mapping_cache_key(data.jwt_claim_name, data.jwt_claim_value, data.jwt_issuer) await evict_and_broadcast(cache_keys=(cache_key,), user_api_key_cache=user_api_key_cache) return _to_response(new_mapping) @@ -131,7 +136,10 @@ async def create_jwt_key_mapping( if "unique" in error_str or "p2002" in error_str: raise HTTPException( status_code=409, - detail=f"A mapping for claim '{data.jwt_claim_name}' = '{data.jwt_claim_value}' already exists.", + detail=( + f"A mapping for claim '{data.jwt_claim_name}' = '{data.jwt_claim_value}' " + f"already exists for issuer '{data.jwt_issuer}'." + ), ) if "foreign" in error_str or "p2003" in error_str: raise HTTPException( @@ -161,6 +169,9 @@ async def update_jwt_key_mapping( update_data: Final = data.model_dump(exclude_unset=True, exclude={"id", "key"}) if data.key is not None: update_data["token"] = hash_token(data.key) + if "jwt_issuer" in update_data: + # DB column is NOT NULL (see schema.prisma); "" is the global/unscoped sentinel. + update_data["jwt_issuer"] = update_data["jwt_issuer"] or "" update_data["updated_by"] = user_api_key_dict.user_id try: @@ -178,9 +189,11 @@ async def update_jwt_key_mapping( # Evict only after the write commits: a concurrent request between an # early eviction and the commit would re-cache the old mapping and keep # it authorized until TTL. - old_cache_key: Final = jwt_key_mapping_cache_key(old_mapping.jwt_claim_name, old_mapping.jwt_claim_value) + old_cache_key: Final = jwt_key_mapping_cache_key( + old_mapping.jwt_claim_name, old_mapping.jwt_claim_value, old_mapping.jwt_issuer + ) new_cache_key: Final = jwt_key_mapping_cache_key( - updated_mapping.jwt_claim_name, updated_mapping.jwt_claim_value + updated_mapping.jwt_claim_name, updated_mapping.jwt_claim_value, updated_mapping.jwt_issuer ) cache_keys: Final = (old_cache_key,) if old_cache_key == new_cache_key else (old_cache_key, new_cache_key) await evict_and_broadcast(cache_keys=cache_keys, user_api_key_cache=user_api_key_cache) @@ -227,7 +240,9 @@ async def delete_jwt_key_mapping( # Evict only after the row is gone, else a concurrent request can # re-cache the deleted mapping and keep it authorized until TTL. - cache_key: Final = jwt_key_mapping_cache_key(old_mapping.jwt_claim_name, old_mapping.jwt_claim_value) + cache_key: Final = jwt_key_mapping_cache_key( + old_mapping.jwt_claim_name, old_mapping.jwt_claim_value, old_mapping.jwt_issuer + ) await evict_and_broadcast(cache_keys=(cache_key,), user_api_key_cache=user_api_key_cache) return {"status": "success"} except HTTPException: diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 8072df5aa5b..62853d8e4b8 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -487,6 +487,10 @@ model LiteLLM_VerificationToken { model LiteLLM_JWTKeyMapping { id String @id @default(uuid()) + jwt_issuer String @default("") // Scopes the mapping to one configured issuer; "" matches any issuer. + // Not nullable: Postgres unique constraints treat every NULL as + // distinct, so a nullable column would let multiple unscoped + // mappings collide on the same claim without a constraint violation. jwt_claim_name String // e.g. "sub", "email" jwt_claim_value String // The claim value to match token String // Hashed virtual key (FK) @@ -499,8 +503,8 @@ model LiteLLM_JWTKeyMapping { litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token], onDelete: Cascade) - @@unique([jwt_claim_name, jwt_claim_value]) - @@index([jwt_claim_name, jwt_claim_value, is_active]) + @@unique([jwt_issuer, jwt_claim_name, jwt_claim_value]) + @@index([jwt_issuer, jwt_claim_name, jwt_claim_value, is_active]) } // Deprecated keys during grace period - allows old key to work until revoke_at diff --git a/schema.prisma b/schema.prisma index 8072df5aa5b..62853d8e4b8 100644 --- a/schema.prisma +++ b/schema.prisma @@ -487,6 +487,10 @@ model LiteLLM_VerificationToken { model LiteLLM_JWTKeyMapping { id String @id @default(uuid()) + jwt_issuer String @default("") // Scopes the mapping to one configured issuer; "" matches any issuer. + // Not nullable: Postgres unique constraints treat every NULL as + // distinct, so a nullable column would let multiple unscoped + // mappings collide on the same claim without a constraint violation. jwt_claim_name String // e.g. "sub", "email" jwt_claim_value String // The claim value to match token String // Hashed virtual key (FK) @@ -499,8 +503,8 @@ model LiteLLM_JWTKeyMapping { litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token], onDelete: Cascade) - @@unique([jwt_claim_name, jwt_claim_value]) - @@index([jwt_claim_name, jwt_claim_value, is_active]) + @@unique([jwt_issuer, jwt_claim_name, jwt_claim_value]) + @@index([jwt_issuer, jwt_claim_name, jwt_claim_value, is_active]) } // Deprecated keys during grace period - allows old key to work until revoke_at diff --git a/tests/proxy_unit_tests/test_jwt_key_mapping.py b/tests/proxy_unit_tests/test_jwt_key_mapping.py index e8db5d1cf7f..3f2c04336a7 100644 --- a/tests/proxy_unit_tests/test_jwt_key_mapping.py +++ b/tests/proxy_unit_tests/test_jwt_key_mapping.py @@ -91,6 +91,154 @@ async def test_jwt_to_virtual_key_mapping_resolution(): prisma_client.db.litellm_jwtkeymapping.find_first.assert_not_called() +@pytest.mark.asyncio +async def test_colliding_claim_value_from_another_issuer_does_not_resolve_to_the_wrong_virtual_key(): + """LIT-7417: a mapping registered for one issuer must not answer a lookup from a + DIFFERENT issuer whose claim value happens to collide, even though both issuers + map the same claim field (``sub``) to a virtual key.""" + issuer_a = "https://issuer-a.example.com" + issuer_b = "https://issuer-b.example.com" + + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( + virtual_key_claim_field="sub", virtual_key_mapping_cache_ttl=3600 + ) + + rows = [ + { + "jwt_issuer": issuer_b, + "jwt_claim_name": "sub", + "jwt_claim_value": "dev-alice", + "token": "hashed-issuer-b-key", + "is_active": True, + } + ] + + async def fake_find_first(where): + for row in rows: + if all(row.get(k) == v for k, v in where.items()): + return MagicMock(**row) + return None + + prisma_client = MagicMock() + prisma_client.db.litellm_jwtkeymapping.find_first = AsyncMock(side_effect=fake_find_first) + + # Dependency-inject the resolved key via the cache (IdentityStore._resolve_key + # reads it from here) instead of monkeypatching IdentityStore itself. + user_api_key_cache = DualCache() + await user_api_key_cache.async_set_cache( + key="hashed-issuer-b-key", + value=UserAPIKeyAuth(token="hashed-issuer-b-key", team_id="issuer-b-team"), + ) + + # The rightful owner: issuer-b's own claim resolves to its mapping. + owner_result = await _resolve_jwt_to_virtual_key( + jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: issuer_b, "sub": "dev-alice"}, + jwt_handler=jwt_handler, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=None, + ) + assert isinstance(owner_result, UserAPIKeyAuth) + assert owner_result.token == "hashed-issuer-b-key" + + # A validly-signed token from issuer-a carrying the SAME claim value must not + # inherit issuer-b's mapping. Default behavior is fallback_team_mapping, so a + # correctly-scoped miss returns None instead of resolving to issuer-b's key. + colliding_result = await _resolve_jwt_to_virtual_key( + jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: issuer_a, "sub": "dev-alice"}, + jwt_handler=jwt_handler, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=None, + ) + assert colliding_result is None + + +@pytest.mark.asyncio +async def test_global_mapping_resolution_is_cached_under_the_global_key_not_the_requesting_issuer(): + """LIT-7417: caching a global (unscoped) mapping's hit under the REQUESTING + issuer's key would leave every issuer that falls back to it holding its own + stale copy after the row is updated/deleted -- CRUD only evicts the cache key + computed from the row's own scope (global), so a copy cached under some other + issuer's key would keep resolving to the old token until TTL. Caching it under + the global key instead means every issuer shares (and CRUD correctly evicts) + the exact same entry.""" + issuer_a = "https://issuer-a.example.com" + issuer_b = "https://issuer-b.example.com" + + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( + issuers=[ + { + "issuer": issuer_a, + "jwks_url": f"{issuer_a}/jwks", + "virtual_key_claim_field": "sub", + "disable_audience_validation": True, + }, + { + "issuer": issuer_b, + "jwks_url": f"{issuer_b}/jwks", + "virtual_key_claim_field": "sub", + "disable_audience_validation": True, + }, + ] + ) + + rows = [ + { + "jwt_issuer": "", + "jwt_claim_name": "sub", + "jwt_claim_value": "legacy-user", + "token": "hashed-legacy-key", + "is_active": True, + } + ] + + async def fake_find_first(where): + for row in rows: + if all(row.get(k) == v for k, v in where.items()): + return MagicMock(**row) + return None + + prisma_client = MagicMock() + find_first = AsyncMock(side_effect=fake_find_first) + prisma_client.db.litellm_jwtkeymapping.find_first = find_first + + user_api_key_cache = DualCache() + await user_api_key_cache.async_set_cache( + key="hashed-legacy-key", + value=UserAPIKeyAuth(token="hashed-legacy-key", team_id="legacy-team"), + ) + + resolved_a = await _resolve_jwt_to_virtual_key( + jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: issuer_a, "sub": "legacy-user"}, + jwt_handler=jwt_handler, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=None, + ) + assert isinstance(resolved_a, UserAPIKeyAuth) + assert find_first.await_count == 2 # issuer-a-scoped miss, then global hit + + # issuer-b resolving the SAME global mapping must hit the cache issuer-a's + # resolution populated, not issue a fresh DB query for the global row again. + resolved_b = await _resolve_jwt_to_virtual_key( + jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: issuer_b, "sub": "legacy-user"}, + jwt_handler=jwt_handler, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=None, + ) + assert isinstance(resolved_b, UserAPIKeyAuth) + assert resolved_b.token == "hashed-legacy-key" + assert find_first.await_count == 3 # +1 for issuer-b's own issuer-scoped miss; global tier served from cache + + @pytest.mark.asyncio async def test_jwt_to_virtual_key_mapping_no_mapping(): """ @@ -223,6 +371,7 @@ def test_to_response_excludes_token(): now = datetime.now(timezone.utc) mock_mapping = MagicMock() mock_mapping.id = "mapping-1" + mock_mapping.jwt_issuer = None mock_mapping.jwt_claim_name = "email" mock_mapping.jwt_claim_value = "user@example.com" mock_mapping.token = "hashed_secret_value" @@ -275,10 +424,12 @@ def _mock_mapping( id="mapping-1", claim_name="email", claim_value="user@example.com", + issuer=None, ): now = datetime.now(timezone.utc) m = MagicMock() m.id = id + m.jwt_issuer = issuer m.jwt_claim_name = claim_name m.jwt_claim_value = claim_value m.token = "hashed_token" @@ -485,6 +636,35 @@ async def test_create_success_returns_response_without_token(): assert result.jwt_claim_name == "email" +@pytest.mark.asyncio +async def test_create_without_issuer_stores_empty_string_not_null(): + """LIT-7417: the DB column is NOT NULL (see schema.prisma). Storing a real NULL + for an unscoped mapping would let Postgres accept unlimited duplicate unscoped + rows for the same claim (NULL is never equal to NULL in a unique constraint), + so two mappings for the same claim value could point at two different keys with + no conflict, and resolution would pick whichever one Postgres returns first.""" + from litellm.proxy._types import CreateJWTKeyMappingRequest + + mock_prisma = _mock_prisma() + mock_prisma.db.litellm_jwtkeymapping.create.return_value = _mock_mapping() + mock_cache = AsyncMock() + + data = CreateJWTKeyMappingRequest(jwt_claim_name="sub", jwt_claim_value="dev-alice", key="sk-test-key") + + with ( + patch( # test-quality-ok: proxy_server module global is the endpoint's only injection point + "litellm.proxy.proxy_server.prisma_client", mock_prisma + ), + patch( # test-quality-ok: proxy_server module global is the endpoint's only injection point + "litellm.proxy.proxy_server.user_api_key_cache", mock_cache + ), + ): + await create_jwt_key_mapping(data=data, user_api_key_dict=_make_admin_auth()) + + sent_data = mock_prisma.db.litellm_jwtkeymapping.create.call_args.kwargs["data"] + assert sent_data["jwt_issuer"] == "" + + # ────────────────────────────────────────────── # Tests: unregistered_jwt_client_behavior # ────────────────────────────────────────────── diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 866ea0b20e4..bd7ff62ac8b 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -32,7 +32,13 @@ from litellm.proxy._types import ( JWTRoutingOverride, ) from litellm.proxy.auth.handle_jwt import JWTHandler -from litellm.proxy.auth.auth_checks import TeamNotFoundError, UserNotFoundError, get_key_object, _cache_key_object +from litellm.proxy.auth.auth_checks import ( + TeamNotFoundError, + UserNotFoundError, + get_key_object, + _cache_key_object, + jwt_key_mapping_cache_key, +) from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.auth.user_api_key_auth import ( _check_key_model_budget_with_fallback, @@ -7948,13 +7954,38 @@ def _per_issuer_virtual_key_jwt_handler( def _fake_prisma_with_jwt_key_mapping(hashed_token: str | None) -> tuple[SimpleNamespace, AsyncMock]: + """Every ``find_first`` call (issuer-scoped or global fallback) resolves the same way.""" find_first = AsyncMock(return_value=None if hashed_token is None else SimpleNamespace(token=hashed_token)) prisma_client = SimpleNamespace(db=SimpleNamespace(litellm_jwtkeymapping=SimpleNamespace(find_first=find_first))) return prisma_client, find_first -def _mapping_where(claim_name: str, claim_value: str) -> dict[str, str | bool]: - return {"jwt_claim_name": claim_name, "jwt_claim_value": claim_value, "is_active": True} +def _fake_prisma_jwt_key_mapping_table(rows: list[dict[str, object]]) -> tuple[SimpleNamespace, AsyncMock]: + """A ``find_first`` whose result depends on the ``where`` clause, like a real table. + + Matches a row when every key present in ``where`` equals that key on the row -- + a key ``get_jwt_key_mapping_object`` omits (e.g. old, issuer-blind code never + sending ``jwt_issuer``) does not constrain the match, exactly like Prisma. + """ + + async def _find_first(where: dict[str, object]) -> SimpleNamespace | None: + for row in rows: + if all(row.get(k) == v for k, v in where.items()): + return SimpleNamespace(**row) + return None + + find_first = AsyncMock(side_effect=_find_first) + prisma_client = SimpleNamespace(db=SimpleNamespace(litellm_jwtkeymapping=SimpleNamespace(find_first=find_first))) + return prisma_client, find_first + + +def _mapping_where(claim_name: str, claim_value: str, jwt_issuer: str | None) -> dict[str, str | bool]: + return { + "jwt_claim_name": claim_name, + "jwt_claim_value": claim_value, + "jwt_issuer": jwt_issuer or "", + "is_active": True, + } @pytest.mark.asyncio @@ -7978,11 +8009,13 @@ async def test_per_issuer_virtual_key_claim_field_selects_the_issuer_mapping_for proxy_logging_obj=MagicMock(), ) - find_first.assert_awaited_once_with(where=_mapping_where("sub", "svc-account-7")) + # Issuer-scoped lookup hits on the first query, so no global fallback query runs. + find_first.assert_awaited_once_with(where=_mapping_where("sub", "svc-account-7", ISSUER_TWO)) assert isinstance(resolved, UserAPIKeyAuth) assert resolved.token == "hashed-mapped-key" assert resolved.team_id == "svc-team" - assert await user_api_key_cache.async_get_cache("jwt_key_mapping:sub:svc-account-7") == "hashed-mapped-key" + cache_key = jwt_key_mapping_cache_key("sub", "svc-account-7", ISSUER_TWO) + assert await user_api_key_cache.async_get_cache(cache_key) == "hashed-mapped-key" @pytest.mark.asyncio @@ -8015,7 +8048,11 @@ async def test_per_issuer_reject_behavior_does_not_leak_into_the_team_issuer(): assert exc.value.status_code == 403 assert "No registered mapping for sub='unknown-svc'" in str(exc.value.detail) - find_first.assert_awaited_once_with(where=_mapping_where("sub", "unknown-svc")) + # REJECT checks the issuer-scoped row first, then falls back to a global (NULL-issuer) row. + assert [c.kwargs["where"] for c in find_first.await_args_list] == [ + _mapping_where("sub", "unknown-svc", ISSUER_TWO), + _mapping_where("sub", "unknown-svc", None), + ] @pytest.mark.asyncio @@ -8025,7 +8062,10 @@ async def test_proxy_admin_sentinel_cached_by_another_issuer_does_not_bypass_rej jwt_handler = _per_issuer_virtual_key_jwt_handler(global_claim_field="sub", global_behavior="auto_register") prisma_client, find_first = _fake_prisma_with_jwt_key_mapping(None) user_api_key_cache = DualCache() - await user_api_key_cache.async_set_cache(key="jwt_key_mapping:sub:admin-7", value=_JWT_PROXY_ADMIN_SENTINEL) + # Sentinel cached under issuer-one's own key -- must never answer issuer-two's lookup. + await user_api_key_cache.async_set_cache( + key=jwt_key_mapping_cache_key("sub", "admin-7", ISSUER_ONE), value=_JWT_PROXY_ADMIN_SENTINEL + ) auto_register_issuer_result = await _resolve_jwt_to_virtual_key( jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_ONE, "sub": "admin-7"}, @@ -8050,7 +8090,10 @@ async def test_proxy_admin_sentinel_cached_by_another_issuer_does_not_bypass_rej assert exc.value.status_code == 403 assert "No registered mapping for sub='admin-7'" in str(exc.value.detail) - find_first.assert_awaited_once_with(where=_mapping_where("sub", "admin-7")) + assert [c.kwargs["where"] for c in find_first.await_args_list] == [ + _mapping_where("sub", "admin-7", ISSUER_TWO), + _mapping_where("sub", "admin-7", None), + ] @pytest.mark.asyncio @@ -8079,7 +8122,125 @@ async def test_issuer_without_virtual_key_claim_field_falls_back_to_the_global_f assert with_claim is None assert without_claim is None - find_first.assert_awaited_once_with(where=_mapping_where("client_id", "app-9")) + # without_claim has no claim value and returns before ever reaching the DB. + assert [c.kwargs["where"] for c in find_first.await_args_list] == [ + _mapping_where("client_id", "app-9", ISSUER_ONE), + _mapping_where("client_id", "app-9", None), + ] + + +@pytest.mark.asyncio +async def test_colliding_claim_value_from_another_issuer_does_not_resolve_to_the_wrong_virtual_key(): + """LIT-7417: a mapping registered for one issuer must not answer a lookup from a + DIFFERENT issuer whose claim value happens to collide, even though both issuers + use the same claim field (``sub``) for their virtual-key mapping.""" + from litellm.proxy.auth.user_api_key_auth import _resolve_jwt_to_virtual_key + + jwt_handler = _per_issuer_virtual_key_jwt_handler(global_claim_field="sub") + prisma_client, find_first = _fake_prisma_jwt_key_mapping_table( + [ + { + "jwt_issuer": ISSUER_TWO, + "jwt_claim_name": "sub", + "jwt_claim_value": "dev-alice", + "token": "hashed-issuer-b-key", + "is_active": True, + } + ] + ) + user_api_key_cache = DualCache() + await user_api_key_cache.async_set_cache( + key="hashed-issuer-b-key", + value=UserAPIKeyAuth(token="hashed-issuer-b-key", api_key="hashed-issuer-b-key", team_id="issuer-b-team"), + ) + + owner_result = await _resolve_jwt_to_virtual_key( + jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_TWO, "sub": "dev-alice"}, + jwt_handler=jwt_handler, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + ) + assert isinstance(owner_result, UserAPIKeyAuth) + assert owner_result.token == "hashed-issuer-b-key" + + # issuer-one's behavior is fallback_team_mapping: a correctly-scoped miss must + # return None (fall through to team-based JWT auth), never issuer-two's key. + colliding_result = await _resolve_jwt_to_virtual_key( + jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_ONE, "sub": "dev-alice"}, + jwt_handler=jwt_handler, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + ) + assert colliding_result is None + assert find_first.await_count == 3 # owner hit (1 call) + colliding miss (issuer-scoped + global fallback) + + +@pytest.mark.asyncio +async def test_cached_resolution_for_one_issuer_does_not_leak_to_a_colliding_issuer(): + """A cached positive resolution must be keyed by issuer too, or a colliding + claim value from another issuer could be served straight from cache without + ever reaching the (correctly issuer-scoped) DB lookup.""" + from litellm.proxy.auth.user_api_key_auth import _resolve_jwt_to_virtual_key + + jwt_handler = _per_issuer_virtual_key_jwt_handler(global_claim_field="sub") + prisma_client, find_first = _fake_prisma_jwt_key_mapping_table([]) + user_api_key_cache = DualCache() + await user_api_key_cache.async_set_cache( + key=jwt_key_mapping_cache_key("sub", "dev-alice", ISSUER_TWO), value="hashed-issuer-b-key" + ) + + colliding_result = await _resolve_jwt_to_virtual_key( + jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_ONE, "sub": "dev-alice"}, + jwt_handler=jwt_handler, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + ) + assert colliding_result is None + # Must have gone to the DB rather than serving issuer-two's cached token. + assert find_first.await_count == 2 + + +@pytest.mark.asyncio +async def test_issuer_agnostic_mapping_matches_every_issuer(): + """A mapping created before issuer scoping existed (``jwt_issuer`` is NULL) keeps + matching any issuer, so existing global mappings are not broken by this fix.""" + from litellm.proxy.auth.user_api_key_auth import _resolve_jwt_to_virtual_key + + jwt_handler = _per_issuer_virtual_key_jwt_handler(global_claim_field="sub") + prisma_client, _find_first = _fake_prisma_jwt_key_mapping_table( + [ + { + "jwt_issuer": "", + "jwt_claim_name": "sub", + "jwt_claim_value": "legacy-user", + "token": "hashed-legacy-key", + "is_active": True, + } + ] + ) + user_api_key_cache = DualCache() + await user_api_key_cache.async_set_cache( + key="hashed-legacy-key", + value=UserAPIKeyAuth(token="hashed-legacy-key", api_key="hashed-legacy-key", team_id="legacy-team"), + ) + + for issuer in (ISSUER_ONE, ISSUER_TWO): + resolved = await _resolve_jwt_to_virtual_key( + jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: issuer, "sub": "legacy-user"}, + jwt_handler=jwt_handler, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + ) + assert isinstance(resolved, UserAPIKeyAuth) + assert resolved.token == "hashed-legacy-key" @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 63055872aa1..4e70063015d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -36,7 +36,11 @@ from litellm.proxy._types import ( UpdateKeyRequest, ) from litellm.models.object_permission import LiteLLM_ObjectPermissionTable -from litellm.proxy.auth.auth_checks import _delete_cache_key_object, _project_cache_key +from litellm.proxy.auth.auth_checks import ( + _delete_cache_key_object, + _project_cache_key, + jwt_key_mapping_cache_key, +) from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.litellm_core_utils.duration_parser import duration_in_seconds @@ -5132,10 +5136,11 @@ async def test_delete_verification_tokens_persists_deleted_keys(monkeypatch): class _JWTMappingRow: - def __init__(self, token, jwt_claim_name, jwt_claim_value): + def __init__(self, token, jwt_claim_name, jwt_claim_value, jwt_issuer=None): self.token = token self.jwt_claim_name = jwt_claim_name self.jwt_claim_value = jwt_claim_value + self.jwt_issuer = jwt_issuer class _CascadingJWTMappingTable: @@ -5226,7 +5231,7 @@ async def test_delete_verification_tokens_evicts_jwt_key_mapping_cache(monkeypat ), ) - assert recording_evict.cache_keys == ("jwt_key_mapping:email:user@example.com",) + assert recording_evict.cache_keys == (jwt_key_mapping_cache_key("email", "user@example.com", None),) @pytest.mark.asyncio @@ -13131,11 +13136,11 @@ async def test_regenerate_evicts_jwt_key_mapping_cache_so_next_jwt_call_gets_new _execute_virtual_key_regeneration, ) - stale_cache_key = "jwt_key_mapping:sub:user1" + stale_cache_key = jwt_key_mapping_cache_key("sub", "user1", None) existing_key = _make_regenerate_existing_key() mock_prisma_client = _make_regenerate_mock_prisma() mock_prisma_client.db.litellm_jwtkeymapping.find_many = AsyncMock( - return_value=[MagicMock(jwt_claim_name="sub", jwt_claim_value="user1")] + return_value=[MagicMock(jwt_claim_name="sub", jwt_claim_value="user1", jwt_issuer=None)] ) mock_prisma_client.db.litellm_jwtkeymapping.find_first = AsyncMock( return_value=MagicMock(token="new-hashed-token") diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 4c16a8613eb..af3ee891cb4 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -27297,6 +27297,8 @@ export interface components { jwt_claim_name: string; /** Jwt Claim Value */ jwt_claim_value: string; + /** Jwt Issuer */ + jwt_issuer?: string | null; /** Key */ key: string; }; @@ -28910,6 +28912,8 @@ export interface components { jwt_claim_name: string; /** Jwt Claim Value */ jwt_claim_value: string; + /** Jwt Issuer */ + jwt_issuer?: string | null; /** * Updated At * Format: date-time @@ -38574,6 +38578,8 @@ export interface components { id: string; /** Is Active */ is_active?: boolean | null; + /** Jwt Issuer */ + jwt_issuer?: string | null; /** Key */ key?: string | null; };