mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
fix(proxy): restore deprecated key auth cache lookups
This commit is contained in:
parent
b8f7d61400
commit
c04a1dc1a2
2 changed files with 63 additions and 5 deletions
|
|
@ -2400,10 +2400,25 @@ async def _lookup_deprecated_key(
|
|||
|
||||
# Check cache first
|
||||
cached = _deprecated_key_cache.get(hashed_token)
|
||||
cached = _deprecated_key_cache.get(hashed_token)
|
||||
if cached is not None:
|
||||
active_token_id, cache_expires_at_ts, revoke_at_ts = cached
|
||||
if now_ts < cache_expires_at_ts and now_ts < revoke_at_ts:
|
||||
if len(cached) == 3:
|
||||
active_token_id, cache_expires_at_ts, revoke_at_ts = cached
|
||||
elif len(cached) == 2:
|
||||
# Backward compatibility for cache entries written before
|
||||
# revoke_at_ts was added to the in-memory tuple.
|
||||
active_token_id, cache_expires_at_ts = cached
|
||||
revoke_at_ts = float("inf")
|
||||
else:
|
||||
_deprecated_key_cache.pop(hashed_token, None)
|
||||
active_token_id = None
|
||||
cache_expires_at_ts = 0.0
|
||||
revoke_at_ts = 0.0
|
||||
|
||||
if (
|
||||
active_token_id is not None
|
||||
and now_ts < cache_expires_at_ts
|
||||
and now_ts < revoke_at_ts
|
||||
):
|
||||
return active_token_id
|
||||
else:
|
||||
_deprecated_key_cache.pop(hashed_token, None)
|
||||
|
|
@ -2414,12 +2429,17 @@ async def _lookup_deprecated_key(
|
|||
"token": hashed_token,
|
||||
"revoke_at": {"gt": now},
|
||||
},
|
||||
select={"active_token_id": True},
|
||||
select={"active_token_id": True, "revoke_at": True},
|
||||
)
|
||||
if deprecated_row and deprecated_row.active_token_id:
|
||||
if (
|
||||
deprecated_row
|
||||
and deprecated_row.active_token_id
|
||||
and deprecated_row.revoke_at is not None
|
||||
):
|
||||
_deprecated_key_cache[hashed_token] = (
|
||||
deprecated_row.active_token_id,
|
||||
now_ts + _DEPRECATED_KEY_CACHE_TTL_SECONDS,
|
||||
deprecated_row.revoke_at.timestamp(),
|
||||
)
|
||||
return deprecated_row.active_token_id
|
||||
# Only cache positive results; negative lookups are fast on indexed columns
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@ import asyncio
|
|||
import json
|
||||
import os
|
||||
import sys
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Tuple
|
||||
from unittest.mock import ANY, AsyncMock, MagicMock, patch
|
||||
|
||||
|
|
@ -903,6 +904,43 @@ def test_proxy_admin_jwt_auth_handles_no_team_object():
|
|||
assert result.end_user_id is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_lookup_deprecated_key_uses_cache_on_subsequent_requests():
|
||||
"""
|
||||
Regression test for deprecated key auth during grace period.
|
||||
|
||||
The in-memory cache entry must be readable on the second lookup; otherwise
|
||||
rotated keys intermittently fail auth with token_not_found_in_db.
|
||||
"""
|
||||
from types import SimpleNamespace
|
||||
|
||||
from litellm.proxy.utils import _deprecated_key_cache, _lookup_deprecated_key
|
||||
|
||||
hashed_old_key = "old-key-hash"
|
||||
hashed_new_key = "new-key-hash"
|
||||
revoke_at = datetime.now(timezone.utc) + timedelta(minutes=10)
|
||||
|
||||
_deprecated_key_cache.clear()
|
||||
|
||||
mock_db = MagicMock()
|
||||
mock_db.litellm_deprecatedverificationtoken.find_first = AsyncMock(
|
||||
return_value=SimpleNamespace(
|
||||
active_token_id=hashed_new_key,
|
||||
revoke_at=revoke_at,
|
||||
)
|
||||
)
|
||||
|
||||
try:
|
||||
first_lookup = await _lookup_deprecated_key(mock_db, hashed_old_key)
|
||||
second_lookup = await _lookup_deprecated_key(mock_db, hashed_old_key)
|
||||
finally:
|
||||
_deprecated_key_cache.clear()
|
||||
|
||||
assert first_lookup == hashed_new_key
|
||||
assert second_lookup == hashed_new_key
|
||||
mock_db.litellm_deprecatedverificationtoken.find_first.assert_awaited_once()
|
||||
|
||||
|
||||
class TestJWTOAuth2Coexistence:
|
||||
"""
|
||||
Test that JWT and OAuth2 auth can coexist on the same instance.
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue