diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 2b328b455dc..f83f0303deb 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -142,6 +142,8 @@ class _PrismaDictableRow(Protocol): class _PrismaJWTKeyMappingRow(Protocol): token: str + jwt_claim_name: str + jwt_claim_value: str class _PrismaModelDumpRow(Protocol): @@ -3466,6 +3468,23 @@ 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}" + + +@log_db_metrics +async def get_jwt_key_mapping_cache_keys_for_token( + hashed_token: str, + prisma_client: PrismaClient, +) -> tuple[str, ...]: + """Cache keys of every JWT claim mapped to the given virtual key.""" + 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) + + @log_db_metrics async def get_jwt_key_mapping_object( jwt_claim_name: str, diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 5fb6dad0cd7..93293db24c6 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -58,6 +58,7 @@ from litellm.proxy.auth.auth_checks import ( get_team_object, get_user_object, is_valid_fallback_model, + jwt_key_mapping_cache_key, resolve_and_validate_end_user_id, ) from litellm.proxy.auth.auth_exception_handler import UserAPIKeyAuthExceptionHandler @@ -970,7 +971,7 @@ async def _resolve_jwt_to_virtual_key( ) return None - cache_key: Final = f"jwt_key_mapping:{virtual_key_claim_field}:{claim_value}" + cache_key: Final = jwt_key_mapping_cache_key(virtual_key_claim_field, str(claim_value)) cached_mapping: Final = await user_api_key_cache.async_get_cache(cache_key) if cached_mapping == _JWT_PROXY_ADMIN_SENTINEL: diff --git a/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py index ccfd5338ec4..694930a543c 100644 --- a/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py +++ b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py @@ -13,7 +13,9 @@ from litellm.proxy._types import ( UserAPIKeyAuth, hash_token, ) +from litellm.proxy.auth.auth_checks import jwt_key_mapping_cache_key from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view from litellm.repositories.table_repositories import JWTKeyMappingRepository @@ -118,9 +120,8 @@ async def create_jwt_key_mapping( new_mapping: Final = await _mapping_table(prisma_client).create(data=create_data) - # Invalidate cache - cache_key: Final = f"jwt_key_mapping:{data.jwt_claim_name}:{data.jwt_claim_value}" - await user_api_key_cache.async_delete_cache(cache_key) + cache_key: Final = jwt_key_mapping_cache_key(data.jwt_claim_name, data.jwt_claim_value) + await evict_and_broadcast(cache_keys=(cache_key,), user_api_key_cache=user_api_key_cache) return _to_response(new_mapping) except HTTPException: @@ -169,17 +170,20 @@ async def update_jwt_key_mapping( if old_mapping is None: raise HTTPException(status_code=404, detail="Mapping not found") - cache_key = f"jwt_key_mapping:{old_mapping.jwt_claim_name}:{old_mapping.jwt_claim_value}" - await user_api_key_cache.async_delete_cache(cache_key) - updated_mapping: Final = await _mapping_table(prisma_client).update(where={"id": data.id}, data=update_data) if updated_mapping is None: raise HTTPException(status_code=404, detail="Mapping not found") - # Invalidate new cache key if claim fields changed - cache_key = f"jwt_key_mapping:{updated_mapping.jwt_claim_name}:{updated_mapping.jwt_claim_value}" - await user_api_key_cache.async_delete_cache(cache_key) + # 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) + new_cache_key: Final = jwt_key_mapping_cache_key( + updated_mapping.jwt_claim_name, updated_mapping.jwt_claim_value + ) + 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) return _to_response(updated_mapping) except HTTPException: @@ -219,10 +223,12 @@ async def delete_jwt_key_mapping( if old_mapping is None: raise HTTPException(status_code=404, detail="Mapping not found") - cache_key: Final = f"jwt_key_mapping:{old_mapping.jwt_claim_name}:{old_mapping.jwt_claim_value}" - await user_api_key_cache.async_delete_cache(cache_key) - await _mapping_table(prisma_client).delete(where={"id": data.id}) + + # 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) + await evict_and_broadcast(cache_keys=(cache_key,), user_api_key_cache=user_api_key_cache) return {"status": "success"} except HTTPException: raise diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 324d380b85b..ebcfab090b5 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -54,6 +54,7 @@ from litellm.proxy._types import Litellm_EntityType, LiteLLM_VerificationToken, from litellm.proxy.auth.auth_checks import ( _delete_cache_key_object, can_team_access_model, + get_jwt_key_mapping_cache_keys_for_token, get_org_object, get_project_object, get_team_object, @@ -65,6 +66,7 @@ from litellm.proxy.auth.auth_utils import ( ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( + evict_and_broadcast, publish_auth_cache_invalidation, ) from litellm.proxy.common_utils.callback_config_validation import logging_metadata_config_error @@ -4975,6 +4977,13 @@ async def _execute_virtual_key_regeneration( update_data.update(non_default_values) jsonified_update_data: Final[Mapping[str, object]] = prisma_client.jsonify_object(data=update_data) + # Snapshot before the token update: the FK cascade rewrites mapping rows to the new hash, + # but their cached jwt_key_mapping entries still point at the old token (LIT-5379). + jwt_mapping_cache_keys: Final = await get_jwt_key_mapping_cache_keys_for_token( + hashed_token=hashed_api_key, + prisma_client=prisma_client, + ) + # If grace period set, insert deprecated key so old key remains valid await _insert_deprecated_key( prisma_client=prisma_client, @@ -5000,6 +5009,8 @@ async def _execute_virtual_key_regeneration( proxy_logging_obj=proxy_logging_obj, ) + await evict_and_broadcast(cache_keys=jwt_mapping_cache_keys, user_api_key_cache=user_api_key_cache) + # After credential invalidation, so a failure here can never keep the old key alive. await sync_key_regeneration_access_group_membership( prisma_client=prisma_client, diff --git a/tests/proxy_unit_tests/test_jwt_key_mapping.py b/tests/proxy_unit_tests/test_jwt_key_mapping.py index 4b50f83e9eb..e8db5d1cf7f 100644 --- a/tests/proxy_unit_tests/test_jwt_key_mapping.py +++ b/tests/proxy_unit_tests/test_jwt_key_mapping.py @@ -1333,3 +1333,86 @@ def test_jwt_client_id_field_does_not_raise_on_duplicate(): virtual_key_claim_field="new_field", ) assert auth.virtual_key_claim_field == "new_field" + + +# ────────────────────────────────────────────── +# Tests: cache eviction must happen AFTER the DB write commits +# ────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_delete_evicts_cache_after_row_is_gone(): + """A JWT request racing the delete must not keep the removed mapping authorized. + + The DB delete simulates a concurrent request re-caching the mapping mid-write. + If the endpoint evicts before the delete commits, that repopulated entry + survives until TTL and the deleted mapping stays usable. + """ + from litellm.proxy._types import DeleteJWTKeyMappingRequest + from litellm.proxy.auth.auth_checks import jwt_key_mapping_cache_key + + cache_key = jwt_key_mapping_cache_key("email", "user@example.com") + user_api_key_cache = DualCache() + await user_api_key_cache.async_set_cache(key=cache_key, value="hashed_token") + + mock_prisma = _mock_prisma() + mock_prisma.db.litellm_jwtkeymapping.find_unique.return_value = _mock_mapping() + + async def concurrent_reader_repopulates(**kwargs): + await user_api_key_cache.async_set_cache(key=cache_key, value="hashed_token") + return _mock_mapping() + + mock_prisma.db.litellm_jwtkeymapping.delete.side_effect = concurrent_reader_repopulates + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), # test-quality-ok: proxy_server module global is the endpoint's only injection point + patch("litellm.proxy.proxy_server.user_api_key_cache", user_api_key_cache), # test-quality-ok: proxy_server module global is the endpoint's only injection point + ): + result = await delete_jwt_key_mapping( + data=DeleteJWTKeyMappingRequest(id="mapping-1"), + user_api_key_dict=_make_admin_auth(), + ) + + assert result == {"status": "success"} + assert await user_api_key_cache.async_get_cache(cache_key) is None + + +@pytest.mark.asyncio +async def test_update_evicts_old_and_new_cache_keys_after_write(): + """Renaming a mapping's claim must leave neither claim serving stale cache. + + The DB update simulates a concurrent request re-caching the OLD mapping + mid-write. Both the old claim's entry (would restore the pre-rename token) + and the new claim's __NO_MAPPING__ sentinel (would 403 the renamed claim) + must be gone once the endpoint returns. + """ + from litellm.proxy._types import UpdateJWTKeyMappingRequest + from litellm.proxy.auth.auth_checks import jwt_key_mapping_cache_key + + old_cache_key = jwt_key_mapping_cache_key("email", "user@example.com") + new_cache_key = jwt_key_mapping_cache_key("email", "renamed@example.com") + user_api_key_cache = DualCache() + await user_api_key_cache.async_set_cache(key=old_cache_key, value="hashed_token") + await user_api_key_cache.async_set_cache(key=new_cache_key, value="__NO_MAPPING__") + + mock_prisma = _mock_prisma() + mock_prisma.db.litellm_jwtkeymapping.find_unique.return_value = _mock_mapping() + + async def concurrent_reader_repopulates(**kwargs): + await user_api_key_cache.async_set_cache(key=old_cache_key, value="hashed_token") + return _mock_mapping(claim_value="renamed@example.com") + + mock_prisma.db.litellm_jwtkeymapping.update.side_effect = concurrent_reader_repopulates + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), # test-quality-ok: proxy_server module global is the endpoint's only injection point + patch("litellm.proxy.proxy_server.user_api_key_cache", user_api_key_cache), # test-quality-ok: proxy_server module global is the endpoint's only injection point + ): + result = await update_jwt_key_mapping( + data=UpdateJWTKeyMappingRequest(id="mapping-1", jwt_claim_value="renamed@example.com"), + user_api_key_dict=_make_admin_auth(), + ) + + assert result.jwt_claim_value == "renamed@example.com" + assert await user_api_key_cache.async_get_cache(old_cache_key) is None + assert await user_api_key_cache.async_get_cache(new_cache_key) is None 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 0e4af9f75a5..47571497f74 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 @@ -11912,6 +11912,109 @@ async def test_execute_virtual_key_regeneration_allows_within_limit_duration(mon assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 1 +@pytest.mark.asyncio +async def test_regenerate_evicts_jwt_key_mapping_cache_so_next_jwt_call_gets_new_token(): + """ + LIT-5379: /key/regenerate rewrites the JWT mapping row to the new token (FK + cascade) but left the jwt_key_mapping cache entry pointing at the old hash, + so JWT calls kept resolving the dead token until the cache TTL expired. + Regenerate must evict the entry locally, broadcast the eviction to other + workers, and the very next JWT resolve must return the rotated token. + """ + from litellm.caching.caching import DualCache + from litellm.proxy._types import LiteLLM_JWTAuth + from litellm.proxy.auth.auth_method import AuthMethod + from litellm.proxy.auth.resolvers.models import CredentialRef + from litellm.proxy.auth.resolvers.store import IdentityStore + from litellm.proxy.auth.user_api_key_auth import _resolve_jwt_to_virtual_key + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _execute_virtual_key_regeneration, + ) + + stale_cache_key = "jwt_key_mapping:sub:user1" + 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")] + ) + mock_prisma_client.db.litellm_jwtkeymapping.find_first = AsyncMock( + return_value=MagicMock(token="new-hashed-token") + ) + user_api_key_cache = DualCache() + await user_api_key_cache.async_set_cache(key=stale_cache_key, value="abc123") + + publish_mock = AsyncMock() + with ( + patch( # test-quality-ok: deterministic token; same pattern as sibling regenerate tests + "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token", + new_callable=AsyncMock, + return_value="sk-newtoken1234ab12", + ), + patch( # test-quality-ok: grace-period path not under test; same pattern as sibling regenerate tests + "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: key-object eviction is separate from the mapping eviction under test + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: background rotation hook is irrelevant to cache eviction + "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: captures the cross-worker broadcast without a redis instance + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation", + publish_mock, + ), + ): + await _execute_virtual_key_regeneration( + prisma_client=mock_prisma_client, + key_in_db=existing_key, + hashed_api_key="abc123", + key="abc123", + data=None, + user_api_key_dict=_make_regenerate_user_api_key_dict(), + litellm_changed_by=None, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=MagicMock(), + ) + + assert await user_api_key_cache.async_get_cache(stale_cache_key) is None + publish_mock.assert_any_await(cache_key=stale_cache_key) + mock_prisma_client.db.litellm_jwtkeymapping.find_many.assert_awaited_once_with(where={"token": "abc123"}) + + rotated_key = UserAPIKeyAuth(token="new-hashed-token", user_id="user-1") + rotated_principal = IdentityStore._principal_from_key( + rotated_key, + auth_method=AuthMethod.API_KEY, + credential_ref=CredentialRef(token_id="new-hashed-token"), + ) + + async def fake_resolve(hashed_token): + assert hashed_token == "new-hashed-token", f"JWT resolved stale token {hashed_token!r} after regenerate" + return rotated_principal + + jwt_handler = MagicMock() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( + virtual_key_claim_field="sub", virtual_key_mapping_cache_ttl=300 + ) + with patch( # test-quality-ok: DB-backed resolve; fake asserts it receives the rotated hash + "litellm.proxy.auth.resolvers.store.IdentityStore.resolve", + new_callable=AsyncMock, + side_effect=fake_resolve, + ): + resolved = await _resolve_jwt_to_virtual_key( + jwt_claims={"sub": "user1"}, + jwt_handler=jwt_handler, + prisma_client=mock_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 == "new-hashed-token" + + @pytest.mark.asyncio async def test_execute_virtual_key_regeneration_rejects_over_limit_max_budget(monkeypatch): """Regenerate must reject max_budget exceeding upperbound — proves the fix covers non-duration fields."""