mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-22 00:31:44 +00:00
fix(jwt): fix AUTO_REGISTER sentinel bypass, race condition, and inline import comment
- AUTO_REGISTER now evicts stale __NO_MAPPING__ sentinel instead of silently returning None when cached under a prior fallback_team_mapping config - Race condition in _auto_register_jwt_mapping: catch P2002 unique-constraint violation on concurrent creates, fetch the winning mapping, proceed cleanly - Added comment on inline generate_key_helper_fn import explaining the circular dependency (key_management_endpoints imports user_api_key_auth at line 51) - 3 new tests: stale sentinel eviction, race condition winner fallback, and the existing auto_register happy path Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
shivam
parent
bef94c74cd
commit
5daf7494c7
2 changed files with 173 additions and 11 deletions
|
|
@ -512,7 +512,14 @@ async def _auto_register_jwt_mapping(
|
|||
"""
|
||||
Auto-register: create a new virtual key + mapping for an unrecognised JWT claim value.
|
||||
The new key carries no model/budget restrictions; admins can tighten it later.
|
||||
|
||||
Race safety: if two concurrent requests both reach here simultaneously (both saw
|
||||
no mapping in the DB), one will win the unique-constraint race on
|
||||
litellm_jwtkeymapping. The loser catches the conflict, fetches the winner's
|
||||
mapping, and proceeds — no orphaned keys and no error surfaced to the caller.
|
||||
"""
|
||||
# Inline import required: key_management_endpoints imports user_api_key_auth
|
||||
# (line 51) so a module-level import here would create a circular dependency.
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
generate_key_helper_fn,
|
||||
)
|
||||
|
|
@ -527,15 +534,38 @@ async def _auto_register_jwt_mapping(
|
|||
)
|
||||
token_hash = key_data["token"]
|
||||
|
||||
await prisma_client.db.litellm_jwtkeymapping.create(
|
||||
data={
|
||||
"jwt_claim_name": virtual_key_claim_field,
|
||||
"jwt_claim_value": claim_value,
|
||||
"token": token_hash,
|
||||
"created_by": "auto_register",
|
||||
"updated_by": "auto_register",
|
||||
}
|
||||
)
|
||||
try:
|
||||
await prisma_client.db.litellm_jwtkeymapping.create(
|
||||
data={
|
||||
"jwt_claim_name": virtual_key_claim_field,
|
||||
"jwt_claim_value": claim_value,
|
||||
"token": token_hash,
|
||||
"created_by": "auto_register",
|
||||
"updated_by": "auto_register",
|
||||
}
|
||||
)
|
||||
except Exception as e:
|
||||
error_str = str(e).lower()
|
||||
if "unique" in error_str or "p2002" in error_str:
|
||||
# A concurrent request won the race — fetch the winning mapping and
|
||||
# use its token. The key we just generated is orphaned but harmless;
|
||||
# it will be excluded from spend tracking since nothing maps to it.
|
||||
verbose_proxy_logger.debug(
|
||||
"JWT Key Mapping (auto_register): unique conflict on create — "
|
||||
"fetching winner's mapping for %s='%s'.",
|
||||
virtual_key_claim_field,
|
||||
claim_value,
|
||||
)
|
||||
token_hash = await get_jwt_key_mapping_object(
|
||||
jwt_claim_name=virtual_key_claim_field,
|
||||
jwt_claim_value=claim_value,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
if token_hash is None:
|
||||
# Should not happen, but guard against a delete racing our fetch.
|
||||
return None
|
||||
else:
|
||||
raise
|
||||
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=cache_key,
|
||||
|
|
@ -544,8 +574,9 @@ async def _auto_register_jwt_mapping(
|
|||
)
|
||||
|
||||
verbose_proxy_logger.info(
|
||||
f"JWT Key Mapping (auto_register): created new virtual key for "
|
||||
f"{virtual_key_claim_field}='{claim_value}'."
|
||||
"JWT Key Mapping (auto_register): created new virtual key for %s='%s'.",
|
||||
virtual_key_claim_field,
|
||||
claim_value,
|
||||
)
|
||||
|
||||
return await get_key_object(
|
||||
|
|
@ -591,6 +622,20 @@ 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 and prisma_client is not None:
|
||||
# Stale sentinel written under a prior fallback_team_mapping config —
|
||||
# evict it and auto-register now that the policy has changed.
|
||||
await user_api_key_cache.async_delete_cache(cache_key)
|
||||
return await _auto_register_jwt_mapping(
|
||||
virtual_key_claim_field=virtual_key_claim_field,
|
||||
claim_value=str(claim_value),
|
||||
jwt_handler=jwt_handler,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
cache_key=cache_key,
|
||||
)
|
||||
return None
|
||||
elif cached_mapping is not None:
|
||||
return await get_key_object(
|
||||
|
|
|
|||
|
|
@ -558,6 +558,123 @@ async def test_auto_register_creates_key_and_mapping():
|
|||
assert cached == "hashed_auto_key"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auto_register_triggers_on_stale_no_mapping_sentinel():
|
||||
"""
|
||||
If the cache holds a stale __NO_MAPPING__ sentinel (written under a prior
|
||||
fallback_team_mapping config) and behavior is now AUTO_REGISTER, the sentinel
|
||||
must be evicted and auto-registration must run — not silently return None.
|
||||
"""
|
||||
from litellm.proxy._types import UnregisteredJWTClientBehavior
|
||||
|
||||
jwt_handler = JWTHandler()
|
||||
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
|
||||
virtual_key_claim_field="email",
|
||||
unregistered_jwt_client_behavior=UnregisteredJWTClientBehavior.AUTO_REGISTER,
|
||||
virtual_key_mapping_cache_ttl=300,
|
||||
)
|
||||
jwt_claims = {"email": "alice@corp.com"}
|
||||
|
||||
prisma_client = MagicMock()
|
||||
prisma_client.db.litellm_jwtkeymapping.find_first = AsyncMock(return_value=None)
|
||||
prisma_client.db.litellm_jwtkeymapping.create = AsyncMock()
|
||||
|
||||
user_api_key_cache = DualCache()
|
||||
# Seed the stale sentinel
|
||||
await user_api_key_cache.async_set_cache(
|
||||
"jwt_key_mapping:email:alice@corp.com", "__NO_MAPPING__"
|
||||
)
|
||||
|
||||
mock_key_obj = UserAPIKeyAuth(token="hashed_auto_key", team_id=None)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.auth.user_api_key_auth.get_key_object", new_callable=AsyncMock
|
||||
) as mock_get_key, patch(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_gen_key:
|
||||
mock_gen_key.return_value = {"token": "hashed_auto_key", "key": "sk-auto-key"}
|
||||
mock_get_key.return_value = mock_key_obj
|
||||
|
||||
result = await _resolve_jwt_to_virtual_key(
|
||||
jwt_claims=jwt_claims,
|
||||
jwt_handler=jwt_handler,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=None,
|
||||
proxy_logging_obj=None,
|
||||
)
|
||||
|
||||
# Must have auto-registered, not returned None
|
||||
assert result == mock_key_obj
|
||||
prisma_client.db.litellm_jwtkeymapping.create.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auto_register_race_condition_unique_conflict():
|
||||
"""
|
||||
If two concurrent requests both call _auto_register_jwt_mapping and the
|
||||
second hits a unique-constraint violation on create, it must fall back to
|
||||
fetching the winner's mapping — no error surfaced to the caller.
|
||||
"""
|
||||
from litellm.proxy.auth.user_api_key_auth import _auto_register_jwt_mapping
|
||||
from litellm.proxy._types import UnregisteredJWTClientBehavior
|
||||
|
||||
jwt_handler = JWTHandler()
|
||||
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
|
||||
virtual_key_claim_field="sub",
|
||||
unregistered_jwt_client_behavior=UnregisteredJWTClientBehavior.AUTO_REGISTER,
|
||||
virtual_key_mapping_cache_ttl=300,
|
||||
)
|
||||
|
||||
prisma_client = MagicMock()
|
||||
prisma_client.db.litellm_jwtkeymapping.create = AsyncMock(
|
||||
side_effect=Exception("Unique constraint failed (P2002)")
|
||||
)
|
||||
# Simulate the winner's mapping already in DB after the conflict
|
||||
winner_mapping = MagicMock()
|
||||
winner_mapping.token = "winner_token_hash"
|
||||
winner_mapping.is_active = True
|
||||
prisma_client.db.litellm_jwtkeymapping.find_first = AsyncMock(
|
||||
return_value=winner_mapping
|
||||
)
|
||||
|
||||
user_api_key_cache = DualCache()
|
||||
mock_key_obj = UserAPIKeyAuth(token="winner_token_hash", team_id=None)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.auth.user_api_key_auth.get_key_object", new_callable=AsyncMock
|
||||
) as mock_get_key, patch(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn",
|
||||
new_callable=AsyncMock,
|
||||
return_value={"token": "loser_token_hash", "key": "sk-loser"},
|
||||
):
|
||||
mock_get_key.return_value = mock_key_obj
|
||||
|
||||
result = await _auto_register_jwt_mapping(
|
||||
virtual_key_claim_field="sub",
|
||||
claim_value="user-42",
|
||||
jwt_handler=jwt_handler,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=None,
|
||||
proxy_logging_obj=None,
|
||||
cache_key="jwt_key_mapping:sub:user-42",
|
||||
)
|
||||
|
||||
assert result == mock_key_obj
|
||||
# Cache should hold the winner's token, not the loser's
|
||||
cached = await user_api_key_cache.async_get_cache("jwt_key_mapping:sub:user-42")
|
||||
assert cached == "winner_token_hash"
|
||||
mock_get_key.assert_called_once_with(
|
||||
hashed_token="winner_token_hash",
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=None,
|
||||
proxy_logging_obj=None,
|
||||
)
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────
|
||||
# Tests: backward-compat alias jwt_client_id_field
|
||||
# ──────────────────────────────────────────────
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue