mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-12 23:01:41 +00:00
fix(jwt): invalidate JWT key mapping cache on /key/regenerate
/key/regenerate carries the JWT-to-key mapping to the new token via FK cascade, but the jwt_key_mapping cache entry kept resolving the old (now invalid) token for up to virtual_key_mapping_cache_ttl. Snapshot the key's mapping cache keys before the token update and evict them with evict_and_broadcast so every worker drops the stale entry. Also share the cache-key format through jwt_key_mapping_cache_key and upgrade the /jwt/key/mapping CRUD endpoints from local-only deletes to evict_and_broadcast, closing the same cross-worker staleness there.
This commit is contained in:
parent
300d335255
commit
2f7ee39545
5 changed files with 147 additions and 11 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,18 @@ 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)
|
||||
old_cache_key: Final = jwt_key_mapping_cache_key(old_mapping.jwt_claim_name, old_mapping.jwt_claim_value)
|
||||
await evict_and_broadcast(cache_keys=(old_cache_key,), user_api_key_cache=user_api_key_cache)
|
||||
|
||||
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)
|
||||
new_cache_key: Final = jwt_key_mapping_cache_key(
|
||||
updated_mapping.jwt_claim_name, updated_mapping.jwt_claim_value
|
||||
)
|
||||
await evict_and_broadcast(cache_keys=(new_cache_key,), user_api_key_cache=user_api_key_cache)
|
||||
|
||||
return _to_response(updated_mapping)
|
||||
except HTTPException:
|
||||
|
|
@ -219,8 +221,8 @@ 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)
|
||||
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)
|
||||
|
||||
await _mapping_table(prisma_client).delete(where={"id": data.id})
|
||||
return {"status": "success"}
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.get_new_token",
|
||||
new_callable=AsyncMock,
|
||||
return_value="sk-newtoken1234ab12",
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key",
|
||||
new_callable=AsyncMock,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
|
||||
new_callable=AsyncMock,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook",
|
||||
new_callable=AsyncMock,
|
||||
),
|
||||
patch(
|
||||
"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(
|
||||
"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."""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue