mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-15 23:31:29 +00:00
Merge remote-tracking branch 'origin/main' into litellm_/release-version-bump-940c92
This commit is contained in:
commit
818fd7bbb1
13 changed files with 556 additions and 43 deletions
|
|
@ -0,0 +1,18 @@
|
|||
-- DropIndex
|
||||
DROP INDEX IF EXISTS "LiteLLM_JWTKeyMapping_jwt_claim_name_jwt_claim_value_is_act_idx";
|
||||
|
||||
-- DropIndex
|
||||
DROP INDEX IF EXISTS "LiteLLM_JWTKeyMapping_jwt_claim_name_jwt_claim_value_key";
|
||||
|
||||
-- AlterTable
|
||||
-- NOT NULL DEFAULT '' (not nullable): Postgres unique constraints treat every
|
||||
-- NULL as distinct, so a nullable column would let multiple unscoped mappings
|
||||
-- collide on the same claim without a constraint violation. The constant
|
||||
-- default is a fast, metadata-only backfill for existing rows, not a rewrite.
|
||||
ALTER TABLE "LiteLLM_JWTKeyMapping" ADD COLUMN IF NOT EXISTS "jwt_issuer" TEXT NOT NULL DEFAULT '';
|
||||
|
||||
-- CreateIndex
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_JWTKeyMapping_jwt_issuer_jwt_claim_name_jwt_claim_v_idx" ON "LiteLLM_JWTKeyMapping"("jwt_issuer", "jwt_claim_name", "jwt_claim_value", "is_active");
|
||||
|
||||
-- CreateIndex
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_JWTKeyMapping_jwt_issuer_jwt_claim_name_jwt_claim_v_key" ON "LiteLLM_JWTKeyMapping"("jwt_issuer", "jwt_claim_name", "jwt_claim_value");
|
||||
|
|
@ -487,6 +487,10 @@ model LiteLLM_VerificationToken {
|
|||
|
||||
model LiteLLM_JWTKeyMapping {
|
||||
id String @id @default(uuid())
|
||||
jwt_issuer String @default("") // Scopes the mapping to one configured issuer; "" matches any issuer.
|
||||
// Not nullable: Postgres unique constraints treat every NULL as
|
||||
// distinct, so a nullable column would let multiple unscoped
|
||||
// mappings collide on the same claim without a constraint violation.
|
||||
jwt_claim_name String // e.g. "sub", "email"
|
||||
jwt_claim_value String // The claim value to match
|
||||
token String // Hashed virtual key (FK)
|
||||
|
|
@ -499,8 +503,8 @@ model LiteLLM_JWTKeyMapping {
|
|||
|
||||
litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token], onDelete: Cascade)
|
||||
|
||||
@@unique([jwt_claim_name, jwt_claim_value])
|
||||
@@index([jwt_claim_name, jwt_claim_value, is_active])
|
||||
@@unique([jwt_issuer, jwt_claim_name, jwt_claim_value])
|
||||
@@index([jwt_issuer, jwt_claim_name, jwt_claim_value, is_active])
|
||||
}
|
||||
|
||||
// Deprecated keys during grace period - allows old key to work until revoke_at
|
||||
|
|
|
|||
|
|
@ -15226,6 +15226,17 @@
|
|||
"title": "Jwt Claim Value",
|
||||
"type": "string"
|
||||
},
|
||||
"jwt_issuer": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "string"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"title": "Jwt Issuer"
|
||||
},
|
||||
"key": {
|
||||
"title": "Key",
|
||||
"type": "string"
|
||||
|
|
@ -15310,6 +15321,17 @@
|
|||
"title": "Jwt Claim Value",
|
||||
"type": "string"
|
||||
},
|
||||
"jwt_issuer": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "string"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"title": "Jwt Issuer"
|
||||
},
|
||||
"updated_at": {
|
||||
"format": "date-time",
|
||||
"title": "Updated At",
|
||||
|
|
@ -15366,6 +15388,17 @@
|
|||
],
|
||||
"title": "Is Active"
|
||||
},
|
||||
"jwt_issuer": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "string"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"title": "Jwt Issuer"
|
||||
},
|
||||
"key": {
|
||||
"anyOf": [
|
||||
{
|
||||
|
|
|
|||
|
|
@ -4486,12 +4486,14 @@ class CreateJWTKeyMappingRequest(LiteLLMPydanticObjectBase):
|
|||
jwt_claim_name: str
|
||||
jwt_claim_value: str
|
||||
key: str
|
||||
jwt_issuer: str | None = None
|
||||
description: str | None = None
|
||||
|
||||
|
||||
class UpdateJWTKeyMappingRequest(LiteLLMPydanticObjectBase):
|
||||
id: str
|
||||
key: str | None = None
|
||||
jwt_issuer: str | None = None
|
||||
description: str | None = None
|
||||
is_active: bool | None = None
|
||||
|
||||
|
|
@ -4502,6 +4504,7 @@ class DeleteJWTKeyMappingRequest(LiteLLMPydanticObjectBase):
|
|||
|
||||
class JWTKeyMappingResponse(LiteLLMPydanticObjectBase):
|
||||
id: str
|
||||
jwt_issuer: str | None = None
|
||||
jwt_claim_name: str
|
||||
jwt_claim_value: str
|
||||
description: str | None = None
|
||||
|
|
|
|||
|
|
@ -148,6 +148,7 @@ class _PrismaDictableRow(Protocol):
|
|||
|
||||
class _PrismaJWTKeyMappingRow(Protocol):
|
||||
token: str
|
||||
jwt_issuer: str
|
||||
jwt_claim_name: str
|
||||
jwt_claim_value: str
|
||||
|
||||
|
|
@ -3601,9 +3602,18 @@ 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}"
|
||||
def jwt_key_mapping_cache_key(jwt_claim_name: str, jwt_claim_value: str, jwt_issuer: str | None = None) -> str:
|
||||
"""Cache key under which a JWT-claim-to-key mapping is stored, scoped to one
|
||||
issuer (or the issuer-agnostic/global scope when ``jwt_issuer`` is falsy).
|
||||
|
||||
Scoped by issuer (when one is configured) so a cached hit or ``__NO_MAPPING__`` miss
|
||||
for one issuer's claim value can never be served to a different issuer whose claim
|
||||
value happens to collide. Unchanged for the global scope, keeping the single-issuer
|
||||
(no ``litellm_jwtauth.issuers`` configured) cache key format stable across this fix.
|
||||
"""
|
||||
if not jwt_issuer:
|
||||
return f"jwt_key_mapping:{jwt_claim_name}:{jwt_claim_value}"
|
||||
return f"jwt_key_mapping:{jwt_issuer}:{jwt_claim_name}:{jwt_claim_value}"
|
||||
|
||||
|
||||
@log_db_metrics
|
||||
|
|
@ -3615,7 +3625,7 @@ async def get_jwt_key_mapping_cache_keys_for_token(
|
|||
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)
|
||||
return tuple(jwt_key_mapping_cache_key(m.jwt_claim_name, m.jwt_claim_value, m.jwt_issuer) for m in mappings)
|
||||
|
||||
|
||||
@log_db_metrics
|
||||
|
|
@ -3623,9 +3633,14 @@ async def get_jwt_key_mapping_object(
|
|||
jwt_claim_name: str,
|
||||
jwt_claim_value: str,
|
||||
prisma_client: PrismaClient,
|
||||
jwt_issuer: str | None = None,
|
||||
) -> str | None:
|
||||
"""
|
||||
Lookup a JWT-to-virtual-key mapping from the database.
|
||||
Lookup a JWT-to-virtual-key mapping from the database for one exact scope:
|
||||
``jwt_issuer`` (or the global/issuer-agnostic scope when falsy). Does not fall
|
||||
back to the global scope itself -- a caller that wants "issuer-scoped mapping,
|
||||
else the global one" queries both scopes itself, so each result can be cached
|
||||
under its own scope's key (see ``_resolve_jwt_to_virtual_key``).
|
||||
|
||||
Returns the hashed token (str) if a matching active mapping is found, else None.
|
||||
"""
|
||||
|
|
@ -3633,6 +3648,7 @@ async def get_jwt_key_mapping_object(
|
|||
where={
|
||||
"jwt_claim_name": jwt_claim_name,
|
||||
"jwt_claim_value": jwt_claim_value,
|
||||
"jwt_issuer": jwt_issuer or "",
|
||||
"is_active": True,
|
||||
}
|
||||
)
|
||||
|
|
|
|||
|
|
@ -269,6 +269,17 @@ class _TokenTeamModels(Protocol):
|
|||
def team_models(self) -> list[str]: ...
|
||||
|
||||
|
||||
class _RawCacheRead(Protocol):
|
||||
async def async_get_cache(self, *, key: str) -> object: ...
|
||||
|
||||
|
||||
def _raw_cache(cache: _RawCacheRead) -> _RawCacheRead:
|
||||
"""View an untyped cache object's ``async_get_cache`` as returning ``object``
|
||||
instead of ``Any``, so a caller can ``isinstance``-narrow it without paying
|
||||
the ``reportAny`` cost of the underlying (unannotated) cache implementation."""
|
||||
return cache
|
||||
|
||||
|
||||
def _token_team_models(valid_token: _TokenTeamModels) -> list[str]:
|
||||
return valid_token.team_models
|
||||
|
||||
|
|
@ -842,6 +853,7 @@ class _PendingAutoRegister(NamedTuple):
|
|||
claim_field: str
|
||||
claim_value: str
|
||||
cache_key: str
|
||||
jwt_issuer: str | None = None
|
||||
|
||||
|
||||
async def _auto_register_jwt_mapping(
|
||||
|
|
@ -853,6 +865,7 @@ async def _auto_register_jwt_mapping(
|
|||
parent_otel_span: Span | None,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
cache_key: str,
|
||||
jwt_issuer: str | None = None,
|
||||
team_id: str | None = None,
|
||||
user_id: str | None = None,
|
||||
org_id: str | None = None,
|
||||
|
|
@ -905,6 +918,7 @@ async def _auto_register_jwt_mapping(
|
|||
try:
|
||||
await prisma_client.db.litellm_jwtkeymapping.create(
|
||||
data={
|
||||
"jwt_issuer": jwt_issuer or "",
|
||||
"jwt_claim_name": virtual_key_claim_field,
|
||||
"jwt_claim_value": claim_value,
|
||||
"token": token_hash,
|
||||
|
|
@ -939,6 +953,7 @@ async def _auto_register_jwt_mapping(
|
|||
jwt_claim_name=virtual_key_claim_field,
|
||||
jwt_claim_value=claim_value,
|
||||
prisma_client=prisma_client,
|
||||
jwt_issuer=jwt_issuer,
|
||||
)
|
||||
if token_hash is None:
|
||||
# The winner's mapping vanished between the unique-constraint
|
||||
|
|
@ -983,6 +998,43 @@ async def _auto_register_jwt_mapping(
|
|||
return auto_registered_key
|
||||
|
||||
|
||||
async def _lookup_jwt_mapping_token_hash(
|
||||
prisma_client: PrismaClient,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
virtual_key_claim_field: str,
|
||||
claim_value: str,
|
||||
normalized_issuer: str | None,
|
||||
cache_key: str,
|
||||
ttl: float,
|
||||
) -> str | None:
|
||||
issuer_scoped: Final = await get_jwt_key_mapping_object(
|
||||
jwt_claim_name=virtual_key_claim_field,
|
||||
jwt_claim_value=claim_value,
|
||||
prisma_client=prisma_client,
|
||||
jwt_issuer=normalized_issuer,
|
||||
)
|
||||
if issuer_scoped is not None:
|
||||
await user_api_key_cache.async_set_cache(key=cache_key, value=issuer_scoped, ttl=ttl)
|
||||
return issuer_scoped
|
||||
if normalized_issuer is None:
|
||||
return None
|
||||
# Another issuer may have already resolved (and cached) this same
|
||||
# global mapping -- check its cache entry before re-querying the DB.
|
||||
global_cache_key: Final = jwt_key_mapping_cache_key(virtual_key_claim_field, claim_value)
|
||||
cached_global: Final = await _raw_cache(user_api_key_cache).async_get_cache(key=global_cache_key)
|
||||
if isinstance(cached_global, str) and cached_global != "__NO_MAPPING__":
|
||||
return cached_global
|
||||
global_row: Final = await get_jwt_key_mapping_object(
|
||||
jwt_claim_name=virtual_key_claim_field,
|
||||
jwt_claim_value=claim_value,
|
||||
prisma_client=prisma_client,
|
||||
jwt_issuer=None,
|
||||
)
|
||||
if global_row is not None:
|
||||
await user_api_key_cache.async_set_cache(key=global_cache_key, value=global_row, ttl=ttl)
|
||||
return global_row
|
||||
|
||||
|
||||
async def _resolve_jwt_to_virtual_key(
|
||||
jwt_claims: dict,
|
||||
jwt_handler: JWTHandler,
|
||||
|
|
@ -1041,7 +1093,7 @@ async def _resolve_jwt_to_virtual_key(
|
|||
)
|
||||
return None
|
||||
|
||||
cache_key: Final = jwt_key_mapping_cache_key(virtual_key_claim_field, str(claim_value))
|
||||
cache_key: Final = jwt_key_mapping_cache_key(virtual_key_claim_field, str(claim_value), normalized_issuer)
|
||||
raw_cached_mapping: Final = await user_api_key_cache.async_get_cache(cache_key)
|
||||
sentinel_written_by_this_policy: Final = behavior == UnregisteredJWTClientBehavior.AUTO_REGISTER
|
||||
cached_mapping: Final = (
|
||||
|
|
@ -1081,6 +1133,7 @@ async def _resolve_jwt_to_virtual_key(
|
|||
claim_field=virtual_key_claim_field,
|
||||
claim_value=str(claim_value),
|
||||
cache_key=cache_key,
|
||||
jwt_issuer=normalized_issuer,
|
||||
)
|
||||
return None
|
||||
elif cached_mapping is not None:
|
||||
|
|
@ -1094,21 +1147,30 @@ async def _resolve_jwt_to_virtual_key(
|
|||
)
|
||||
|
||||
# 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: str | None = 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),
|
||||
# miss (no DB → no mapping can exist → apply no-match policy below). An
|
||||
# issuer-scoped row wins; falling back to the global (no-issuer) row keeps
|
||||
# mappings created before issuer scoping existed working for every issuer.
|
||||
# Each tier is cached under ITS OWN key (the global tier under the
|
||||
# issuer-less cache key, not under `cache_key`/this issuer's key) so that
|
||||
# updating or deleting either row invalidates exactly the cache entries it
|
||||
# can affect. Caching a global-row hit under the requesting issuer's key
|
||||
# would leave every OTHER issuer that had fallen back to that same global
|
||||
# mapping serving its stale token until TTL after the row changes.
|
||||
token_hash: Final = (
|
||||
await _lookup_jwt_mapping_token_hash(
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
|
||||
if token_hash is not None:
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=cache_key,
|
||||
value=token_hash,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
virtual_key_claim_field=virtual_key_claim_field,
|
||||
claim_value=str(claim_value),
|
||||
normalized_issuer=normalized_issuer,
|
||||
cache_key=cache_key,
|
||||
ttl=jwt_handler.litellm_jwtauth.virtual_key_mapping_cache_ttl,
|
||||
)
|
||||
if prisma_client is not None
|
||||
else None
|
||||
)
|
||||
|
||||
if token_hash is not None:
|
||||
return IdentityStore.key_from_principal(
|
||||
await IdentityStore(
|
||||
prisma_client,
|
||||
|
|
@ -1149,6 +1211,7 @@ async def _resolve_jwt_to_virtual_key(
|
|||
claim_field=virtual_key_claim_field,
|
||||
claim_value=str(claim_value),
|
||||
cache_key=cache_key,
|
||||
jwt_issuer=normalized_issuer,
|
||||
)
|
||||
|
||||
# FALLBACK_TEAM_MAPPING (default): cache the miss and return None so the
|
||||
|
|
@ -1641,6 +1704,7 @@ async def _user_api_key_auth_builder(
|
|||
parent_otel_span=parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
cache_key=pending_auto_register.cache_key,
|
||||
jwt_issuer=pending_auto_register.jwt_issuer,
|
||||
team_id=team_id,
|
||||
user_id=user_id,
|
||||
org_id=org_id,
|
||||
|
|
|
|||
|
|
@ -28,6 +28,9 @@ class _JWTKeyMappingRecord(Protocol):
|
|||
@property
|
||||
def id(self) -> str: ...
|
||||
|
||||
@property
|
||||
def jwt_issuer(self) -> str: ...
|
||||
|
||||
@property
|
||||
def jwt_claim_name(self) -> str: ...
|
||||
|
||||
|
|
@ -78,6 +81,7 @@ def _to_response(mapping: _JWTKeyMappingRecord) -> JWTKeyMappingResponse:
|
|||
"""Convert a Prisma mapping object to a safe response (no hashed token)."""
|
||||
return JWTKeyMappingResponse(
|
||||
id=mapping.id,
|
||||
jwt_issuer=mapping.jwt_issuer or None,
|
||||
jwt_claim_name=mapping.jwt_claim_name,
|
||||
jwt_claim_value=mapping.jwt_claim_value,
|
||||
description=mapping.description,
|
||||
|
|
@ -109,6 +113,7 @@ async def create_jwt_key_mapping(
|
|||
try:
|
||||
hashed_key: Final = hash_token(data.key)
|
||||
create_data: Final = {
|
||||
"jwt_issuer": data.jwt_issuer or "",
|
||||
"jwt_claim_name": data.jwt_claim_name,
|
||||
"jwt_claim_value": data.jwt_claim_value,
|
||||
"token": hashed_key,
|
||||
|
|
@ -120,7 +125,7 @@ async def create_jwt_key_mapping(
|
|||
|
||||
new_mapping: Final = await _mapping_table(prisma_client).create(data=create_data)
|
||||
|
||||
cache_key: Final = jwt_key_mapping_cache_key(data.jwt_claim_name, data.jwt_claim_value)
|
||||
cache_key: Final = jwt_key_mapping_cache_key(data.jwt_claim_name, data.jwt_claim_value, data.jwt_issuer)
|
||||
await evict_and_broadcast(cache_keys=(cache_key,), user_api_key_cache=user_api_key_cache)
|
||||
|
||||
return _to_response(new_mapping)
|
||||
|
|
@ -131,7 +136,10 @@ async def create_jwt_key_mapping(
|
|||
if "unique" in error_str or "p2002" in error_str:
|
||||
raise HTTPException(
|
||||
status_code=409,
|
||||
detail=f"A mapping for claim '{data.jwt_claim_name}' = '{data.jwt_claim_value}' already exists.",
|
||||
detail=(
|
||||
f"A mapping for claim '{data.jwt_claim_name}' = '{data.jwt_claim_value}' "
|
||||
f"already exists for issuer '{data.jwt_issuer}'."
|
||||
),
|
||||
)
|
||||
if "foreign" in error_str or "p2003" in error_str:
|
||||
raise HTTPException(
|
||||
|
|
@ -161,6 +169,9 @@ async def update_jwt_key_mapping(
|
|||
update_data: Final = data.model_dump(exclude_unset=True, exclude={"id", "key"})
|
||||
if data.key is not None:
|
||||
update_data["token"] = hash_token(data.key)
|
||||
if "jwt_issuer" in update_data:
|
||||
# DB column is NOT NULL (see schema.prisma); "" is the global/unscoped sentinel.
|
||||
update_data["jwt_issuer"] = update_data["jwt_issuer"] or ""
|
||||
update_data["updated_by"] = user_api_key_dict.user_id
|
||||
|
||||
try:
|
||||
|
|
@ -178,9 +189,11 @@ async def update_jwt_key_mapping(
|
|||
# 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)
|
||||
old_cache_key: Final = jwt_key_mapping_cache_key(
|
||||
old_mapping.jwt_claim_name, old_mapping.jwt_claim_value, old_mapping.jwt_issuer
|
||||
)
|
||||
new_cache_key: Final = jwt_key_mapping_cache_key(
|
||||
updated_mapping.jwt_claim_name, updated_mapping.jwt_claim_value
|
||||
updated_mapping.jwt_claim_name, updated_mapping.jwt_claim_value, updated_mapping.jwt_issuer
|
||||
)
|
||||
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)
|
||||
|
|
@ -227,7 +240,9 @@ async def delete_jwt_key_mapping(
|
|||
|
||||
# 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)
|
||||
cache_key: Final = jwt_key_mapping_cache_key(
|
||||
old_mapping.jwt_claim_name, old_mapping.jwt_claim_value, old_mapping.jwt_issuer
|
||||
)
|
||||
await evict_and_broadcast(cache_keys=(cache_key,), user_api_key_cache=user_api_key_cache)
|
||||
return {"status": "success"}
|
||||
except HTTPException:
|
||||
|
|
|
|||
|
|
@ -487,6 +487,10 @@ model LiteLLM_VerificationToken {
|
|||
|
||||
model LiteLLM_JWTKeyMapping {
|
||||
id String @id @default(uuid())
|
||||
jwt_issuer String @default("") // Scopes the mapping to one configured issuer; "" matches any issuer.
|
||||
// Not nullable: Postgres unique constraints treat every NULL as
|
||||
// distinct, so a nullable column would let multiple unscoped
|
||||
// mappings collide on the same claim without a constraint violation.
|
||||
jwt_claim_name String // e.g. "sub", "email"
|
||||
jwt_claim_value String // The claim value to match
|
||||
token String // Hashed virtual key (FK)
|
||||
|
|
@ -499,8 +503,8 @@ model LiteLLM_JWTKeyMapping {
|
|||
|
||||
litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token], onDelete: Cascade)
|
||||
|
||||
@@unique([jwt_claim_name, jwt_claim_value])
|
||||
@@index([jwt_claim_name, jwt_claim_value, is_active])
|
||||
@@unique([jwt_issuer, jwt_claim_name, jwt_claim_value])
|
||||
@@index([jwt_issuer, jwt_claim_name, jwt_claim_value, is_active])
|
||||
}
|
||||
|
||||
// Deprecated keys during grace period - allows old key to work until revoke_at
|
||||
|
|
|
|||
|
|
@ -487,6 +487,10 @@ model LiteLLM_VerificationToken {
|
|||
|
||||
model LiteLLM_JWTKeyMapping {
|
||||
id String @id @default(uuid())
|
||||
jwt_issuer String @default("") // Scopes the mapping to one configured issuer; "" matches any issuer.
|
||||
// Not nullable: Postgres unique constraints treat every NULL as
|
||||
// distinct, so a nullable column would let multiple unscoped
|
||||
// mappings collide on the same claim without a constraint violation.
|
||||
jwt_claim_name String // e.g. "sub", "email"
|
||||
jwt_claim_value String // The claim value to match
|
||||
token String // Hashed virtual key (FK)
|
||||
|
|
@ -499,8 +503,8 @@ model LiteLLM_JWTKeyMapping {
|
|||
|
||||
litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token], onDelete: Cascade)
|
||||
|
||||
@@unique([jwt_claim_name, jwt_claim_value])
|
||||
@@index([jwt_claim_name, jwt_claim_value, is_active])
|
||||
@@unique([jwt_issuer, jwt_claim_name, jwt_claim_value])
|
||||
@@index([jwt_issuer, jwt_claim_name, jwt_claim_value, is_active])
|
||||
}
|
||||
|
||||
// Deprecated keys during grace period - allows old key to work until revoke_at
|
||||
|
|
|
|||
|
|
@ -91,6 +91,154 @@ async def test_jwt_to_virtual_key_mapping_resolution():
|
|||
prisma_client.db.litellm_jwtkeymapping.find_first.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_colliding_claim_value_from_another_issuer_does_not_resolve_to_the_wrong_virtual_key():
|
||||
"""LIT-7417: a mapping registered for one issuer must not answer a lookup from a
|
||||
DIFFERENT issuer whose claim value happens to collide, even though both issuers
|
||||
map the same claim field (``sub``) to a virtual key."""
|
||||
issuer_a = "https://issuer-a.example.com"
|
||||
issuer_b = "https://issuer-b.example.com"
|
||||
|
||||
jwt_handler = JWTHandler()
|
||||
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
|
||||
virtual_key_claim_field="sub", virtual_key_mapping_cache_ttl=3600
|
||||
)
|
||||
|
||||
rows = [
|
||||
{
|
||||
"jwt_issuer": issuer_b,
|
||||
"jwt_claim_name": "sub",
|
||||
"jwt_claim_value": "dev-alice",
|
||||
"token": "hashed-issuer-b-key",
|
||||
"is_active": True,
|
||||
}
|
||||
]
|
||||
|
||||
async def fake_find_first(where):
|
||||
for row in rows:
|
||||
if all(row.get(k) == v for k, v in where.items()):
|
||||
return MagicMock(**row)
|
||||
return None
|
||||
|
||||
prisma_client = MagicMock()
|
||||
prisma_client.db.litellm_jwtkeymapping.find_first = AsyncMock(side_effect=fake_find_first)
|
||||
|
||||
# Dependency-inject the resolved key via the cache (IdentityStore._resolve_key
|
||||
# reads it from here) instead of monkeypatching IdentityStore itself.
|
||||
user_api_key_cache = DualCache()
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key="hashed-issuer-b-key",
|
||||
value=UserAPIKeyAuth(token="hashed-issuer-b-key", team_id="issuer-b-team"),
|
||||
)
|
||||
|
||||
# The rightful owner: issuer-b's own claim resolves to its mapping.
|
||||
owner_result = await _resolve_jwt_to_virtual_key(
|
||||
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: issuer_b, "sub": "dev-alice"},
|
||||
jwt_handler=jwt_handler,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=None,
|
||||
proxy_logging_obj=None,
|
||||
)
|
||||
assert isinstance(owner_result, UserAPIKeyAuth)
|
||||
assert owner_result.token == "hashed-issuer-b-key"
|
||||
|
||||
# A validly-signed token from issuer-a carrying the SAME claim value must not
|
||||
# inherit issuer-b's mapping. Default behavior is fallback_team_mapping, so a
|
||||
# correctly-scoped miss returns None instead of resolving to issuer-b's key.
|
||||
colliding_result = await _resolve_jwt_to_virtual_key(
|
||||
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: issuer_a, "sub": "dev-alice"},
|
||||
jwt_handler=jwt_handler,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=None,
|
||||
proxy_logging_obj=None,
|
||||
)
|
||||
assert colliding_result is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_global_mapping_resolution_is_cached_under_the_global_key_not_the_requesting_issuer():
|
||||
"""LIT-7417: caching a global (unscoped) mapping's hit under the REQUESTING
|
||||
issuer's key would leave every issuer that falls back to it holding its own
|
||||
stale copy after the row is updated/deleted -- CRUD only evicts the cache key
|
||||
computed from the row's own scope (global), so a copy cached under some other
|
||||
issuer's key would keep resolving to the old token until TTL. Caching it under
|
||||
the global key instead means every issuer shares (and CRUD correctly evicts)
|
||||
the exact same entry."""
|
||||
issuer_a = "https://issuer-a.example.com"
|
||||
issuer_b = "https://issuer-b.example.com"
|
||||
|
||||
jwt_handler = JWTHandler()
|
||||
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
|
||||
issuers=[
|
||||
{
|
||||
"issuer": issuer_a,
|
||||
"jwks_url": f"{issuer_a}/jwks",
|
||||
"virtual_key_claim_field": "sub",
|
||||
"disable_audience_validation": True,
|
||||
},
|
||||
{
|
||||
"issuer": issuer_b,
|
||||
"jwks_url": f"{issuer_b}/jwks",
|
||||
"virtual_key_claim_field": "sub",
|
||||
"disable_audience_validation": True,
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
rows = [
|
||||
{
|
||||
"jwt_issuer": "",
|
||||
"jwt_claim_name": "sub",
|
||||
"jwt_claim_value": "legacy-user",
|
||||
"token": "hashed-legacy-key",
|
||||
"is_active": True,
|
||||
}
|
||||
]
|
||||
|
||||
async def fake_find_first(where):
|
||||
for row in rows:
|
||||
if all(row.get(k) == v for k, v in where.items()):
|
||||
return MagicMock(**row)
|
||||
return None
|
||||
|
||||
prisma_client = MagicMock()
|
||||
find_first = AsyncMock(side_effect=fake_find_first)
|
||||
prisma_client.db.litellm_jwtkeymapping.find_first = find_first
|
||||
|
||||
user_api_key_cache = DualCache()
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key="hashed-legacy-key",
|
||||
value=UserAPIKeyAuth(token="hashed-legacy-key", team_id="legacy-team"),
|
||||
)
|
||||
|
||||
resolved_a = await _resolve_jwt_to_virtual_key(
|
||||
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: issuer_a, "sub": "legacy-user"},
|
||||
jwt_handler=jwt_handler,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=None,
|
||||
proxy_logging_obj=None,
|
||||
)
|
||||
assert isinstance(resolved_a, UserAPIKeyAuth)
|
||||
assert find_first.await_count == 2 # issuer-a-scoped miss, then global hit
|
||||
|
||||
# issuer-b resolving the SAME global mapping must hit the cache issuer-a's
|
||||
# resolution populated, not issue a fresh DB query for the global row again.
|
||||
resolved_b = await _resolve_jwt_to_virtual_key(
|
||||
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: issuer_b, "sub": "legacy-user"},
|
||||
jwt_handler=jwt_handler,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=None,
|
||||
proxy_logging_obj=None,
|
||||
)
|
||||
assert isinstance(resolved_b, UserAPIKeyAuth)
|
||||
assert resolved_b.token == "hashed-legacy-key"
|
||||
assert find_first.await_count == 3 # +1 for issuer-b's own issuer-scoped miss; global tier served from cache
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_jwt_to_virtual_key_mapping_no_mapping():
|
||||
"""
|
||||
|
|
@ -223,6 +371,7 @@ def test_to_response_excludes_token():
|
|||
now = datetime.now(timezone.utc)
|
||||
mock_mapping = MagicMock()
|
||||
mock_mapping.id = "mapping-1"
|
||||
mock_mapping.jwt_issuer = None
|
||||
mock_mapping.jwt_claim_name = "email"
|
||||
mock_mapping.jwt_claim_value = "user@example.com"
|
||||
mock_mapping.token = "hashed_secret_value"
|
||||
|
|
@ -275,10 +424,12 @@ def _mock_mapping(
|
|||
id="mapping-1",
|
||||
claim_name="email",
|
||||
claim_value="user@example.com",
|
||||
issuer=None,
|
||||
):
|
||||
now = datetime.now(timezone.utc)
|
||||
m = MagicMock()
|
||||
m.id = id
|
||||
m.jwt_issuer = issuer
|
||||
m.jwt_claim_name = claim_name
|
||||
m.jwt_claim_value = claim_value
|
||||
m.token = "hashed_token"
|
||||
|
|
@ -485,6 +636,35 @@ async def test_create_success_returns_response_without_token():
|
|||
assert result.jwt_claim_name == "email"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_without_issuer_stores_empty_string_not_null():
|
||||
"""LIT-7417: the DB column is NOT NULL (see schema.prisma). Storing a real NULL
|
||||
for an unscoped mapping would let Postgres accept unlimited duplicate unscoped
|
||||
rows for the same claim (NULL is never equal to NULL in a unique constraint),
|
||||
so two mappings for the same claim value could point at two different keys with
|
||||
no conflict, and resolution would pick whichever one Postgres returns first."""
|
||||
from litellm.proxy._types import CreateJWTKeyMappingRequest
|
||||
|
||||
mock_prisma = _mock_prisma()
|
||||
mock_prisma.db.litellm_jwtkeymapping.create.return_value = _mock_mapping()
|
||||
mock_cache = AsyncMock()
|
||||
|
||||
data = CreateJWTKeyMappingRequest(jwt_claim_name="sub", jwt_claim_value="dev-alice", key="sk-test-key")
|
||||
|
||||
with (
|
||||
patch( # test-quality-ok: proxy_server module global is the endpoint's only injection point
|
||||
"litellm.proxy.proxy_server.prisma_client", mock_prisma
|
||||
),
|
||||
patch( # test-quality-ok: proxy_server module global is the endpoint's only injection point
|
||||
"litellm.proxy.proxy_server.user_api_key_cache", mock_cache
|
||||
),
|
||||
):
|
||||
await create_jwt_key_mapping(data=data, user_api_key_dict=_make_admin_auth())
|
||||
|
||||
sent_data = mock_prisma.db.litellm_jwtkeymapping.create.call_args.kwargs["data"]
|
||||
assert sent_data["jwt_issuer"] == ""
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────
|
||||
# Tests: unregistered_jwt_client_behavior
|
||||
# ──────────────────────────────────────────────
|
||||
|
|
|
|||
|
|
@ -32,7 +32,13 @@ from litellm.proxy._types import (
|
|||
JWTRoutingOverride,
|
||||
)
|
||||
from litellm.proxy.auth.handle_jwt import JWTHandler
|
||||
from litellm.proxy.auth.auth_checks import TeamNotFoundError, UserNotFoundError, get_key_object, _cache_key_object
|
||||
from litellm.proxy.auth.auth_checks import (
|
||||
TeamNotFoundError,
|
||||
UserNotFoundError,
|
||||
get_key_object,
|
||||
_cache_key_object,
|
||||
jwt_key_mapping_cache_key,
|
||||
)
|
||||
from litellm.proxy.auth.route_checks import RouteChecks
|
||||
from litellm.proxy.auth.user_api_key_auth import (
|
||||
_check_key_model_budget_with_fallback,
|
||||
|
|
@ -7948,13 +7954,38 @@ def _per_issuer_virtual_key_jwt_handler(
|
|||
|
||||
|
||||
def _fake_prisma_with_jwt_key_mapping(hashed_token: str | None) -> tuple[SimpleNamespace, AsyncMock]:
|
||||
"""Every ``find_first`` call (issuer-scoped or global fallback) resolves the same way."""
|
||||
find_first = AsyncMock(return_value=None if hashed_token is None else SimpleNamespace(token=hashed_token))
|
||||
prisma_client = SimpleNamespace(db=SimpleNamespace(litellm_jwtkeymapping=SimpleNamespace(find_first=find_first)))
|
||||
return prisma_client, find_first
|
||||
|
||||
|
||||
def _mapping_where(claim_name: str, claim_value: str) -> dict[str, str | bool]:
|
||||
return {"jwt_claim_name": claim_name, "jwt_claim_value": claim_value, "is_active": True}
|
||||
def _fake_prisma_jwt_key_mapping_table(rows: list[dict[str, object]]) -> tuple[SimpleNamespace, AsyncMock]:
|
||||
"""A ``find_first`` whose result depends on the ``where`` clause, like a real table.
|
||||
|
||||
Matches a row when every key present in ``where`` equals that key on the row --
|
||||
a key ``get_jwt_key_mapping_object`` omits (e.g. old, issuer-blind code never
|
||||
sending ``jwt_issuer``) does not constrain the match, exactly like Prisma.
|
||||
"""
|
||||
|
||||
async def _find_first(where: dict[str, object]) -> SimpleNamespace | None:
|
||||
for row in rows:
|
||||
if all(row.get(k) == v for k, v in where.items()):
|
||||
return SimpleNamespace(**row)
|
||||
return None
|
||||
|
||||
find_first = AsyncMock(side_effect=_find_first)
|
||||
prisma_client = SimpleNamespace(db=SimpleNamespace(litellm_jwtkeymapping=SimpleNamespace(find_first=find_first)))
|
||||
return prisma_client, find_first
|
||||
|
||||
|
||||
def _mapping_where(claim_name: str, claim_value: str, jwt_issuer: str | None) -> dict[str, str | bool]:
|
||||
return {
|
||||
"jwt_claim_name": claim_name,
|
||||
"jwt_claim_value": claim_value,
|
||||
"jwt_issuer": jwt_issuer or "",
|
||||
"is_active": True,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -7978,11 +8009,13 @@ async def test_per_issuer_virtual_key_claim_field_selects_the_issuer_mapping_for
|
|||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
find_first.assert_awaited_once_with(where=_mapping_where("sub", "svc-account-7"))
|
||||
# Issuer-scoped lookup hits on the first query, so no global fallback query runs.
|
||||
find_first.assert_awaited_once_with(where=_mapping_where("sub", "svc-account-7", ISSUER_TWO))
|
||||
assert isinstance(resolved, UserAPIKeyAuth)
|
||||
assert resolved.token == "hashed-mapped-key"
|
||||
assert resolved.team_id == "svc-team"
|
||||
assert await user_api_key_cache.async_get_cache("jwt_key_mapping:sub:svc-account-7") == "hashed-mapped-key"
|
||||
cache_key = jwt_key_mapping_cache_key("sub", "svc-account-7", ISSUER_TWO)
|
||||
assert await user_api_key_cache.async_get_cache(cache_key) == "hashed-mapped-key"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -8015,7 +8048,11 @@ async def test_per_issuer_reject_behavior_does_not_leak_into_the_team_issuer():
|
|||
|
||||
assert exc.value.status_code == 403
|
||||
assert "No registered mapping for sub='unknown-svc'" in str(exc.value.detail)
|
||||
find_first.assert_awaited_once_with(where=_mapping_where("sub", "unknown-svc"))
|
||||
# REJECT checks the issuer-scoped row first, then falls back to a global (NULL-issuer) row.
|
||||
assert [c.kwargs["where"] for c in find_first.await_args_list] == [
|
||||
_mapping_where("sub", "unknown-svc", ISSUER_TWO),
|
||||
_mapping_where("sub", "unknown-svc", None),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -8025,7 +8062,10 @@ async def test_proxy_admin_sentinel_cached_by_another_issuer_does_not_bypass_rej
|
|||
jwt_handler = _per_issuer_virtual_key_jwt_handler(global_claim_field="sub", global_behavior="auto_register")
|
||||
prisma_client, find_first = _fake_prisma_with_jwt_key_mapping(None)
|
||||
user_api_key_cache = DualCache()
|
||||
await user_api_key_cache.async_set_cache(key="jwt_key_mapping:sub:admin-7", value=_JWT_PROXY_ADMIN_SENTINEL)
|
||||
# Sentinel cached under issuer-one's own key -- must never answer issuer-two's lookup.
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=jwt_key_mapping_cache_key("sub", "admin-7", ISSUER_ONE), value=_JWT_PROXY_ADMIN_SENTINEL
|
||||
)
|
||||
|
||||
auto_register_issuer_result = await _resolve_jwt_to_virtual_key(
|
||||
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_ONE, "sub": "admin-7"},
|
||||
|
|
@ -8050,7 +8090,10 @@ async def test_proxy_admin_sentinel_cached_by_another_issuer_does_not_bypass_rej
|
|||
|
||||
assert exc.value.status_code == 403
|
||||
assert "No registered mapping for sub='admin-7'" in str(exc.value.detail)
|
||||
find_first.assert_awaited_once_with(where=_mapping_where("sub", "admin-7"))
|
||||
assert [c.kwargs["where"] for c in find_first.await_args_list] == [
|
||||
_mapping_where("sub", "admin-7", ISSUER_TWO),
|
||||
_mapping_where("sub", "admin-7", None),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -8079,7 +8122,125 @@ async def test_issuer_without_virtual_key_claim_field_falls_back_to_the_global_f
|
|||
|
||||
assert with_claim is None
|
||||
assert without_claim is None
|
||||
find_first.assert_awaited_once_with(where=_mapping_where("client_id", "app-9"))
|
||||
# without_claim has no claim value and returns before ever reaching the DB.
|
||||
assert [c.kwargs["where"] for c in find_first.await_args_list] == [
|
||||
_mapping_where("client_id", "app-9", ISSUER_ONE),
|
||||
_mapping_where("client_id", "app-9", None),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_colliding_claim_value_from_another_issuer_does_not_resolve_to_the_wrong_virtual_key():
|
||||
"""LIT-7417: a mapping registered for one issuer must not answer a lookup from a
|
||||
DIFFERENT issuer whose claim value happens to collide, even though both issuers
|
||||
use the same claim field (``sub``) for their virtual-key mapping."""
|
||||
from litellm.proxy.auth.user_api_key_auth import _resolve_jwt_to_virtual_key
|
||||
|
||||
jwt_handler = _per_issuer_virtual_key_jwt_handler(global_claim_field="sub")
|
||||
prisma_client, find_first = _fake_prisma_jwt_key_mapping_table(
|
||||
[
|
||||
{
|
||||
"jwt_issuer": ISSUER_TWO,
|
||||
"jwt_claim_name": "sub",
|
||||
"jwt_claim_value": "dev-alice",
|
||||
"token": "hashed-issuer-b-key",
|
||||
"is_active": True,
|
||||
}
|
||||
]
|
||||
)
|
||||
user_api_key_cache = DualCache()
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key="hashed-issuer-b-key",
|
||||
value=UserAPIKeyAuth(token="hashed-issuer-b-key", api_key="hashed-issuer-b-key", team_id="issuer-b-team"),
|
||||
)
|
||||
|
||||
owner_result = await _resolve_jwt_to_virtual_key(
|
||||
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_TWO, "sub": "dev-alice"},
|
||||
jwt_handler=jwt_handler,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=None,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
assert isinstance(owner_result, UserAPIKeyAuth)
|
||||
assert owner_result.token == "hashed-issuer-b-key"
|
||||
|
||||
# issuer-one's behavior is fallback_team_mapping: a correctly-scoped miss must
|
||||
# return None (fall through to team-based JWT auth), never issuer-two's key.
|
||||
colliding_result = await _resolve_jwt_to_virtual_key(
|
||||
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_ONE, "sub": "dev-alice"},
|
||||
jwt_handler=jwt_handler,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=None,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
assert colliding_result is None
|
||||
assert find_first.await_count == 3 # owner hit (1 call) + colliding miss (issuer-scoped + global fallback)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cached_resolution_for_one_issuer_does_not_leak_to_a_colliding_issuer():
|
||||
"""A cached positive resolution must be keyed by issuer too, or a colliding
|
||||
claim value from another issuer could be served straight from cache without
|
||||
ever reaching the (correctly issuer-scoped) DB lookup."""
|
||||
from litellm.proxy.auth.user_api_key_auth import _resolve_jwt_to_virtual_key
|
||||
|
||||
jwt_handler = _per_issuer_virtual_key_jwt_handler(global_claim_field="sub")
|
||||
prisma_client, find_first = _fake_prisma_jwt_key_mapping_table([])
|
||||
user_api_key_cache = DualCache()
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=jwt_key_mapping_cache_key("sub", "dev-alice", ISSUER_TWO), value="hashed-issuer-b-key"
|
||||
)
|
||||
|
||||
colliding_result = await _resolve_jwt_to_virtual_key(
|
||||
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_ONE, "sub": "dev-alice"},
|
||||
jwt_handler=jwt_handler,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=None,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
assert colliding_result is None
|
||||
# Must have gone to the DB rather than serving issuer-two's cached token.
|
||||
assert find_first.await_count == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_issuer_agnostic_mapping_matches_every_issuer():
|
||||
"""A mapping created before issuer scoping existed (``jwt_issuer`` is NULL) keeps
|
||||
matching any issuer, so existing global mappings are not broken by this fix."""
|
||||
from litellm.proxy.auth.user_api_key_auth import _resolve_jwt_to_virtual_key
|
||||
|
||||
jwt_handler = _per_issuer_virtual_key_jwt_handler(global_claim_field="sub")
|
||||
prisma_client, _find_first = _fake_prisma_jwt_key_mapping_table(
|
||||
[
|
||||
{
|
||||
"jwt_issuer": "",
|
||||
"jwt_claim_name": "sub",
|
||||
"jwt_claim_value": "legacy-user",
|
||||
"token": "hashed-legacy-key",
|
||||
"is_active": True,
|
||||
}
|
||||
]
|
||||
)
|
||||
user_api_key_cache = DualCache()
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key="hashed-legacy-key",
|
||||
value=UserAPIKeyAuth(token="hashed-legacy-key", api_key="hashed-legacy-key", team_id="legacy-team"),
|
||||
)
|
||||
|
||||
for issuer in (ISSUER_ONE, ISSUER_TWO):
|
||||
resolved = await _resolve_jwt_to_virtual_key(
|
||||
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: issuer, "sub": "legacy-user"},
|
||||
jwt_handler=jwt_handler,
|
||||
prisma_client=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 == "hashed-legacy-key"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -36,7 +36,11 @@ from litellm.proxy._types import (
|
|||
UpdateKeyRequest,
|
||||
)
|
||||
from litellm.models.object_permission import LiteLLM_ObjectPermissionTable
|
||||
from litellm.proxy.auth.auth_checks import _delete_cache_key_object, _project_cache_key
|
||||
from litellm.proxy.auth.auth_checks import (
|
||||
_delete_cache_key_object,
|
||||
_project_cache_key,
|
||||
jwt_key_mapping_cache_key,
|
||||
)
|
||||
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
|
||||
|
|
@ -5132,10 +5136,11 @@ async def test_delete_verification_tokens_persists_deleted_keys(monkeypatch):
|
|||
|
||||
|
||||
class _JWTMappingRow:
|
||||
def __init__(self, token, jwt_claim_name, jwt_claim_value):
|
||||
def __init__(self, token, jwt_claim_name, jwt_claim_value, jwt_issuer=None):
|
||||
self.token = token
|
||||
self.jwt_claim_name = jwt_claim_name
|
||||
self.jwt_claim_value = jwt_claim_value
|
||||
self.jwt_issuer = jwt_issuer
|
||||
|
||||
|
||||
class _CascadingJWTMappingTable:
|
||||
|
|
@ -5226,7 +5231,7 @@ async def test_delete_verification_tokens_evicts_jwt_key_mapping_cache(monkeypat
|
|||
),
|
||||
)
|
||||
|
||||
assert recording_evict.cache_keys == ("jwt_key_mapping:email:user@example.com",)
|
||||
assert recording_evict.cache_keys == (jwt_key_mapping_cache_key("email", "user@example.com", None),)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -13131,11 +13136,11 @@ async def test_regenerate_evicts_jwt_key_mapping_cache_so_next_jwt_call_gets_new
|
|||
_execute_virtual_key_regeneration,
|
||||
)
|
||||
|
||||
stale_cache_key = "jwt_key_mapping:sub:user1"
|
||||
stale_cache_key = jwt_key_mapping_cache_key("sub", "user1", None)
|
||||
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")]
|
||||
return_value=[MagicMock(jwt_claim_name="sub", jwt_claim_value="user1", jwt_issuer=None)]
|
||||
)
|
||||
mock_prisma_client.db.litellm_jwtkeymapping.find_first = AsyncMock(
|
||||
return_value=MagicMock(token="new-hashed-token")
|
||||
|
|
|
|||
6
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
6
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -27297,6 +27297,8 @@ export interface components {
|
|||
jwt_claim_name: string;
|
||||
/** Jwt Claim Value */
|
||||
jwt_claim_value: string;
|
||||
/** Jwt Issuer */
|
||||
jwt_issuer?: string | null;
|
||||
/** Key */
|
||||
key: string;
|
||||
};
|
||||
|
|
@ -28910,6 +28912,8 @@ export interface components {
|
|||
jwt_claim_name: string;
|
||||
/** Jwt Claim Value */
|
||||
jwt_claim_value: string;
|
||||
/** Jwt Issuer */
|
||||
jwt_issuer?: string | null;
|
||||
/**
|
||||
* Updated At
|
||||
* Format: date-time
|
||||
|
|
@ -38643,6 +38647,8 @@ export interface components {
|
|||
id: string;
|
||||
/** Is Active */
|
||||
is_active?: boolean | null;
|
||||
/** Jwt Issuer */
|
||||
jwt_issuer?: string | null;
|
||||
/** Key */
|
||||
key?: string | null;
|
||||
};
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue