mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-04 07:27:51 +00:00
feat(auth): add link-only JWT mapping and remove route eligibility filtering (#49956)
* feat(auth): allow skipping JWT key route eligibility check * feat(auth): support link-only JWT key mapping * fix(auth): model JWT mapping rejection as a failure value * feat(auth): select existing JWT keys without route filtering
This commit is contained in:
parent
ce73cf8b9d
commit
7fa438df06
4 changed files with 406 additions and 44 deletions
|
|
@ -4796,11 +4796,14 @@ class UnregisteredJWTClientBehavior(str, enum.Enum):
|
|||
- auto_register: Automatically create a new virtual key and mapping on first
|
||||
encounter. The new key has no budget/model restrictions; admins can tighten
|
||||
it later via /jwt_client/update.
|
||||
- link_only: Map an eligible existing user key or return HTTP 403, without
|
||||
creating a key. Proxy admins authenticate without a mapped key.
|
||||
"""
|
||||
|
||||
FALLBACK_TEAM_MAPPING = "fallback_team_mapping"
|
||||
REJECT = "reject"
|
||||
AUTO_REGISTER = "auto_register"
|
||||
LINK_ONLY = "link_only"
|
||||
|
||||
|
||||
class JWTIssuerConfig(LiteLLMBaseModel):
|
||||
|
|
@ -5038,7 +5041,12 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase):
|
|||
"What to do when virtual_key_claim_field is set but the JWT claim value "
|
||||
"has no registered mapping. 'fallback_team_mapping' (default): fall through "
|
||||
"to team-based JWT auth. 'reject': return HTTP 403. "
|
||||
"'auto_register': auto-create a virtual key and mapping on first encounter."
|
||||
"'auto_register': auto-create a virtual key and mapping on first encounter. "
|
||||
"'link_only': map an eligible existing user key or return HTTP 403, without creating a key. "
|
||||
"link_only ignores auto_register_map_existing_key; user provisioning still follows user_id_upsert. "
|
||||
"Existing-key selection ignores allowed_routes; JWTs inherit the selected key's permissions, "
|
||||
"including management routes, with request-time key restrictions still enforced. "
|
||||
"Unmapped proxy admins authenticate directly in auto_register and link_only modes."
|
||||
),
|
||||
)
|
||||
auto_register_map_existing_key: bool = Field(
|
||||
|
|
@ -5048,8 +5056,10 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase):
|
|||
"field is the user_id_jwt_field or user_email_jwt_field, the JWT claim is mapped to a virtual key the "
|
||||
"JWT-resolved user already owns instead of minting a new one. If the user owns several, the most recently created key in the "
|
||||
"JWT-resolved team (or with no team when the JWT resolves none) is chosen among keys that never "
|
||||
"expire, are not blocked, are not Admin UI session keys, were not minted by auto_register, and "
|
||||
"have no allowed_routes or include llm_api_routes. Otherwise a new key is minted as usual."
|
||||
"expire, are not blocked, are not Admin UI session keys, and were not minted by auto_register. "
|
||||
"allowed_routes does not affect key selection. JWTs inherit the selected key's permissions, "
|
||||
"including management routes; request-time key restrictions still apply. "
|
||||
"Otherwise a new key is minted as usual."
|
||||
),
|
||||
)
|
||||
routing_overrides: list[JWTRoutingOverride] | None = Field(
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ import re
|
|||
import secrets
|
||||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime, timezone
|
||||
from enum import Enum
|
||||
from types import MappingProxyType
|
||||
from typing import Final, NamedTuple, Protocol, Union, cast
|
||||
|
||||
|
|
@ -1576,10 +1577,14 @@ _JWT_AUTH_DISABLED_HINT = (
|
|||
)
|
||||
|
||||
|
||||
class _JWTKeyMappingFailure(Enum):
|
||||
NO_ELIGIBLE_KEY = "JWT Key Mapping: No eligible existing virtual key for this identity. Access denied."
|
||||
|
||||
|
||||
class _PendingAutoRegister(NamedTuple):
|
||||
"""
|
||||
Signal returned by ``_resolve_jwt_to_virtual_key`` when the JWT's claim is
|
||||
unmapped and ``unregistered_jwt_client_behavior`` is AUTO_REGISTER.
|
||||
unmapped and ``unregistered_jwt_client_behavior`` is AUTO_REGISTER or LINK_ONLY.
|
||||
|
||||
The caller MUST run standard ``JWTAuthManager.auth_builder`` to apply RBAC,
|
||||
scope mappings, ``custom_validate``, and ``user_allowed_email_domain``
|
||||
|
|
@ -1593,23 +1598,22 @@ class _PendingAutoRegister(NamedTuple):
|
|||
claim_value: str
|
||||
cache_key: str
|
||||
jwt_issuer: str | None = None
|
||||
behavior: UnregisteredJWTClientBehavior = UnregisteredJWTClientBehavior.AUTO_REGISTER
|
||||
|
||||
|
||||
def _claim_identifies_user(jwt_handler: JWTHandler, claim_field: str, jwt_issuer: str | None) -> bool:
|
||||
if not jwt_handler.litellm_jwtauth.auto_register_map_existing_key:
|
||||
return False
|
||||
if jwt_handler.litellm_jwtauth.is_user_identity_claim(claim_field, jwt_issuer):
|
||||
return True
|
||||
verbose_proxy_logger.warning(
|
||||
"JWT Key Mapping (auto_register_map_existing_key): claim '%s' is not the user_id or user_email JWT field "
|
||||
"and may be shared by several users, so a new key is minted instead of reusing one the user owns.",
|
||||
"JWT Key Mapping: claim '%s' is not the user_id or user_email JWT field "
|
||||
"and may be shared by several users, so existing user keys cannot be linked automatically.",
|
||||
claim_field,
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
async def _reusable_key_hash_for_user(prisma_client: PrismaClient, user_id: str, team_id: str | None) -> str | None:
|
||||
key: Final = await VerificationTokenRepository(prisma_client).find_newest_reusable_llm_api_key(user_id, team_id)
|
||||
key: Final = await VerificationTokenRepository(prisma_client).find_newest_reusable_key(user_id, team_id)
|
||||
return None if key is None else key.token
|
||||
|
||||
|
||||
|
|
@ -1628,11 +1632,13 @@ async def _auto_register_jwt_mapping(
|
|||
org_id: str | None = None,
|
||||
end_user_id: str | None = None,
|
||||
agent_id: str | None = None,
|
||||
) -> UserAPIKeyAuth | None:
|
||||
behavior: UnregisteredJWTClientBehavior = UnregisteredJWTClientBehavior.AUTO_REGISTER,
|
||||
) -> UserAPIKeyAuth | _JWTKeyMappingFailure | None:
|
||||
"""
|
||||
Auto-register: create a new virtual key + mapping for an unrecognised JWT
|
||||
claim value, or point the mapping at a key the resolved user already owns
|
||||
when ``auto_register_map_existing_key`` is set. ``team_id`` and ``user_id``
|
||||
when ``auto_register_map_existing_key`` is set. LINK_ONLY always searches
|
||||
for an existing key and rejects when none is eligible. ``team_id`` and ``user_id``
|
||||
must come from a successful ``JWTAuthManager.auth_builder`` run — they
|
||||
encode the JWT identity AFTER
|
||||
RBAC/scope/custom_validate/email-domain policy has been enforced. The key
|
||||
|
|
@ -1650,11 +1656,21 @@ async def _auto_register_jwt_mapping(
|
|||
generate_key_helper_fn,
|
||||
)
|
||||
|
||||
link_only: Final = behavior == UnregisteredJWTClientBehavior.LINK_ONLY
|
||||
reuse_existing: Final = link_only or jwt_handler.litellm_jwtauth.auto_register_map_existing_key
|
||||
existing_token_hash: Final = (
|
||||
await _reusable_key_hash_for_user(prisma_client, user_id, team_id)
|
||||
if user_id is not None and _claim_identifies_user(jwt_handler, virtual_key_claim_field, jwt_issuer)
|
||||
await _reusable_key_hash_for_user(
|
||||
prisma_client,
|
||||
user_id,
|
||||
team_id,
|
||||
)
|
||||
if reuse_existing
|
||||
and user_id is not None
|
||||
and _claim_identifies_user(jwt_handler, virtual_key_claim_field, jwt_issuer)
|
||||
else None
|
||||
)
|
||||
if link_only and existing_token_hash is None:
|
||||
return _JWTKeyMappingFailure.NO_ELIGIBLE_KEY
|
||||
minted: Final = existing_token_hash is None
|
||||
if existing_token_hash is not None:
|
||||
token_hash = existing_token_hash
|
||||
|
|
@ -1691,8 +1707,8 @@ async def _auto_register_jwt_mapping(
|
|||
"jwt_claim_name": virtual_key_claim_field,
|
||||
"jwt_claim_value": claim_value,
|
||||
"token": token_hash,
|
||||
"created_by": "auto_register",
|
||||
"updated_by": "auto_register",
|
||||
"created_by": behavior.value,
|
||||
"updated_by": behavior.value,
|
||||
}
|
||||
)
|
||||
except Exception as e:
|
||||
|
|
@ -1704,8 +1720,8 @@ async def _auto_register_jwt_mapping(
|
|||
# the DB and the cleartext is in memory on this request. Delete it
|
||||
# so orphans don't accumulate under sustained concurrency.
|
||||
verbose_proxy_logger.debug(
|
||||
"JWT Key Mapping (auto_register): unique conflict on create — "
|
||||
"deleting orphaned virtual key and fetching winner's mapping for %s='%s'.",
|
||||
"JWT Key Mapping (%s): unique conflict on create, fetching winner's mapping for %s='%s'.",
|
||||
behavior.value,
|
||||
virtual_key_claim_field,
|
||||
claim_value,
|
||||
)
|
||||
|
|
@ -1735,7 +1751,7 @@ async def _auto_register_jwt_mapping(
|
|||
raise HTTPException(
|
||||
status_code=503,
|
||||
detail=(
|
||||
"JWT Key Mapping: AUTO_REGISTER race resolution failed — "
|
||||
f"JWT Key Mapping: {behavior.value.upper()} race resolution failed — "
|
||||
"winner's mapping was concurrently removed. Retry the request."
|
||||
),
|
||||
)
|
||||
|
|
@ -1749,7 +1765,8 @@ async def _auto_register_jwt_mapping(
|
|||
)
|
||||
|
||||
verbose_proxy_logger.info(
|
||||
"JWT Key Mapping (auto_register): %s virtual key for %s='%s'.",
|
||||
"JWT Key Mapping (%s): %s virtual key for %s='%s'.",
|
||||
behavior.value,
|
||||
"created new" if minted else "mapped existing",
|
||||
virtual_key_claim_field,
|
||||
claim_value,
|
||||
|
|
@ -1821,7 +1838,7 @@ async def _resolve_jwt_to_virtual_key(
|
|||
- ``UserAPIKeyAuth``: a resolved virtual key (cache hit or DB hit). The
|
||||
caller may use this directly; JWT policy has been enforced previously
|
||||
(at key-creation time or, for cached results, before caching).
|
||||
- ``_PendingAutoRegister``: claim is unmapped and behavior is AUTO_REGISTER.
|
||||
- ``_PendingAutoRegister``: claim is unmapped and behavior is AUTO_REGISTER or LINK_ONLY.
|
||||
The caller MUST run ``JWTAuthManager.auth_builder`` to enforce JWT
|
||||
policy (RBAC, scope, custom_validate, email-domain), then invoke
|
||||
``_auto_register_jwt_mapping`` with the validated team_id/user_id.
|
||||
|
|
@ -1829,7 +1846,7 @@ async def _resolve_jwt_to_virtual_key(
|
|||
The caller falls through to standard team-based JWT auth (which itself
|
||||
enforces full JWT policy via auth_builder).
|
||||
- Raises HTTPException: REJECT policy hit, missing claim under
|
||||
REJECT/AUTO_REGISTER, or other policy violations.
|
||||
REJECT/AUTO_REGISTER/LINK_ONLY, or other policy violations.
|
||||
"""
|
||||
raw_issuer: Final = jwt_claims.get(JWTHandler.LITELLM_JWT_ISSUER_CLAIM)
|
||||
normalized_issuer: Final = raw_issuer if isinstance(raw_issuer, str) else None
|
||||
|
|
@ -1856,6 +1873,7 @@ async def _resolve_jwt_to_virtual_key(
|
|||
if behavior in (
|
||||
UnregisteredJWTClientBehavior.REJECT,
|
||||
UnregisteredJWTClientBehavior.AUTO_REGISTER,
|
||||
UnregisteredJWTClientBehavior.LINK_ONLY,
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
|
|
@ -1887,7 +1905,7 @@ 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:
|
||||
if behavior in (UnregisteredJWTClientBehavior.AUTO_REGISTER, UnregisteredJWTClientBehavior.LINK_ONLY):
|
||||
# Stale sentinel written under a prior fallback_team_mapping config —
|
||||
# evict it and defer auto-register to after auth_builder runs. Raise
|
||||
# the same 500 as the fresh-path AUTO_REGISTER branch when there is
|
||||
|
|
@ -1897,7 +1915,7 @@ async def _resolve_jwt_to_virtual_key(
|
|||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail=(
|
||||
"JWT Key Mapping: AUTO_REGISTER requires a database connection. "
|
||||
f"JWT Key Mapping: {behavior.value.upper()} requires a database connection. "
|
||||
"Configure a database or change unregistered_jwt_client_behavior."
|
||||
),
|
||||
)
|
||||
|
|
@ -1907,6 +1925,7 @@ async def _resolve_jwt_to_virtual_key(
|
|||
claim_value=str(claim_value),
|
||||
cache_key=cache_key,
|
||||
jwt_issuer=normalized_issuer,
|
||||
behavior=behavior,
|
||||
)
|
||||
return None
|
||||
elif cached_mapping is not None:
|
||||
|
|
@ -1967,12 +1986,12 @@ async def _resolve_jwt_to_virtual_key(
|
|||
detail=f"JWT Key Mapping: No registered mapping for {virtual_key_claim_field}='{claim_value}'. Access denied.",
|
||||
)
|
||||
|
||||
if behavior == UnregisteredJWTClientBehavior.AUTO_REGISTER:
|
||||
if behavior in (UnregisteredJWTClientBehavior.AUTO_REGISTER, UnregisteredJWTClientBehavior.LINK_ONLY):
|
||||
if prisma_client is None:
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail=(
|
||||
"JWT Key Mapping: AUTO_REGISTER requires a database connection. "
|
||||
f"JWT Key Mapping: {behavior.value.upper()} requires a database connection. "
|
||||
"Configure a database or change unregistered_jwt_client_behavior."
|
||||
),
|
||||
)
|
||||
|
|
@ -1985,6 +2004,7 @@ async def _resolve_jwt_to_virtual_key(
|
|||
claim_value=str(claim_value),
|
||||
cache_key=cache_key,
|
||||
jwt_issuer=normalized_issuer,
|
||||
behavior=behavior,
|
||||
)
|
||||
|
||||
# FALLBACK_TEAM_MAPPING (default): cache the miss and return None so the
|
||||
|
|
@ -2477,10 +2497,13 @@ async def user_api_key_auth_builder(
|
|||
org_id=org_id,
|
||||
end_user_id=end_user_id,
|
||||
agent_id=agent_id,
|
||||
behavior=pending_auto_register.behavior,
|
||||
)
|
||||
if pending_auto_register is not None and prisma_client is not None
|
||||
else None
|
||||
)
|
||||
if isinstance(auto_registered, _JWTKeyMappingFailure):
|
||||
raise HTTPException(status_code=403, detail=auto_registered.value)
|
||||
if auto_registered is not None:
|
||||
auto_registered.jwt_claims = jwt_claims
|
||||
auto_registered.user_email = user_email
|
||||
|
|
@ -2496,7 +2519,13 @@ async def user_api_key_auth_builder(
|
|||
|
||||
falls_through_to_key_checks: Final = (
|
||||
auto_registered is not None
|
||||
and jwt_handler.litellm_jwtauth.auto_register_map_existing_key
|
||||
and (
|
||||
jwt_handler.litellm_jwtauth.auto_register_map_existing_key
|
||||
or (
|
||||
pending_auto_register is not None
|
||||
and pending_auto_register.behavior == UnregisteredJWTClientBehavior.LINK_ONLY
|
||||
)
|
||||
)
|
||||
and master_key is not None
|
||||
)
|
||||
if not falls_through_to_key_checks:
|
||||
|
|
|
|||
|
|
@ -124,9 +124,7 @@ class VerificationTokenRepository(BaseRepository[LiteLLM_VerificationToken]):
|
|||
records: Final[Sequence[PrismaVerificationToken]] = await self.table.find_many(where={"user_id": user_id})
|
||||
return self._to_model_list(records)
|
||||
|
||||
async def find_newest_reusable_llm_api_key(
|
||||
self, user_id: str, team_id: str | None
|
||||
) -> LiteLLM_VerificationToken | None:
|
||||
async def find_newest_reusable_key(self, user_id: str, team_id: str | None) -> LiteLLM_VerificationToken | None:
|
||||
records: Final[Sequence[PrismaVerificationToken]] = await self.table.find_many(
|
||||
where={
|
||||
"user_id": user_id,
|
||||
|
|
@ -140,12 +138,6 @@ class VerificationTokenRepository(BaseRepository[LiteLLM_VerificationToken]):
|
|||
{"team_id": {"not": UI_SESSION_TOKEN_TEAM_ID}},
|
||||
]
|
||||
},
|
||||
{
|
||||
"OR": [
|
||||
{"allowed_routes": {"is_empty": True}},
|
||||
{"allowed_routes": {"has": "llm_api_routes"}},
|
||||
]
|
||||
},
|
||||
],
|
||||
},
|
||||
order={"created_at": "desc"},
|
||||
|
|
|
|||
|
|
@ -7,13 +7,15 @@ from concurrent.futures import ThreadPoolExecutor
|
|||
from contextlib import contextmanager
|
||||
from hashlib import sha256
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
from typing import Final, Literal
|
||||
|
||||
import httpx
|
||||
import jwt
|
||||
import pytest
|
||||
import yaml
|
||||
from cryptography.hazmat.primitives.asymmetric import rsa
|
||||
from pydantic import JsonValue
|
||||
from redis import Redis
|
||||
|
||||
from tests.integration._support.client import Gateway, eventually, string_value
|
||||
from tests.integration._support.database import read_rows
|
||||
|
|
@ -31,7 +33,15 @@ def _hash(key: str) -> str:
|
|||
return sha256(key.encode()).hexdigest()
|
||||
|
||||
|
||||
def _config(directory: Path, claim_field: str) -> Path:
|
||||
def _config(
|
||||
directory: Path,
|
||||
claim_field: str,
|
||||
*,
|
||||
behavior: Literal["auto_register", "link_only", "reject"] = "auto_register",
|
||||
map_existing_key: bool = True,
|
||||
user_id_upsert: bool = False,
|
||||
issuers: tuple[dict[str, JsonValue], ...] = (),
|
||||
) -> Path:
|
||||
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
config["general_settings"] = {
|
||||
**config["general_settings"],
|
||||
|
|
@ -40,8 +50,10 @@ def _config(directory: Path, claim_field: str) -> Path:
|
|||
"user_id_jwt_field": "sub",
|
||||
"user_email_jwt_field": "email",
|
||||
"virtual_key_claim_field": claim_field,
|
||||
"unregistered_jwt_client_behavior": "auto_register",
|
||||
"auto_register_map_existing_key": True,
|
||||
"unregistered_jwt_client_behavior": behavior,
|
||||
"auto_register_map_existing_key": map_existing_key,
|
||||
"user_id_upsert": user_id_upsert,
|
||||
"issuers": list(issuers),
|
||||
},
|
||||
}
|
||||
path: Final = directory / f"jwt_map_existing_key_{claim_field}.yaml"
|
||||
|
|
@ -107,7 +119,7 @@ def _billed_key(response: httpx.Response) -> str:
|
|||
return string_value(rows[0]["api_key"])
|
||||
|
||||
|
||||
def test_first_jwt_call_reuses_the_newest_durable_llm_key_and_skips_every_ineligible_newer_key(
|
||||
def test_first_jwt_call_reuses_the_newest_durable_key_and_skips_every_ineligible_newer_key(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
with _issuer() as (private_key, jwks_url), gateway.scenario() as scenario:
|
||||
|
|
@ -118,7 +130,6 @@ def test_first_jwt_call_reuses_the_newest_durable_llm_key_and_skips_every_inelig
|
|||
skipped: Final = {
|
||||
"older_durable": older_durable,
|
||||
"expiring": scenario.key(user_id=user, duration="1h"),
|
||||
"management_only": scenario.key(user_id=user, allowed_routes=["management_routes"]),
|
||||
"auto_registered_look_alike": scenario.key(user_id=user, metadata={"auto_registered": True}),
|
||||
"other_team": scenario.key(user_id=user, team_id=scenario.team()),
|
||||
"blocked": scenario.key(user_id=user),
|
||||
|
|
@ -140,6 +151,84 @@ def test_first_jwt_call_reuses_the_newest_durable_llm_key_and_skips_every_inelig
|
|||
assert _billed_key(response) == _hash(durable)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("behavior", "map_existing_key"), [("auto_register", True), ("link_only", False), ("link_only", True)]
|
||||
)
|
||||
def test_route_restricted_keys_are_linked_without_bypassing_other_key_checks(
|
||||
gateway: Gateway,
|
||||
tmp_path: Path,
|
||||
behavior: Literal["auto_register", "link_only"],
|
||||
map_existing_key: bool,
|
||||
) -> None:
|
||||
with _issuer() as (private_key, jwks_url), gateway.scenario() as scenario:
|
||||
model: Final = scenario.model()
|
||||
user: Final = scenario.user(user_role="internal_user")
|
||||
scenario.key(user_id=user)
|
||||
restricted: Final = scenario.key(user_id=user, allowed_routes=["/v1/chat/completions"])
|
||||
scenario.key(user_id=user, duration="1h")
|
||||
scenario.key(user_id=user, metadata={"auto_registered": True})
|
||||
scenario.key(user_id=user, team_id=scenario.team())
|
||||
blocked: Final = scenario.key(user_id=user)
|
||||
gateway.post("/key/block", {"key": blocked})
|
||||
keys_before: Final = _user_key_hashes(user)
|
||||
token: Final = _token(private_key, user)
|
||||
|
||||
with owned_proxy(
|
||||
gateway,
|
||||
tmp_path,
|
||||
{"JWT_PUBLIC_KEY_URL": jwks_url},
|
||||
config=_config(
|
||||
tmp_path,
|
||||
"sub",
|
||||
behavior=behavior,
|
||||
map_existing_key=map_existing_key,
|
||||
),
|
||||
) as candidate:
|
||||
response: Final = _chat(candidate, model, token)
|
||||
repeated: Final = _chat(candidate, model, token)
|
||||
restricted_route: Final = candidate.request("GET", "/v1/models", key=token)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert repeated.status_code == 200, repeated.text
|
||||
assert restricted_route.status_code == 403, restricted_route.text
|
||||
assert _mapped_token("sub", user) == _hash(restricted)
|
||||
assert _user_key_hashes(user) == keys_before, "a new key was minted instead of linking an existing key"
|
||||
assert _billed_key(response) == _hash(restricted)
|
||||
assert _billed_key(repeated) == _hash(restricted)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("behavior", ["auto_register", "link_only"])
|
||||
def test_management_route_keys_are_linked_by_default_and_keep_their_route_permissions(
|
||||
gateway: Gateway, tmp_path: Path, behavior: Literal["auto_register", "link_only"]
|
||||
) -> None:
|
||||
with _issuer() as (private_key, jwks_url), gateway.scenario() as scenario:
|
||||
model: Final = scenario.model()
|
||||
user: Final = scenario.user(user_role="internal_user")
|
||||
target: Final = f"integration-{uuid.uuid4().hex}"
|
||||
scenario.key(user_id=user)
|
||||
management_key: Final = scenario.key(user_id=user, allowed_routes=["/user/new"])
|
||||
keys_before: Final = _user_key_hashes(user)
|
||||
token: Final = _token(private_key, user)
|
||||
with owned_proxy(
|
||||
gateway, tmp_path, {"JWT_PUBLIC_KEY_URL": jwks_url}, config=_config(tmp_path, "sub", behavior=behavior)
|
||||
) as candidate:
|
||||
denied: Final = _chat(candidate, model, token)
|
||||
allowed: Final = candidate.request(
|
||||
"POST",
|
||||
"/user/new",
|
||||
{"user_id": target, "auto_create_key": False, "user_role": "internal_user"},
|
||||
key=token,
|
||||
)
|
||||
assert allowed.status_code == 200, allowed.text
|
||||
scenario.cleanups.callback(scenario.delete_user, target)
|
||||
repeated: Final = _chat(candidate, model, token)
|
||||
assert denied.status_code == 403, denied.text
|
||||
assert repeated.status_code == 403, repeated.text
|
||||
assert allowed.json()["user_id"] == target
|
||||
assert _mapped_token("sub", user) == _hash(management_key)
|
||||
assert _user_key_hashes(user) == keys_before
|
||||
|
||||
|
||||
def test_user_matched_by_email_instead_of_sub_still_reuses_their_existing_key(gateway: Gateway, tmp_path: Path) -> None:
|
||||
with _issuer() as (private_key, jwks_url), gateway.scenario() as scenario:
|
||||
model: Final = scenario.model()
|
||||
|
|
@ -183,8 +272,9 @@ def test_shared_client_claim_never_maps_a_second_user_onto_the_first_users_perso
|
|||
assert read_rows('SELECT request_id FROM "LiteLLM_SpendLogs" WHERE api_key = %s', (_hash(personal),)) == []
|
||||
|
||||
|
||||
def test_concurrent_first_jwt_calls_of_a_keyless_user_both_succeed_on_one_surviving_mapped_key(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
@pytest.mark.parametrize("behavior", ["auto_register", "link_only"])
|
||||
def test_concurrent_first_jwt_calls_both_succeed_on_one_surviving_mapped_key(
|
||||
gateway: Gateway, tmp_path: Path, behavior: Literal["auto_register", "link_only"]
|
||||
) -> None:
|
||||
writer_url: Final = os.environ.get("INTEGRATION_PROXY_DATABASE_URL") or os.environ["DATABASE_URL"]
|
||||
with (
|
||||
|
|
@ -194,6 +284,7 @@ def test_concurrent_first_jwt_calls_of_a_keyless_user_both_succeed_on_one_surviv
|
|||
):
|
||||
model: Final = scenario.model()
|
||||
user: Final = scenario.user(user_role="internal_user")
|
||||
existing: Final = scenario.key(user_id=user) if behavior == "link_only" else None
|
||||
token: Final = _token(private_key, user)
|
||||
overrides: Final = {
|
||||
"JWT_PUBLIC_KEY_URL": jwks_url,
|
||||
|
|
@ -202,7 +293,12 @@ def test_concurrent_first_jwt_calls_of_a_keyless_user_both_succeed_on_one_surviv
|
|||
}
|
||||
|
||||
with (
|
||||
owned_proxy(gateway, tmp_path, overrides, config=_config(tmp_path, "sub")) as candidate,
|
||||
owned_proxy(
|
||||
gateway,
|
||||
tmp_path,
|
||||
overrides,
|
||||
config=_config(tmp_path, "sub", behavior=behavior, map_existing_key=False),
|
||||
) as candidate,
|
||||
ThreadPoolExecutor(max_workers=1) as pool,
|
||||
):
|
||||
held_call: Final = pool.submit(_chat, candidate, model, token)
|
||||
|
|
@ -215,5 +311,240 @@ def test_concurrent_first_jwt_calls_of_a_keyless_user_both_succeed_on_one_surviv
|
|||
assert held.status_code == 200, held.text
|
||||
keys: Final = _user_key_hashes(user)
|
||||
assert len(keys) == 1, keys
|
||||
if existing is not None:
|
||||
assert keys == frozenset({_hash(existing)})
|
||||
assert _mapped_token("sub", user) in keys
|
||||
assert (_billed_key(held), _billed_key(racing)) == (_mapped_token("sub", user),) * 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize("ineligible_keys", [False, True])
|
||||
def test_link_only_rejects_without_writes_then_links_a_newly_available_key(
|
||||
gateway: Gateway, tmp_path: Path, ineligible_keys: bool
|
||||
) -> None:
|
||||
with _issuer() as (private_key, jwks_url), gateway.scenario() as scenario:
|
||||
model: Final = scenario.model()
|
||||
user: Final = scenario.user(user_role="internal_user")
|
||||
if ineligible_keys:
|
||||
scenario.key(user_id=user, duration="1h")
|
||||
scenario.key(user_id=user, metadata={"auto_registered": True})
|
||||
scenario.key(user_id=user, team_id=scenario.team())
|
||||
blocked: Final = scenario.key(user_id=user)
|
||||
gateway.post("/key/block", {"key": blocked})
|
||||
keys_before: Final = _user_key_hashes(user)
|
||||
token: Final = _token(private_key, user)
|
||||
with owned_proxy(
|
||||
gateway,
|
||||
tmp_path,
|
||||
{"JWT_PUBLIC_KEY_URL": jwks_url},
|
||||
config=_config(tmp_path, "sub", behavior="link_only", map_existing_key=False),
|
||||
) as candidate:
|
||||
first: Final = _chat(candidate, model, token)
|
||||
repeated: Final = _chat(candidate, model, token)
|
||||
assert first.status_code == 403, first.text
|
||||
assert repeated.status_code == 403, repeated.text
|
||||
assert _user_key_hashes(user) == keys_before
|
||||
assert read_rows('SELECT token FROM "LiteLLM_JWTKeyMapping" WHERE jwt_claim_value=%s', (user,)) == []
|
||||
available: Final = scenario.key(user_id=user)
|
||||
accepted: Final = _chat(candidate, model, token)
|
||||
assert accepted.status_code == 200, accepted.text
|
||||
assert _mapped_token("sub", user) == _hash(available)
|
||||
assert _user_key_hashes(user) == keys_before | {_hash(available)}
|
||||
assert _billed_key(accepted) == _hash(available)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("claim_field", ["client_id", "email"])
|
||||
def test_link_only_rejects_shared_or_missing_claims_without_linking_a_personal_key(
|
||||
gateway: Gateway, tmp_path: Path, claim_field: str
|
||||
) -> None:
|
||||
with _issuer() as (private_key, jwks_url), gateway.scenario() as scenario:
|
||||
model: Final = scenario.model()
|
||||
user: Final = scenario.user(user_role="internal_user")
|
||||
key: Final = scenario.key(user_id=user)
|
||||
shared_claim: Final = f"integration-client-{uuid.uuid4().hex}"
|
||||
with owned_proxy(
|
||||
gateway,
|
||||
tmp_path,
|
||||
{"JWT_PUBLIC_KEY_URL": jwks_url},
|
||||
config=_config(tmp_path, claim_field, behavior="link_only", map_existing_key=False),
|
||||
) as candidate:
|
||||
response: Final = _chat(candidate, model, _token(private_key, user, client_id=shared_claim))
|
||||
assert response.status_code == 403, response.text
|
||||
assert _user_key_hashes(user) == frozenset({_hash(key)})
|
||||
assert read_rows('SELECT token FROM "LiteLLM_JWTKeyMapping" WHERE token=%s', (_hash(key),)) == []
|
||||
assert read_rows('SELECT token FROM "LiteLLM_JWTKeyMapping" WHERE jwt_claim_value=%s', (shared_claim,)) == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize("restriction", ["models", "allowed_routes"])
|
||||
def test_link_only_enforces_key_permissions_on_the_first_and_cached_requests(
|
||||
gateway: Gateway, tmp_path: Path, restriction: str
|
||||
) -> None:
|
||||
with _issuer() as (private_key, jwks_url), gateway.scenario() as scenario:
|
||||
allowed: Final = scenario.model()
|
||||
denied: Final = scenario.model()
|
||||
user: Final = scenario.user(user_role="internal_user")
|
||||
key: Final = scenario.key(
|
||||
user_id=user,
|
||||
**({"models": [allowed]} if restriction == "models" else {"allowed_routes": ["/v1/models"]}),
|
||||
)
|
||||
token: Final = _token(private_key, user)
|
||||
with owned_proxy(
|
||||
gateway,
|
||||
tmp_path,
|
||||
{"JWT_PUBLIC_KEY_URL": jwks_url},
|
||||
config=_config(tmp_path, "sub", behavior="link_only", map_existing_key=False),
|
||||
) as candidate:
|
||||
first: Final = _chat(candidate, denied, token)
|
||||
repeated: Final = _chat(candidate, denied, token)
|
||||
control: Final = (
|
||||
_chat(candidate, allowed, token)
|
||||
if restriction == "models"
|
||||
else candidate.request("GET", "/v1/models", key=token)
|
||||
)
|
||||
assert first.status_code == 403, first.text
|
||||
assert repeated.status_code == 403, repeated.text
|
||||
assert control.status_code == 200, control.text
|
||||
assert _mapped_token("sub", user) == _hash(key)
|
||||
assert _user_key_hashes(user) == frozenset({_hash(key)})
|
||||
|
||||
|
||||
def test_link_only_does_not_replace_an_existing_mapping_to_a_blocked_key(gateway: Gateway, tmp_path: Path) -> None:
|
||||
with _issuer() as (private_key, jwks_url), gateway.scenario() as scenario:
|
||||
model: Final = scenario.model()
|
||||
user: Final = scenario.user(user_role="internal_user")
|
||||
blocked: Final = scenario.key(user_id=user)
|
||||
available: Final = scenario.key(user_id=user)
|
||||
gateway.post("/jwt/key/mapping/new", {"jwt_claim_name": "sub", "jwt_claim_value": user, "key": blocked})
|
||||
gateway.post("/key/block", {"key": blocked})
|
||||
token: Final = _token(private_key, user)
|
||||
with owned_proxy(
|
||||
gateway,
|
||||
tmp_path,
|
||||
{"JWT_PUBLIC_KEY_URL": jwks_url},
|
||||
config=_config(tmp_path, "sub", behavior="link_only", map_existing_key=False),
|
||||
) as candidate:
|
||||
first: Final = _chat(candidate, model, token)
|
||||
repeated: Final = _chat(candidate, model, token)
|
||||
assert first.status_code == 401, first.text
|
||||
assert repeated.status_code == 401, repeated.text
|
||||
assert _mapped_token("sub", user) == _hash(blocked)
|
||||
assert _user_key_hashes(user) == frozenset({_hash(blocked), _hash(available)})
|
||||
|
||||
|
||||
def test_link_only_exempts_proxy_admins_but_rechecks_the_role_after_a_cached_exemption(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
with _issuer() as (private_key, jwks_url), gateway.scenario() as scenario:
|
||||
model: Final = scenario.model()
|
||||
user: Final = scenario.user(user_role="internal_user")
|
||||
admin: Final = _token(private_key, user, scope="litellm_proxy_admin")
|
||||
ordinary: Final = _token(private_key, user)
|
||||
with owned_proxy(
|
||||
gateway,
|
||||
tmp_path,
|
||||
{"JWT_PUBLIC_KEY_URL": jwks_url},
|
||||
config=_config(tmp_path, "sub", behavior="link_only", map_existing_key=False),
|
||||
) as candidate:
|
||||
first: Final = candidate.request("GET", "/v1/models", key=admin)
|
||||
repeated: Final = candidate.request("GET", "/v1/models", key=admin)
|
||||
denied: Final = candidate.request("GET", "/v1/models", key=ordinary)
|
||||
forbidden_route: Final = _chat(candidate, model, admin)
|
||||
assert first.status_code == 200, first.text
|
||||
assert repeated.status_code == 200, repeated.text
|
||||
assert denied.status_code == 403, denied.text
|
||||
assert forbidden_route.status_code == 401, forbidden_route.text
|
||||
assert _user_key_hashes(user) == frozenset()
|
||||
assert read_rows('SELECT token FROM "LiteLLM_JWTKeyMapping" WHERE jwt_claim_value=%s', (user,)) == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize("user_id_upsert", [False, True])
|
||||
def test_link_only_keeps_user_provisioning_independent_of_key_creation(
|
||||
gateway: Gateway, tmp_path: Path, user_id_upsert: bool
|
||||
) -> None:
|
||||
with _issuer() as (private_key, jwks_url), gateway.scenario() as scenario:
|
||||
model: Final = scenario.model()
|
||||
subject: Final = f"integration-link-only-{uuid.uuid4().hex}"
|
||||
try:
|
||||
with owned_proxy(
|
||||
gateway,
|
||||
tmp_path,
|
||||
{"JWT_PUBLIC_KEY_URL": jwks_url},
|
||||
config=_config(
|
||||
tmp_path, "sub", behavior="link_only", map_existing_key=False, user_id_upsert=user_id_upsert
|
||||
),
|
||||
) as candidate:
|
||||
response: Final = _chat(candidate, model, _token(private_key, subject))
|
||||
assert response.status_code == (403 if user_id_upsert else 401), response.text
|
||||
users: Final = read_rows('SELECT user_id FROM "LiteLLM_UserTable" WHERE user_id=%s', (subject,))
|
||||
assert users == ([{"user_id": subject}] if user_id_upsert else [])
|
||||
assert _user_key_hashes(subject) == frozenset()
|
||||
assert read_rows('SELECT token FROM "LiteLLM_JWTKeyMapping" WHERE jwt_claim_value=%s', (subject,)) == []
|
||||
finally:
|
||||
if read_rows('SELECT user_id FROM "LiteLLM_UserTable" WHERE user_id=%s', (subject,)):
|
||||
scenario.delete_user(subject)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("issuer_behavior", ["auto_register", "link_only"])
|
||||
def test_issuer_specific_registration_mode_overrides_the_global_mode(
|
||||
gateway: Gateway, tmp_path: Path, issuer_behavior: Literal["auto_register", "link_only"]
|
||||
) -> None:
|
||||
with _issuer() as (private_key, jwks_url), gateway.scenario() as scenario:
|
||||
model: Final = scenario.model()
|
||||
user: Final = scenario.user(user_role="internal_user")
|
||||
issuer: Final = f"https://integration.example.com/{uuid.uuid4().hex}"
|
||||
token: Final = _token(private_key, user, iss=issuer, aud="integration")
|
||||
path: Final = _config(
|
||||
tmp_path,
|
||||
"sub",
|
||||
behavior="auto_register" if issuer_behavior == "link_only" else "link_only",
|
||||
map_existing_key=False,
|
||||
issuers=(
|
||||
{
|
||||
"issuer": issuer,
|
||||
"jwks_url": jwks_url,
|
||||
"audience": "integration",
|
||||
"user_id_jwt_field": "sub",
|
||||
"virtual_key_claim_field": "sub",
|
||||
"unregistered_jwt_client_behavior": issuer_behavior,
|
||||
},
|
||||
),
|
||||
)
|
||||
with owned_proxy(gateway, tmp_path, {"JWT_PUBLIC_KEY_URL": jwks_url}, config=path) as candidate:
|
||||
response: Final = _chat(candidate, model, token)
|
||||
assert response.status_code == (403 if issuer_behavior == "link_only" else 200), response.text
|
||||
if issuer_behavior == "link_only":
|
||||
assert _user_key_hashes(user) == frozenset()
|
||||
assert read_rows('SELECT token FROM "LiteLLM_JWTKeyMapping" WHERE jwt_claim_value=%s', (user,)) == []
|
||||
else:
|
||||
assert len(_user_key_hashes(user)) == 1
|
||||
assert _mapped_token("sub", user) in _user_key_hashes(user)
|
||||
assert _billed_key(response) == _mapped_token("sub", user)
|
||||
|
||||
|
||||
def test_link_only_links_after_a_cached_rejection_from_the_previous_mode(gateway: Gateway, tmp_path: Path) -> None:
|
||||
with _issuer() as (private_key, jwks_url), gateway.scenario() as scenario:
|
||||
model: Final = scenario.model()
|
||||
user: Final = scenario.user(user_role="internal_user")
|
||||
key: Final = scenario.key(user_id=user)
|
||||
token: Final = _token(private_key, user)
|
||||
with owned_proxy(
|
||||
gateway,
|
||||
tmp_path,
|
||||
{"JWT_PUBLIC_KEY_URL": jwks_url},
|
||||
config=_config(tmp_path, "sub", behavior="reject", map_existing_key=False),
|
||||
) as rejecting:
|
||||
denied: Final = _chat(rejecting, model, token)
|
||||
assert denied.status_code == 403, denied.text
|
||||
with Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) as cache:
|
||||
cached: Final = cache.get(f"jwt_key_mapping:sub:{user}")
|
||||
assert cached is not None and b"__NO_MAPPING__" in cached, cached
|
||||
with owned_proxy(
|
||||
gateway,
|
||||
tmp_path,
|
||||
{"JWT_PUBLIC_KEY_URL": jwks_url},
|
||||
config=_config(tmp_path, "sub", behavior="link_only", map_existing_key=False),
|
||||
) as linking:
|
||||
accepted: Final = _chat(linking, model, token)
|
||||
assert accepted.status_code == 200, accepted.text
|
||||
assert _mapped_token("sub", user) == _hash(key)
|
||||
assert _user_key_hashes(user) == frozenset({_hash(key)})
|
||||
assert _billed_key(accepted) == _hash(key)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue