fix(jwt): enforce no-match policy when prisma_client is None

The early `if prisma_client is None: return None` guard ran before the
no-match policy check, silently bypassing REJECT and AUTO_REGISTER — every
JWT client fell through to team auth regardless of configuration.

Fix: treat prisma_client=None as a definitive DB miss and fall through to the
same policy block as a real miss. REJECT now raises 403, AUTO_REGISTER raises
500 with a clear message (can't create keys without a DB), FALLBACK_TEAM_MAPPING
returns None unchanged.

Adds three tests: REJECT/403 with no DB, FALLBACK returns None with no DB,
AUTO_REGISTER/500 with no DB.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
shivam 2026-04-11 13:29:20 -07:00
parent f5fabe1fe4
commit 86d8147df5
No known key found for this signature in database
2 changed files with 107 additions and 9 deletions

View file

@ -646,14 +646,15 @@ async def _resolve_jwt_to_virtual_key(
proxy_logging_obj=proxy_logging_obj,
)
if prisma_client is None:
return 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,
)
# 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).
token_hash: Optional[str] = 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,
)
if token_hash is not None:
await user_api_key_cache.async_set_cache(
@ -669,7 +670,7 @@ async def _resolve_jwt_to_virtual_key(
proxy_logging_obj=proxy_logging_obj,
)
# No mapping found — apply no-match policy
# No mapping found (DB miss or no DB) — apply no-match policy.
behavior = jwt_handler.litellm_jwtauth.unregistered_jwt_client_behavior
if behavior == UnregisteredJWTClientBehavior.REJECT:
@ -686,6 +687,14 @@ async def _resolve_jwt_to_virtual_key(
)
if behavior == UnregisteredJWTClientBehavior.AUTO_REGISTER:
if prisma_client is None:
raise HTTPException(
status_code=500,
detail=(
"JWT Key Mapping: AUTO_REGISTER requires a database connection. "
"Configure a database or change unregistered_jwt_client_behavior."
),
)
return await _auto_register_jwt_mapping(
virtual_key_claim_field=virtual_key_claim_field,
claim_value=str(claim_value),

View file

@ -731,6 +731,95 @@ async def test_auto_register_race_condition_unique_conflict():
)
# ──────────────────────────────────────────────
# Tests: prisma_client=None does not bypass no-match policy
# ──────────────────────────────────────────────
@pytest.mark.asyncio
async def test_reject_behavior_enforced_when_prisma_client_is_none():
"""
When prisma_client is None and behavior is REJECT, a 403 must be raised
not silently fallen through to team auth.
"""
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.REJECT,
)
jwt_claims = {"email": "unknown@example.com"}
user_api_key_cache = DualCache()
with pytest.raises(HTTPException) as exc_info:
await _resolve_jwt_to_virtual_key(
jwt_claims=jwt_claims,
jwt_handler=jwt_handler,
prisma_client=None, # no DB
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=None,
)
assert exc_info.value.status_code == 403
assert "unknown@example.com" in exc_info.value.detail
@pytest.mark.asyncio
async def test_fallback_team_mapping_returns_none_when_prisma_client_is_none():
"""
When prisma_client is None and behavior is FALLBACK_TEAM_MAPPING, the
function must return None (fall through to team auth) not raise.
"""
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.FALLBACK_TEAM_MAPPING,
)
jwt_claims = {"email": "anyone@example.com"}
result = await _resolve_jwt_to_virtual_key(
jwt_claims=jwt_claims,
jwt_handler=jwt_handler,
prisma_client=None,
user_api_key_cache=DualCache(),
parent_otel_span=None,
proxy_logging_obj=None,
)
assert result is None
@pytest.mark.asyncio
async def test_auto_register_raises_500_when_prisma_client_is_none():
"""
AUTO_REGISTER without a DB connection must raise HTTP 500 with a clear
message it cannot create keys without a database.
"""
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,
)
jwt_claims = {"sub": "new-user-42"}
with pytest.raises(HTTPException) as exc_info:
await _resolve_jwt_to_virtual_key(
jwt_claims=jwt_claims,
jwt_handler=jwt_handler,
prisma_client=None,
user_api_key_cache=DualCache(),
parent_otel_span=None,
proxy_logging_obj=None,
)
assert exc_info.value.status_code == 500
assert "AUTO_REGISTER requires a database" in exc_info.value.detail
# ──────────────────────────────────────────────
# Tests: backward-compat alias jwt_client_id_field
# ──────────────────────────────────────────────