From e5c446067f27f656cab662df19344748a1a0025e Mon Sep 17 00:00:00 2001 From: yassin Date: Sat, 31 Oct 2026 06:31:42 +0000 Subject: [PATCH] feat(control_plane): service skeleton, settings, license gate, JWT auth and scope RBAC Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/test-unit.yml | 1 + litellm/proxy/control_plane/__init__.py | 0 litellm/proxy/control_plane/app.py | 85 ++++ litellm/proxy/control_plane/auth/__init__.py | 0 .../proxy/control_plane/auth/dependencies.py | 43 ++ litellm/proxy/control_plane/auth/jwt.py | 330 +++++++++++++ litellm/proxy/control_plane/auth/login.py | 130 ++++++ litellm/proxy/control_plane/auth/rbac.py | 191 ++++++++ litellm/proxy/control_plane/auth/router.py | 54 +++ .../proxy/control_plane/common/__init__.py | 0 .../proxy/control_plane/common/exceptions.py | 88 ++++ litellm/proxy/control_plane/container.py | 83 ++++ .../proxy/control_plane/health/__init__.py | 0 litellm/proxy/control_plane/health/router.py | 17 + litellm/proxy/control_plane/settings.py | 98 ++++ litellm/types/proxy/control_plane/__init__.py | 0 litellm/types/proxy/control_plane/auth.py | 34 ++ litellm/types/proxy/control_plane/health.py | 11 + tests/unit/proxy/control_plane/__init__.py | 0 .../unit/proxy/control_plane/auth/__init__.py | 0 .../control_plane/auth/test_dependencies.py | 92 ++++ .../unit/proxy/control_plane/auth/test_jwt.py | 437 ++++++++++++++++++ .../proxy/control_plane/auth/test_login.py | 134 ++++++ .../proxy/control_plane/auth/test_rbac.py | 135 ++++++ .../proxy/control_plane/common/__init__.py | 0 .../proxy/control_plane/health/__init__.py | 0 tests/unit/proxy/control_plane/test_app.py | 259 +++++++++++ .../unit/proxy/control_plane/test_settings.py | 53 +++ 28 files changed, 2275 insertions(+) create mode 100644 litellm/proxy/control_plane/__init__.py create mode 100644 litellm/proxy/control_plane/app.py create mode 100644 litellm/proxy/control_plane/auth/__init__.py create mode 100644 litellm/proxy/control_plane/auth/dependencies.py create mode 100644 litellm/proxy/control_plane/auth/jwt.py create mode 100644 litellm/proxy/control_plane/auth/login.py create mode 100644 litellm/proxy/control_plane/auth/rbac.py create mode 100644 litellm/proxy/control_plane/auth/router.py create mode 100644 litellm/proxy/control_plane/common/__init__.py create mode 100644 litellm/proxy/control_plane/common/exceptions.py create mode 100644 litellm/proxy/control_plane/container.py create mode 100644 litellm/proxy/control_plane/health/__init__.py create mode 100644 litellm/proxy/control_plane/health/router.py create mode 100644 litellm/proxy/control_plane/settings.py create mode 100644 litellm/types/proxy/control_plane/__init__.py create mode 100644 litellm/types/proxy/control_plane/auth.py create mode 100644 litellm/types/proxy/control_plane/health.py create mode 100644 tests/unit/proxy/control_plane/__init__.py create mode 100644 tests/unit/proxy/control_plane/auth/__init__.py create mode 100644 tests/unit/proxy/control_plane/auth/test_dependencies.py create mode 100644 tests/unit/proxy/control_plane/auth/test_jwt.py create mode 100644 tests/unit/proxy/control_plane/auth/test_login.py create mode 100644 tests/unit/proxy/control_plane/auth/test_rbac.py create mode 100644 tests/unit/proxy/control_plane/common/__init__.py create mode 100644 tests/unit/proxy/control_plane/health/__init__.py create mode 100644 tests/unit/proxy/control_plane/test_app.py create mode 100644 tests/unit/proxy/control_plane/test_settings.py diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 0a490bb8f45..39d57f56308 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -499,6 +499,7 @@ jobs: - shard: proxy-auth test-path: |- tests/unit/proxy/auth + tests/unit/proxy/control_plane --ignore=tests/unit/proxy/auth/test_auth_checks.py --ignore=tests/unit/proxy/auth/test_user_api_key_auth.py --ignore=tests/unit/proxy/auth/test_default_end_user_budget_simple.py diff --git a/litellm/proxy/control_plane/__init__.py b/litellm/proxy/control_plane/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/proxy/control_plane/app.py b/litellm/proxy/control_plane/app.py new file mode 100644 index 00000000000..6aa52e2618f --- /dev/null +++ b/litellm/proxy/control_plane/app.py @@ -0,0 +1,85 @@ +"""App factory for the LiteLLM global control plane. Enterprise only: refuses to build without a valid license.""" + +from contextvars import ContextVar +from typing import Final, Protocol + +from fastapi import FastAPI +from starlette.types import ASGIApp, Message, Receive, Scope, Send + +from litellm.proxy.control_plane.auth.dependencies import get_token_verifier +from litellm.proxy.control_plane.auth.router import get_login_service +from litellm.proxy.control_plane.auth.router import router as auth_router +from litellm.proxy.control_plane.common.exceptions import ( + error_body, + error_response, + install_exception_handlers, +) +from litellm.proxy.control_plane.container import Container +from litellm.proxy.control_plane.health.router import router as health_router +from litellm.proxy.middleware.request_size_limit_middleware import RequestSizeLimitMiddleware + +TITLE: Final = "LiteLLM Control Plane" +MAX_REQUEST_SIZE_MB: Final = 1 +REQUEST_TOO_LARGE: Final = error_response( + error_body("request-too-large", "Payload Too Large", 413, "Request size is too large. Max size is 1 MB.") +) + + +class License(Protocol): + def is_premium(self) -> bool: ... + + +class LicenseRequired(RuntimeError): + def __init__(self) -> None: + super().__init__( + "The LiteLLM control plane is an enterprise feature. Set LITELLM_LICENSE to a valid license key." + ) + + +class ControlPlaneRequestSizeLimitMiddleware: + def __init__(self, app: ASGIApp) -> None: + self._middleware = RequestSizeLimitMiddleware( + app, + get_max_request_size_mb=lambda: MAX_REQUEST_SIZE_MB, + is_request_size_limit_enabled=lambda: True, + ) + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + replacing_413: Final = ContextVar("control_plane_request_too_large", default=False) + + async def send_problem(message: Message) -> None: + if message["type"] == "http.response.start" and message["status"] == 413: + replacing_413.set(True) + await send( + { + "type": "http.response.start", + "status": REQUEST_TOO_LARGE.status_code, + "headers": list(REQUEST_TOO_LARGE.raw_headers), + } + ) + return + if replacing_413.get() and message["type"] == "http.response.body": + if not message.get("more_body", False): + await send({"type": "http.response.body", "body": REQUEST_TOO_LARGE.body, "more_body": False}) + return + await send(message) + + await self._middleware(scope, receive, send_problem) + + +def create_app(container: Container, license_check: License) -> FastAPI: + if not license_check.is_premium(): + raise LicenseRequired() + app: Final = FastAPI( + title=TITLE, + version="1", + docs_url=f"{container.settings.api_prefix}/docs", + openapi_url=f"{container.settings.api_prefix}/openapi.json", + ) + app.add_middleware(ControlPlaneRequestSizeLimitMiddleware) + app.dependency_overrides[get_token_verifier] = lambda: container.token_verifier + app.dependency_overrides[get_login_service] = lambda: container.login_service + install_exception_handlers(app) + app.include_router(health_router) + app.include_router(auth_router, prefix=container.settings.api_prefix) + return app diff --git a/litellm/proxy/control_plane/auth/__init__.py b/litellm/proxy/control_plane/auth/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/proxy/control_plane/auth/dependencies.py b/litellm/proxy/control_plane/auth/dependencies.py new file mode 100644 index 00000000000..5c93828e6f1 --- /dev/null +++ b/litellm/proxy/control_plane/auth/dependencies.py @@ -0,0 +1,43 @@ +"""FastAPI glue: bearer header -> Principal, scope checks via Security.""" + +from typing import Annotated, Final + +from fastapi import Depends +from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer, SecurityScopes +from starlette.concurrency import run_in_threadpool + +from litellm._logging import verbose_proxy_logger +from litellm.proxy.control_plane.auth.jwt import Rejected, TokenVerifier +from litellm.proxy.control_plane.auth.rbac import Principal, Scope +from litellm.proxy.control_plane.common.exceptions import ControlPlaneException, forbidden, unauthorized + +bearer_scheme: Final = HTTPBearer(auto_error=False, scheme_name="ControlPlaneBearer") + + +def get_token_verifier() -> TokenVerifier: + """Placeholder bound in create_app via dependency_overrides, the FastAPI way to inject a per-process singleton.""" + raise RuntimeError("get_token_verifier is not bound, build the app with create_app") + + +async def get_principal( + credentials: Annotated[HTTPAuthorizationCredentials | None, Depends(bearer_scheme)], + verifier: Annotated[TokenVerifier, Depends(get_token_verifier)], +) -> Principal: + if credentials is None: + raise ControlPlaneException(unauthorized("Missing bearer credentials.")) + verified: Final = await run_in_threadpool(verifier.verify, credentials.credentials) + if isinstance(verified, Rejected): + verbose_proxy_logger.debug("Control-plane bearer verification rejected: %s", verified.reason) + raise ControlPlaneException(unauthorized("Invalid bearer credentials.")) + return verified + + +async def require_scopes( + security_scopes: SecurityScopes, + principal: Annotated[Principal, Depends(get_principal)], +) -> Principal: + required: Final = frozenset(Scope(scope) for scope in security_scopes.scopes) + missing: Final = principal.missing_scopes(required) + if missing: + raise ControlPlaneException(forbidden(f"Missing scope: {', '.join(sorted(scope.value for scope in missing))}.")) + return principal diff --git a/litellm/proxy/control_plane/auth/jwt.py b/litellm/proxy/control_plane/auth/jwt.py new file mode 100644 index 00000000000..0ff78dd65e8 --- /dev/null +++ b/litellm/proxy/control_plane/auth/jwt.py @@ -0,0 +1,330 @@ +"""Token verification, independent of FastAPI. Session JWTs are HS256 under the master key, OIDC JWTs are +verified against the IdP's JWKS. Failures are values so the HTTP layer decides the status code.""" + +import secrets +import threading +import time +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone +from typing import Final, Literal, Protocol, TypeAlias + +import httpx +import jwt +from jwt import PyJWK, PyJWKClient +from pydantic import TypeAdapter, ValidationError + +from litellm.proxy.control_plane.auth.rbac import ( + ADMIN_ROLE, + ALL_SCOPES, + DataplaneBinding, + Principal, + RolePolicy, + build_principal, + parse_scopes, +) +from litellm.proxy.control_plane.settings import ControlPlaneSettings, OidcSettings + +SESSION_ALGORITHM: Final = "HS256" +SESSION_AUDIENCE: Final = "litellm-control-plane-api" +OIDC_ALGORITHMS: Final = ("RS256", "RS384", "RS512", "ES256", "ES384", "ES512", "PS256", "PS384", "PS512") + +_STRING_LIST: Final = TypeAdapter(list[str]) +_STRING: Final = TypeAdapter(str) +_CLAIMS: Final = TypeAdapter(Mapping[str, object]) +_DISCOVERY: Final = TypeAdapter(Mapping[str, object]) + + +@dataclass(frozen=True, slots=True) +class Rejected: + reason: str + + +VerifiedToken: TypeAlias = Principal | Rejected + + +SigningKey: TypeAlias = PyJWK | str +JWKS_FETCH_TIMEOUT_SECONDS: Final = 10 +DISCOVERY_FAILURE_COOLDOWN_SECONDS: Final = 30 + + +class JwksClient(Protocol): + def get_signing_key_from_jwt(self, token: str) -> SigningKey: ... + + +class JwksKeyLookup(Protocol): + def signing_key(self, token: str) -> SigningKey | Rejected: ... + + +DiscoveryFetch: TypeAlias = Callable[[str], Mapping[str, object]] +JwksClientFactory: TypeAlias = Callable[[str], JwksClient] + + +def _fetch_oidc_configuration(url: str) -> Mapping[str, object]: + response: Final = httpx.get(url, timeout=JWKS_FETCH_TIMEOUT_SECONDS) + response.raise_for_status() + return _DISCOVERY.validate_json(response.content) + + +def _new_jwks_client(url: str) -> JwksClient: + return PyJWKClient(url, cache_keys=False, timeout=JWKS_FETCH_TIMEOUT_SECONDS) + + +class PyJwkClientLookup: + def __init__( + self, + jwks_url: str | None, + issuer: str, + fetch: DiscoveryFetch = _fetch_oidc_configuration, + client_factory: JwksClientFactory = _new_jwks_client, + clock: Callable[[], float] = time.monotonic, + ) -> None: + self._jwks_url = jwks_url + self._issuer = issuer + self._fetch = fetch + self._client_factory = client_factory + self._clock = clock + self._discovery_lock = threading.Lock() + self._client: JwksClient | None = None + self._failure_until: float | None = None + + def signing_key(self, token: str) -> SigningKey | Rejected: + client: Final = self._get_client() + if isinstance(client, Rejected): + return client + try: + return client.get_signing_key_from_jwt(token) + except (jwt.PyJWKClientError, OSError, ValueError, jwt.PyJWTError) as error: + return Rejected(f"jwks lookup failed: {error}") + + def _get_client(self) -> JwksClient | Rejected: + with self._discovery_lock: + if self._client is not None: + return self._client + if self._failure_until is not None and self._clock() < self._failure_until: + return Rejected("oidc discovery is cooling down after a previous failure") + try: + jwks_url: Final = self._jwks_url or _as_optional_string( + self._fetch(f"{self._issuer.rstrip('/')}/.well-known/openid-configuration").get("jwks_uri") + ) + if jwks_url is None: + self._failure_until = self._clock() + DISCOVERY_FAILURE_COOLDOWN_SECONDS + return Rejected("oidc discovery did not include a valid jwks_uri") + client: Final = self._client_factory(jwks_url) + except (httpx.HTTPError, jwt.PyJWKClientError, OSError, ValueError, jwt.PyJWTError) as error: + self._failure_until = self._clock() + DISCOVERY_FAILURE_COOLDOWN_SECONDS + return Rejected(f"oidc discovery failed: {error}") + self._jwks_url = jwks_url + self._client = client + return client + + +def _as_string_tuple(value: object) -> tuple[str, ...]: + if isinstance(value, str): + return tuple(value.split()) + try: + return tuple(_STRING_LIST.validate_python(value)) + except ValidationError: + return () + + +def _as_optional_string(value: object) -> str | None: + try: + return _STRING.validate_python(value) + except ValidationError: + return None + + +def _decode( + token: str, key: SigningKey, algorithms: tuple[str, ...], audience: str | None, issuer: str, leeway: int +) -> Mapping[str, object] | Rejected: + try: + raw: Final = jwt.decode( + token, + key, + algorithms=list(algorithms), + audience=audience, + issuer=issuer, + leeway=leeway, + options={"require": ["exp", "sub"]}, + ) + return _CLAIMS.validate_python(raw) + except jwt.PyJWTError as error: + return Rejected(f"invalid token: {error}") + except ValidationError: + return Rejected("invalid token: claims are not an object") + + +class SessionTokenCodec: + """Mints and verifies the control plane's own session JWTs. The audience pins a token to this API, so a + token minted under the same master key by another component never authorizes a control plane route.""" + + def __init__( + self, master_key: str, issuer: str, ttl: timedelta, policy: RolePolicy, audience: str = SESSION_AUDIENCE + ) -> None: + self._master_key = master_key + self._issuer = issuer + self._audience = audience + self._ttl = ttl + self._policy = policy + + def mint(self, principal: Principal, now: datetime | None = None) -> str: + issued_at: Final = now or datetime.now(timezone.utc) + claims: Final = { + "iss": self._issuer, + "aud": self._audience, + "sub": principal.subject, + "email": principal.email, + "roles": sorted(principal.roles), + "scope": " ".join(sorted(scope.value for scope in principal.scopes)), + "litellm_dataplanes": ["*"] if principal.dataplanes.all else sorted(principal.dataplanes.ids), + "iat": int(issued_at.timestamp()), + "exp": int((issued_at + self._ttl).timestamp()), + "jti": secrets.token_urlsafe(16), + } + return jwt.encode(claims, self._master_key, algorithm=SESSION_ALGORITHM) + + def verify(self, token: str) -> VerifiedToken: + claims: Final = _decode(token, self._master_key, (SESSION_ALGORITHM,), self._audience, self._issuer, 0) + if isinstance(claims, Rejected): + return claims + return _principal_from_claims(self._policy, claims, "roles", "scope", "litellm_dataplanes", "email", "session") + + +class OidcTokenVerifier: + def __init__(self, settings: OidcSettings, keys: JwksKeyLookup, policy: RolePolicy) -> None: + self._settings = settings + self._keys = keys + self._policy = policy + + @property + def issuer(self) -> str: + return self._settings.issuer + + def verify(self, token: str) -> VerifiedToken: + key: Final = self._keys.signing_key(token) + if isinstance(key, Rejected): + return key + claims: Final = _decode( + token, key, OIDC_ALGORITHMS, self._settings.audience, self._settings.issuer, self._settings.leeway_seconds + ) + if isinstance(claims, Rejected): + return claims + return _principal_from_claims( + self._policy, + claims, + self._settings.roles_claim, + self._settings.scope_claim, + self._settings.dataplanes_claim, + self._settings.email_claim, + "oidc", + ) + + +def _principal_from_claims( + policy: RolePolicy, + claims: Mapping[str, object], + roles_claim: str, + scope_claim: str, + dataplanes_claim: str, + email_claim: str, + token_kind: Literal["session", "oidc"], +) -> VerifiedToken: + subject: Final = _as_optional_string(claims.get("sub")) + if subject is None: + return Rejected("invalid token: sub is not a string") + roles: Final = policy.known(_as_string_tuple(claims.get(roles_claim))) + if not roles: + return Rejected("token carries no control plane role") + dataplanes: Final = _dataplane_ids(claims, dataplanes_claim) + if dataplanes is None: + return Rejected(f"invalid token: {dataplanes_claim} is not a string list") + raw_scopes: Final = claims.get(scope_claim) + return build_principal( + policy=policy, + subject=subject, + email=_as_optional_string(claims.get(email_claim)), + roles=roles, + claimed_scopes=None if raw_scopes is None else parse_scopes(_as_string_tuple(raw_scopes)), + dataplane_ids=frozenset(dataplanes), + token_kind=token_kind, + ) + + +def _dataplane_ids(claims: Mapping[str, object], claim: str) -> tuple[str, ...] | None: + if claim not in claims: + return () + raw: Final = claims[claim] + if isinstance(raw, str): + return tuple(raw.split()) + try: + return tuple(_STRING_LIST.validate_python(raw)) + except ValidationError: + return None + + +MASTER_KEY_SUBJECT: Final = "master_key" + + +def master_key_principal() -> Principal: + """The master key is the break-glass credential and holds every scope even when operators drop the default + ``admin`` role from the policy.""" + return Principal( + subject=MASTER_KEY_SUBJECT, + email=None, + roles=frozenset({ADMIN_ROLE}), + scopes=ALL_SCOPES, + dataplanes=DataplaneBinding(all=True, ids=frozenset()), + token_kind="master_key", + ) + + +class TokenVerifier: + """Resolves a bearer credential of any supported kind to a Principal.""" + + def __init__( + self, master_key: str, sessions: SessionTokenCodec, oidc: OidcTokenVerifier | None, policy: RolePolicy + ) -> None: + self._master_key = master_key + self._sessions = sessions + self._oidc = oidc + self._policy = policy + + def verify(self, credential: str) -> VerifiedToken: + if secrets.compare_digest(credential.encode("utf-8"), self._master_key.encode("utf-8")): + return master_key_principal() + if not _looks_like_jwt(credential): + return Rejected("credential is neither the master key nor a JWT") + if self._oidc is not None and _issuer_of(credential) == self._oidc.issuer: + return self._oidc.verify(credential) + return self._sessions.verify(credential) + + +def _looks_like_jwt(credential: str) -> bool: + return credential.count(".") == 2 + + +def _issuer_of(credential: str) -> str | None: + try: + unverified: Final = jwt.decode(credential, options={"verify_signature": False}) + except jwt.PyJWTError: + return None + return _as_optional_string(unverified.get("iss")) + + +def build_token_verifier( + settings: ControlPlaneSettings, policy: RolePolicy, keys: JwksKeyLookup | None = None +) -> TokenVerifier: + master_key: Final = settings.master_key.get_secret_value() + sessions: Final = SessionTokenCodec( + master_key=master_key, + issuer=settings.session_issuer, + audience=settings.session_audience, + ttl=timedelta(seconds=settings.session_ttl_seconds), + policy=policy, + ) + oidc_settings: Final = settings.oidc + if oidc_settings is None: + return TokenVerifier(master_key, sessions, None, policy) + lookup: Final = keys or PyJwkClientLookup(oidc_settings.jwks_url, oidc_settings.issuer) + return TokenVerifier(master_key, sessions, OidcTokenVerifier(oidc_settings, lookup, policy), policy) diff --git a/litellm/proxy/control_plane/auth/login.py b/litellm/proxy/control_plane/auth/login.py new file mode 100644 index 00000000000..9218adeb70b --- /dev/null +++ b/litellm/proxy/control_plane/auth/login.py @@ -0,0 +1,130 @@ +"""Username and password login for the admin configured through UI_USERNAME / UI_PASSWORD.""" + +import secrets +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from math import ceil +from types import MappingProxyType +from typing import Final, TypeAlias + +from litellm.proxy.control_plane.auth.jwt import SessionTokenCodec +from litellm.proxy.control_plane.auth.rbac import ADMIN_ROLE, RolePolicy, build_principal +from litellm.types.proxy.control_plane.auth import SessionTokenResponse + +LOCKOUT_FAILURES: Final = 5 +LOCKOUT_SECONDS: Final = 300.0 +MAX_TRACKED_LOGIN_KEYS: Final = 10_000 + + +@dataclass(frozen=True, slots=True) +class InvalidCredentials: + pass + + +@dataclass(frozen=True, slots=True) +class Throttled: + retry_after_seconds: int + + +LoginResult: TypeAlias = SessionTokenResponse | InvalidCredentials | Throttled + + +class LoginThrottle: + """Locks a source-and-account pair after repeated misses.""" + + def __init__( + self, now: Callable[[], float], failures: int = LOCKOUT_FAILURES, lockout_seconds: float = LOCKOUT_SECONDS + ) -> None: + self._now = now + self._failures = failures + self._lockout_seconds = lockout_seconds + self._misses: Mapping[tuple[str, str], tuple[float, ...]] = MappingProxyType({}) + + def _recent(self, key: tuple[str, str]) -> tuple[float, ...]: + cutoff: Final = self._now() - self._lockout_seconds + return tuple(at for at in self._misses.get(key, ()) if at > cutoff) + + def locked_for(self, key: tuple[str, str]) -> float: + recent: Final = self._recent(key) + if len(recent) < self._failures: + return 0.0 + return max(0.0, recent[-1] + self._lockout_seconds - self._now()) + + def record_failure(self, key: tuple[str, str]) -> None: + now: Final = self._now() + cutoff: Final = now - self._lockout_seconds + active: Final = { + tracked_key: tuple(at for at in misses if at > cutoff) + for tracked_key, misses in self._misses.items() + if any(at > cutoff for at in misses) + } + updated: Final = {**active, key: (*active.get(key, ()), now)} + self._misses = self._bounded(updated) + + @staticmethod + def _bounded( + misses: Mapping[tuple[str, str], tuple[float, ...]], + ) -> Mapping[tuple[str, str], tuple[float, ...]]: + if len(misses) <= MAX_TRACKED_LOGIN_KEYS: + return MappingProxyType(dict(misses)) + oldest: Final = min(misses, key=lambda tracked_key: misses[tracked_key][-1]) + return MappingProxyType( + {tracked_key: values for tracked_key, values in misses.items() if tracked_key != oldest} + ) + + def tracked(self) -> int: + return len(self._misses) + + def reset(self, key: tuple[str, str]) -> None: + self._misses = MappingProxyType( + {tracked_key: misses for tracked_key, misses in self._misses.items() if tracked_key != key} + ) + + +class LoginService: + def __init__( + self, + sessions: SessionTokenCodec, + policy: RolePolicy, + ttl_seconds: int, + admin_username: str | None, + admin_password: str | None, + throttle: LoginThrottle, + ) -> None: + self._sessions = sessions + self._policy = policy + self._ttl_seconds = ttl_seconds + self._admin_username = admin_username + self._admin_password = admin_password + self._throttle = throttle + + def login(self, username: str, password: str, client_source: str = "") -> LoginResult: + if self._admin_username is None or self._admin_password is None: + return InvalidCredentials() + if not secrets.compare_digest(username.encode("utf-8"), self._admin_username.encode("utf-8")): + return InvalidCredentials() + throttle_key: Final = (client_source, self._admin_username) + locked: Final = self._throttle.locked_for(throttle_key) + if locked > 0: + return Throttled(retry_after_seconds=max(1, ceil(locked))) + if not self._credentials_match(username, password): + self._throttle.record_failure(throttle_key) + return InvalidCredentials() + self._throttle.reset(throttle_key) + principal: Final = build_principal( + policy=self._policy, + subject=username, + email=None, + roles=frozenset({ADMIN_ROLE}), + claimed_scopes=None, + dataplane_ids=frozenset(), + token_kind="session", + ) + return SessionTokenResponse(access_token=self._sessions.mint(principal), expires_in=self._ttl_seconds) + + def _credentials_match(self, username: str, password: str) -> bool: + if self._admin_username is None or self._admin_password is None: + return False + username_ok: Final = secrets.compare_digest(username.encode("utf-8"), self._admin_username.encode("utf-8")) + password_ok: Final = secrets.compare_digest(password.encode("utf-8"), self._admin_password.encode("utf-8")) + return username_ok and password_ok diff --git a/litellm/proxy/control_plane/auth/rbac.py b/litellm/proxy/control_plane/auth/rbac.py new file mode 100644 index 00000000000..4e8f680dfde --- /dev/null +++ b/litellm/proxy/control_plane/auth/rbac.py @@ -0,0 +1,191 @@ +"""Role and scope policy for the control plane. Roles grant scopes, scopes gate routes, the dataplane binding +on the token decides where the scopes apply (``*`` or no binding is fleet-wide).""" + +from collections.abc import Mapping +from dataclasses import dataclass +from enum import Enum +from types import MappingProxyType +from typing import Final, Literal, TypeAlias + + +class Scope(str, Enum): + DATAPLANES_READ = "dataplanes:read" + DATAPLANES_WRITE = "dataplanes:write" + USAGE_READ = "usage:read" + FLEET_READ = "fleet:read" + PROXY_READ = "proxy:read" + PROXY_WRITE = "proxy:write" + AUDIT_READ = "audit:read" + ROLES_WRITE = "roles:write" + + +ALL_SCOPES: Final = frozenset(Scope) +READ_SCOPES: Final = frozenset({Scope.DATAPLANES_READ, Scope.USAGE_READ, Scope.FLEET_READ, Scope.PROXY_READ}) + +ADMIN_ROLE: Final = "admin" +EDITOR_ROLE: Final = "editor" +VIEWER_ROLE: Final = "viewer" +USAGE_VIEWER_ROLE: Final = "usage_viewer" + +DEFAULT_ROLES: Final[Mapping[str, frozenset[Scope]]] = MappingProxyType( + { + ADMIN_ROLE: ALL_SCOPES, + EDITOR_ROLE: READ_SCOPES | {Scope.PROXY_WRITE}, + VIEWER_ROLE: READ_SCOPES, + USAGE_VIEWER_ROLE: frozenset({Scope.USAGE_READ}), + } +) + +FLEET_WIDE_BINDING: Final = "*" + +TokenKind: TypeAlias = Literal["session", "oidc", "master_key"] + + +@dataclass(frozen=True, slots=True) +class RoleDefinition: + scopes: frozenset[Scope] + inherits: frozenset[str] + + +@dataclass(frozen=True, slots=True) +class UnknownScope: + role: str + scope: str + + +@dataclass(frozen=True, slots=True) +class UnknownParent: + role: str + parent: str + + +RolePolicyError: TypeAlias = UnknownScope | UnknownParent + + +class RolePolicy: + """Resolved role -> scopes table. Built once from defaults plus operator-defined roles.""" + + def __init__(self, roles: Mapping[str, frozenset[Scope]]) -> None: + self._roles: Final = MappingProxyType(dict(roles)) + + @property + def names(self) -> frozenset[str]: + return frozenset(self._roles) + + def scopes_for(self, role: str) -> frozenset[Scope]: + return self._roles.get(role, frozenset()) + + def known(self, raw: tuple[str, ...]) -> frozenset[str]: + return frozenset(name for name in raw if name in self._roles) + + def grant(self, roles: frozenset[str]) -> frozenset[Scope]: + return frozenset[Scope]().union(*(self.scopes_for(role) for role in roles)) + + +def default_policy() -> RolePolicy: + return RolePolicy(DEFAULT_ROLES) + + +def _resolve(name: str, definitions: Mapping[str, RoleDefinition], seen: frozenset[str]) -> frozenset[Scope]: + definition: Final = definitions[name] + parents: Final = definition.inherits - seen + return definition.scopes.union(*(_resolve(parent, definitions, seen | {name}) for parent in parents)) + + +def build_policy(custom: Mapping[str, RoleDefinition], include_defaults: bool = True) -> RolePolicy | RolePolicyError: + """Operator roles override same-named defaults. Inheritance may reference defaults or other custom roles.""" + base: Final[Mapping[str, RoleDefinition]] = ( + {name: RoleDefinition(scopes=scopes, inherits=frozenset()) for name, scopes in DEFAULT_ROLES.items()} + if include_defaults + else {} + ) + definitions: Final = {**base, **custom} + unknown_parents: Final = tuple( + UnknownParent(role=name, parent=parent) + for name, definition in custom.items() + for parent in sorted(definition.inherits - frozenset(definitions)) + ) # comprehension-ok: flattening (role, parent) pairs + if unknown_parents: + return unknown_parents[0] + return RolePolicy({name: _resolve(name, definitions, frozenset()) for name in definitions}) + + +def parse_role_definition( + role: str, scopes: tuple[str, ...], inherits: tuple[str, ...] +) -> RoleDefinition | UnknownScope: + known: Final = frozenset(scope.value for scope in Scope) + unknown: Final = tuple(scope for scope in scopes if scope not in known) + if unknown: + return UnknownScope(role=role, scope=unknown[0]) + return RoleDefinition(scopes=frozenset(Scope(scope) for scope in scopes), inherits=frozenset(inherits)) + + +@dataclass(frozen=True, slots=True) +class DataplaneBinding: + """Which dataplanes a principal may touch. ``all`` is fleet-wide, ``ids`` is an explicit set.""" + + all: bool + ids: frozenset[str] + + def allows(self, dataplane_id: str) -> bool: + return self.all or dataplane_id in self.ids + + def restrict(self, requested: frozenset[str] | None) -> frozenset[str] | None: + """Intersect a caller's requested dataplane filter with the binding. ``None`` means every bound dataplane.""" + if requested is None: + return None if self.all else self.ids + return requested if self.all else requested & self.ids + + +@dataclass(frozen=True, slots=True) +class Principal: + subject: str + email: str | None + roles: frozenset[str] + scopes: frozenset[Scope] + dataplanes: DataplaneBinding + token_kind: TokenKind + + def has_scope(self, scope: Scope) -> bool: + return scope in self.scopes + + def missing_scopes(self, required: frozenset[Scope]) -> frozenset[Scope]: + return required - self.scopes + + +def clamp_scopes(granted: frozenset[Scope], claimed: frozenset[Scope] | None) -> frozenset[Scope]: + """A token may narrow what its roles grant, never widen it. ``None`` is a token without a scope claim and + keeps the full role grant. A claim that names only unknown scopes narrows the grant to nothing.""" + return granted if claimed is None else granted & claimed + + +def binding_for(dataplane_ids: frozenset[str]) -> DataplaneBinding: + fleet_wide: Final = not dataplane_ids or FLEET_WIDE_BINDING in dataplane_ids + return DataplaneBinding(all=fleet_wide, ids=frozenset() if fleet_wide else dataplane_ids) + + +def parse_scopes(raw: tuple[str, ...] | None) -> frozenset[Scope] | None: + """Unknown scope names (``openid``, ``profile``) are dropped. Only an absent claim is ``None``.""" + if raw is None: + return None + known: Final = frozenset(scope.value for scope in Scope) + return frozenset(Scope(value) for value in raw if value in known) + + +def build_principal( + policy: RolePolicy, + subject: str, + email: str | None, + roles: frozenset[str], + claimed_scopes: frozenset[Scope] | None, + dataplane_ids: frozenset[str], + token_kind: TokenKind, +) -> Principal: + return Principal( + subject=subject, + email=email, + roles=roles, + scopes=clamp_scopes(policy.grant(roles), claimed_scopes), + dataplanes=binding_for(dataplane_ids), + token_kind=token_kind, + ) diff --git a/litellm/proxy/control_plane/auth/router.py b/litellm/proxy/control_plane/auth/router.py new file mode 100644 index 00000000000..90f541dd641 --- /dev/null +++ b/litellm/proxy/control_plane/auth/router.py @@ -0,0 +1,54 @@ +from typing import Annotated, Final + +from fastapi import APIRouter, Depends, Request, Security +from fastapi.responses import JSONResponse + +from litellm.proxy.control_plane.auth.dependencies import require_scopes +from litellm.proxy.control_plane.auth.login import InvalidCredentials, LoginService, Throttled +from litellm.proxy.control_plane.auth.rbac import Principal +from litellm.proxy.control_plane.common.exceptions import ControlPlaneException, too_many_requests, unauthorized +from litellm.types.proxy.control_plane.auth import LoginRequest, PrincipalResponse, SessionTokenResponse + +router: Final = APIRouter(prefix="/auth", tags=["control plane: auth"]) + + +def get_login_service() -> LoginService: + raise RuntimeError("get_login_service is not bound, build the app with create_app") + + +def principal_response(principal: Principal) -> PrincipalResponse: + return PrincipalResponse( + subject=principal.subject, + email=principal.email, + roles=tuple(sorted(principal.roles)), + scopes=tuple(sorted(scope.value for scope in principal.scopes)), + dataplanes=None if principal.dataplanes.all else tuple(sorted(principal.dataplanes.ids)), + token_kind=principal.token_kind, + ) + + +@router.post("/login", response_model=SessionTokenResponse) +async def login( + body: LoginRequest, + service: Annotated[LoginService, Depends(get_login_service)], + request: Request, +) -> JSONResponse: + client_source: Final = request.client.host if request.client is not None else "" + match service.login(body.username, body.password, client_source): + case SessionTokenResponse() as token: + return JSONResponse( + content=token.model_dump(mode="json"), + headers={"Cache-Control": "no-store", "Pragma": "no-cache"}, + ) + case Throttled(retry_after_seconds=seconds): + raise ControlPlaneException( + too_many_requests(f"Too many failed logins, retry in {seconds} seconds."), + headers={"Retry-After": str(seconds)}, + ) + case InvalidCredentials(): + raise ControlPlaneException(unauthorized("Invalid username or password.")) + + +@router.get("/me", response_model=PrincipalResponse) +async def me(principal: Annotated[Principal, Security(require_scopes)]) -> PrincipalResponse: + return principal_response(principal) diff --git a/litellm/proxy/control_plane/common/__init__.py b/litellm/proxy/control_plane/common/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/proxy/control_plane/common/exceptions.py b/litellm/proxy/control_plane/common/exceptions.py new file mode 100644 index 00000000000..d069f087537 --- /dev/null +++ b/litellm/proxy/control_plane/common/exceptions.py @@ -0,0 +1,88 @@ +"""RFC 9457 error bodies for the control plane, same contract as /management/v1.""" + +from collections.abc import Mapping +from http import HTTPStatus +from types import MappingProxyType +from typing import Final + +from fastapi import FastAPI, Request +from fastapi.exceptions import RequestValidationError +from fastapi.responses import JSONResponse +from starlette.exceptions import HTTPException as StarletteHTTPException + +from litellm.proxy.common_utils.validation_error_body import public_validation_errors +from litellm.proxy.list_api.common import PROBLEM_TYPE_BASE as ERROR_TYPE_BASE +from litellm.proxy.list_api.common import problem_response as error_response +from litellm.proxy.list_api.common import request_validation_problem as request_validation_error +from litellm.types.proxy.management_endpoints.management_v1 import ProblemDetail as ErrorBody + +WWW_AUTHENTICATE: Final = "WWW-Authenticate" +BEARER_CHALLENGE: Final = 'Bearer realm="litellm-control-plane"' +EMPTY_HEADERS: Final[Mapping[str, str]] = MappingProxyType({}) + + +def error_body(slug: str, title: str, status: int, detail: str) -> ErrorBody: + return ErrorBody(type=f"{ERROR_TYPE_BASE}{slug}", title=title, status=status, detail=detail) + + +def unauthorized(detail: str) -> ErrorBody: + return error_body("unauthorized", "Unauthorized", 401, detail) + + +def forbidden(detail: str) -> ErrorBody: + return error_body("forbidden", "Forbidden", 403, detail) + + +def not_found(detail: str) -> ErrorBody: + return error_body("not-found", "Not Found", 404, detail) + + +def too_many_requests(detail: str) -> ErrorBody: + return error_body("too-many-requests", "Too Many Requests", 429, detail) + + +class ControlPlaneException(Exception): + """An error raised from a dependency or handler, rendered by the app-level handler below.""" + + def __init__(self, body: ErrorBody, *, headers: Mapping[str, str] = EMPTY_HEADERS) -> None: + self.body = body + self.headers = MappingProxyType(dict(headers)) + super().__init__(body.detail) + + +def _render(body: ErrorBody) -> JSONResponse: + response: Final = error_response(body) + if body.status == 401: + response.headers[WWW_AUTHENTICATE] = BEARER_CHALLENGE + return response + + +async def _handle_exception(_request: Request, exc: Exception) -> JSONResponse: + if not isinstance(exc, ControlPlaneException): + raise exc + response: Final = _render(exc.body) + response.headers.update(exc.headers) + return response + + +async def _handle_validation(_request: Request, exc: Exception) -> JSONResponse: + if not isinstance(exc, RequestValidationError): + raise exc + return _render(request_validation_error(public_validation_errors(exc.errors()))) + + +async def _handle_http_exception(_request: Request, exc: Exception) -> JSONResponse: + if not isinstance(exc, StarletteHTTPException): + raise exc + phrase: Final = HTTPStatus(exc.status_code).phrase + slug: Final = phrase.lower().replace("'", "").replace(" ", "-") + body: Final = error_body(slug, phrase, exc.status_code, str(exc.detail)) + response: Final = error_response(body) + response.headers.update(exc.headers or {}) + return response + + +def install_exception_handlers(app: FastAPI) -> None: + app.add_exception_handler(ControlPlaneException, _handle_exception) + app.add_exception_handler(RequestValidationError, _handle_validation) + app.add_exception_handler(StarletteHTTPException, _handle_http_exception) diff --git a/litellm/proxy/control_plane/container.py b/litellm/proxy/control_plane/container.py new file mode 100644 index 00000000000..e1caa26f485 --- /dev/null +++ b/litellm/proxy/control_plane/container.py @@ -0,0 +1,83 @@ +"""Long-lived objects built once per process and handed to the app. Tests build their own.""" + +import time +from collections.abc import Mapping +from dataclasses import dataclass +from datetime import timedelta +from types import MappingProxyType +from typing import Final + +from litellm.proxy.control_plane.auth.jwt import JwksKeyLookup, SessionTokenCodec, TokenVerifier, build_token_verifier +from litellm.proxy.control_plane.auth.login import LoginService, LoginThrottle +from litellm.proxy.control_plane.auth.rbac import ( + RoleDefinition, + RolePolicy, + RolePolicyError, + UnknownParent, + UnknownScope, + build_policy, + parse_role_definition, +) +from litellm.proxy.control_plane.settings import ControlPlaneSettings + + +@dataclass(frozen=True, slots=True) +class Container: + settings: ControlPlaneSettings + policy: RolePolicy + token_verifier: TokenVerifier + login_service: LoginService + + +class InvalidRolePolicy(ValueError): + def __init__(self, error: RolePolicyError) -> None: + match error: + case UnknownScope(role=role, scope=scope): + super().__init__(f"role {role!r} grants unknown scope {scope!r}") + case UnknownParent(role=role, parent=parent): + super().__init__(f"role {role!r} inherits unknown role {parent!r}") + + +def _parse_roles(settings: ControlPlaneSettings) -> Mapping[str, RoleDefinition]: + definitions: Final = MappingProxyType( + {name: parse_role_definition(name, spec.scopes, spec.inherits) for name, spec in settings.roles.items()} + ) + for error in definitions.values(): + if isinstance(error, UnknownScope): + raise InvalidRolePolicy(error) + return MappingProxyType( + {name: definition for name, definition in definitions.items() if isinstance(definition, RoleDefinition)} + ) + + +def build_role_policy(settings: ControlPlaneSettings) -> RolePolicy: + policy: Final = build_policy(_parse_roles(settings), include_defaults=settings.include_default_roles) + if isinstance(policy, RolePolicy): + return policy + raise InvalidRolePolicy(policy) + + +def build_container(settings: ControlPlaneSettings, jwks: JwksKeyLookup | None = None) -> Container: + master_key: Final = settings.master_key.get_secret_value() + policy: Final = build_role_policy(settings) + sessions: Final = SessionTokenCodec( + master_key=master_key, + issuer=settings.session_issuer, + audience=settings.session_audience, + ttl=timedelta(seconds=settings.session_ttl_seconds), + policy=policy, + ) + password: Final = settings.ui_password + return Container( + settings=settings, + policy=policy, + token_verifier=build_token_verifier(settings, policy, jwks), + login_service=LoginService( + sessions=sessions, + policy=policy, + ttl_seconds=settings.session_ttl_seconds, + admin_username=settings.ui_username, + admin_password=None if password is None else password.get_secret_value(), + throttle=LoginThrottle(time.monotonic), + ), + ) diff --git a/litellm/proxy/control_plane/health/__init__.py b/litellm/proxy/control_plane/health/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/proxy/control_plane/health/router.py b/litellm/proxy/control_plane/health/router.py new file mode 100644 index 00000000000..50305295801 --- /dev/null +++ b/litellm/proxy/control_plane/health/router.py @@ -0,0 +1,17 @@ +from typing import Final + +from fastapi import APIRouter + +from litellm.types.proxy.control_plane.health import HealthResponse + +router: Final = APIRouter(tags=["control plane: health"]) + + +@router.get("/healthz", response_model=HealthResponse) +async def healthz() -> HealthResponse: + return HealthResponse(status="ok") + + +@router.get("/readyz", response_model=HealthResponse) +async def readyz() -> HealthResponse: + return HealthResponse(status="ok") diff --git a/litellm/proxy/control_plane/settings.py b/litellm/proxy/control_plane/settings.py new file mode 100644 index 00000000000..6ff7929f0b6 --- /dev/null +++ b/litellm/proxy/control_plane/settings.py @@ -0,0 +1,98 @@ +"""Process-level configuration for the control plane, read once at startup.""" + +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Final + +from pydantic import BaseModel, ConfigDict, Field, SecretStr, field_validator, model_validator +from pydantic_settings import BaseSettings, SettingsConfigDict + +from litellm.proxy.control_plane.auth.rbac import ADMIN_ROLE +from litellm.types.proxy.control_plane.auth import PASSWORD_MAX_LENGTH, USERNAME_MAX_LENGTH + +DEFAULT_SESSION_TTL_SECONDS: Final = 12 * 60 * 60 +DEFAULT_DATAPLANE_TIMEOUT_SECONDS: Final = 10.0 + + +@dataclass(frozen=True, slots=True) +class OidcSettings: + """External IdP used to verify caller JWTs and, later, to exchange them for dataplane tokens.""" + + issuer: str + audience: str + jwks_url: str | None + roles_claim: str + scope_claim: str + dataplanes_claim: str + email_claim: str + leeway_seconds: int + + +class RoleSpec(BaseModel): + """An operator-defined role, e.g. CONTROL_PLANE_ROLES='{"finops": {"scopes": ["usage:read", "audit:read"]}}'.""" + + model_config = ConfigDict(frozen=True, extra="forbid") + + scopes: tuple[str, ...] = () + inherits: tuple[str, ...] = () + + +class ControlPlaneSettings(BaseSettings): + model_config = SettingsConfigDict(frozen=True, env_prefix="CONTROL_PLANE_", extra="ignore") + + master_key: SecretStr = Field(validation_alias="LITELLM_MASTER_KEY") + ui_username: str | None = Field(default=None, max_length=USERNAME_MAX_LENGTH, validation_alias="UI_USERNAME") + ui_password: SecretStr | None = Field(default=None, max_length=PASSWORD_MAX_LENGTH, validation_alias="UI_PASSWORD") + session_ttl_seconds: int = DEFAULT_SESSION_TTL_SECONDS + session_issuer: str = "litellm-control-plane" + session_audience: str = "litellm-control-plane-api" + roles: Mapping[str, RoleSpec] = Field(default_factory=dict) + include_default_roles: bool = True + oidc_issuer: str | None = None + oidc_audience: str | None = None + oidc_jwks_url: str | None = None + oidc_roles_claim: str = "roles" + oidc_scope_claim: str = "scope" + oidc_dataplanes_claim: str = "litellm_dataplanes" + oidc_email_claim: str = "email" + oidc_leeway_seconds: int = 30 + dataplane_timeout_seconds: float = DEFAULT_DATAPLANE_TIMEOUT_SECONDS + fanout_concurrency: int = 16 + api_prefix: str = "/api/v1" + + @field_validator("master_key") + @classmethod + def _master_key_has_minimum_length(cls, value: SecretStr) -> SecretStr: + if len(value.get_secret_value().encode("utf-8")) < 32: + raise ValueError("LITELLM_MASTER_KEY must be at least 32 UTF-8 bytes.") + return value + + @model_validator(mode="after") + def _password_login_has_admin_role(self) -> "ControlPlaneSettings": + if ( + self.ui_username is None + or self.ui_password is None + or self.include_default_roles + or ADMIN_ROLE in self.roles + ): + return self + raise ValueError("password login requires a custom admin role when default roles are disabled.") + + @property + def oidc(self) -> OidcSettings | None: + if self.oidc_issuer is None or self.oidc_audience is None: + return None + return OidcSettings( + issuer=self.oidc_issuer, + audience=self.oidc_audience, + jwks_url=self.oidc_jwks_url, + roles_claim=self.oidc_roles_claim, + scope_claim=self.oidc_scope_claim, + dataplanes_claim=self.oidc_dataplanes_claim, + email_claim=self.oidc_email_claim, + leeway_seconds=self.oidc_leeway_seconds, + ) + + +def load_settings() -> ControlPlaneSettings: + return ControlPlaneSettings() # pyright: ignore[reportCallIssue] # pydantic-settings fills fields from env diff --git a/litellm/types/proxy/control_plane/__init__.py b/litellm/types/proxy/control_plane/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/types/proxy/control_plane/auth.py b/litellm/types/proxy/control_plane/auth.py new file mode 100644 index 00000000000..39b942d0f67 --- /dev/null +++ b/litellm/types/proxy/control_plane/auth.py @@ -0,0 +1,34 @@ +from typing import Final + +from pydantic import ConfigDict, Field + +from litellm.types.llms.base import LiteLLMBaseModel + +USERNAME_MAX_LENGTH: Final = 256 +PASSWORD_MAX_LENGTH: Final = 1024 + + +class LoginRequest(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + username: str = Field(min_length=1, max_length=USERNAME_MAX_LENGTH) + password: str = Field(min_length=1, max_length=PASSWORD_MAX_LENGTH) + + +class SessionTokenResponse(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + access_token: str + token_type: str = "bearer" + expires_in: int + + +class PrincipalResponse(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + subject: str + email: str | None + roles: tuple[str, ...] + scopes: tuple[str, ...] + dataplanes: tuple[str, ...] | None + token_kind: str diff --git a/litellm/types/proxy/control_plane/health.py b/litellm/types/proxy/control_plane/health.py new file mode 100644 index 00000000000..792fc8abc9a --- /dev/null +++ b/litellm/types/proxy/control_plane/health.py @@ -0,0 +1,11 @@ +from typing import Literal + +from pydantic import ConfigDict + +from litellm.types.llms.base import LiteLLMBaseModel + + +class HealthResponse(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + status: Literal["ok"] diff --git a/tests/unit/proxy/control_plane/__init__.py b/tests/unit/proxy/control_plane/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/proxy/control_plane/auth/__init__.py b/tests/unit/proxy/control_plane/auth/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/proxy/control_plane/auth/test_dependencies.py b/tests/unit/proxy/control_plane/auth/test_dependencies.py new file mode 100644 index 00000000000..ced7509d310 --- /dev/null +++ b/tests/unit/proxy/control_plane/auth/test_dependencies.py @@ -0,0 +1,92 @@ +import asyncio +import threading +from datetime import timedelta +from typing import Annotated, Final + +import httpx +from fastapi import Depends, FastAPI, Security +from fastapi.testclient import TestClient + +from litellm.proxy.control_plane.auth.dependencies import get_principal, get_token_verifier, require_scopes +from litellm.proxy.control_plane.auth.jwt import Rejected, SessionTokenCodec, TokenVerifier +from litellm.proxy.control_plane.auth.rbac import ADMIN_ROLE, VIEWER_ROLE, Principal, build_principal, default_policy +from litellm.proxy.control_plane.common.exceptions import install_exception_handlers + +MASTER_KEY: Final = "sk-master" +POLICY: Final = default_policy() + + +def _app() -> tuple[TestClient, SessionTokenCodec]: + codec: Final = SessionTokenCodec(MASTER_KEY, "cp", timedelta(hours=1), POLICY) + app: Final = FastAPI() + verifier: Final = TokenVerifier(MASTER_KEY, codec, None, POLICY) + app.dependency_overrides[get_token_verifier] = lambda: verifier + install_exception_handlers(app) + + @app.get("/usage") + async def usage(principal: Annotated[Principal, Security(require_scopes, scopes=["usage:read"])]) -> str: + return principal.subject + + @app.post("/dataplanes") + async def write(principal: Annotated[Principal, Security(require_scopes, scopes=["dataplanes:write"])]) -> str: + return principal.subject + + return TestClient(app), codec + + +def _token(codec: SessionTokenCodec, role: str) -> dict[str, str]: + principal: Final = build_principal(POLICY, "u", None, frozenset({role}), None, frozenset(), "session") + return {"Authorization": f"Bearer {codec.mint(principal)}"} + + +def test_viewer_reads_usage_but_cannot_write_dataplanes() -> None: + client, codec = _app() + assert client.get("/usage", headers=_token(codec, VIEWER_ROLE)).status_code == 200 + denied: Final = client.post("/dataplanes", headers=_token(codec, VIEWER_ROLE)) + assert denied.status_code == 403 + assert denied.json()["detail"] == "Missing scope: dataplanes:write." + assert denied.headers["content-type"] == "application/problem+json" + + +def test_admin_passes_both_scope_gates() -> None: + client, codec = _app() + assert client.post("/dataplanes", headers=_token(codec, ADMIN_ROLE)).status_code == 200 + + +class _GatedVerifier: + """Blocks the slow credential on an Event and records whether that Event was set before the wait gave up.""" + + def __init__(self, gate: threading.Event) -> None: + self._gate = gate + self.slow_saw_the_gate_open: bool | None = None + + def verify(self, credential: str) -> Principal | Rejected: + if credential == "slow": + self.slow_saw_the_gate_open = self._gate.wait(timeout=5) + return build_principal(POLICY, credential, None, frozenset({VIEWER_ROLE}), None, frozenset(), "session") + + +async def test_a_slow_token_verification_does_not_block_other_requests() -> None: + gate: Final = threading.Event() + verifier: Final = _GatedVerifier(gate) + app: Final = FastAPI() + app.dependency_overrides[get_token_verifier] = lambda: verifier + + async def me(principal: Annotated[Principal, Depends(get_principal)]) -> str: + return principal.subject + + app.add_api_route("/me", me) + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://cp") as client: + slow: Final = asyncio.create_task(client.get("/me", headers={"Authorization": "Bearer slow"})) + await asyncio.sleep(0) + fast: Final = await client.get("/me", headers={"Authorization": "Bearer fast"}) + gate.set() + assert (await slow).json() == "slow" + assert fast.json() == "fast" + assert verifier.slow_saw_the_gate_open is True, "the fast request ran while the slow JWKS lookup was still waiting" + + +def test_security_scopes_show_up_in_openapi() -> None: + client, _ = _app() + security: Final = client.get("/openapi.json").json()["paths"]["/usage"]["get"]["security"] + assert security == [{"ControlPlaneBearer": ["usage:read"]}] diff --git a/tests/unit/proxy/control_plane/auth/test_jwt.py b/tests/unit/proxy/control_plane/auth/test_jwt.py new file mode 100644 index 00000000000..ebe33caca25 --- /dev/null +++ b/tests/unit/proxy/control_plane/auth/test_jwt.py @@ -0,0 +1,437 @@ +import threading +from collections.abc import Mapping +from concurrent.futures import ThreadPoolExecutor +from datetime import timedelta +from typing import Final + +import jwt +import pytest +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric import rsa + +from litellm.proxy.control_plane.auth.jwt import ( + SESSION_AUDIENCE, + OidcTokenVerifier, + PyJwkClientLookup, + Rejected, + SessionTokenCodec, + TokenVerifier, + master_key_principal, +) +from litellm.proxy.control_plane.auth.rbac import ( + ADMIN_ROLE, + ALL_SCOPES, + VIEWER_ROLE, + Principal, + RoleDefinition, + RolePolicy, + Scope, + build_policy, + build_principal, + default_policy, +) +from litellm.proxy.control_plane.settings import OidcSettings + +MASTER_KEY: Final = "sk-master-test" +ISSUER: Final = "https://idp.example.test/" +AUDIENCE: Final = "litellm-control-plane" +TOKEN_IAT: Final = 1_700_000_000 +TOKEN_EXP: Final = 4_102_444_800 +EXPIRED_TOKEN_EXP: Final = 1_600_000_000 + +_PRIVATE_KEY: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) +_PUBLIC_PEM: Final = ( + _PRIVATE_KEY.public_key() + .public_bytes(serialization.Encoding.PEM, serialization.PublicFormat.SubjectPublicKeyInfo) + .decode() +) +_PRIVATE_PEM: Final = _PRIVATE_KEY.private_bytes( + serialization.Encoding.PEM, serialization.PrivateFormat.PKCS8, serialization.NoEncryption() +).decode() + + +class StaticKeys: + def __init__(self, key: str | Rejected) -> None: + self._key = key + + def signing_key(self, token: str) -> str | Rejected: # str is a valid SigningKey + return self._key + + +class StaticJwksClient: + def get_signing_key_from_jwt(self, _token: str) -> str: + return _PUBLIC_PEM + + +class _Clock: + def __init__(self) -> None: + self.current = 100.0 + + def __call__(self) -> float: + return self.current + + def advance(self, seconds: float) -> None: + self.current += seconds + + +class _FetchCounter: + def __init__(self) -> None: + self._lock = threading.Lock() + self._calls = 0 + + @property + def count(self) -> int: + with self._lock: + return self._calls + + def record(self) -> int: + with self._lock: + self._calls += 1 + return self._calls + + +def _oidc_settings(**overrides: str) -> OidcSettings: + return OidcSettings( + issuer=ISSUER, + audience=AUDIENCE, + jwks_url=None, + roles_claim=overrides.get("roles_claim", "roles"), + scope_claim="scope", + dataplanes_claim="litellm_dataplanes", + email_claim="email", + leeway_seconds=0, + ) + + +def _oidc_token(claims: dict[str, object]) -> str: + return jwt.encode( + {"iss": ISSUER, "aud": AUDIENCE, "sub": "user-1", "exp": TOKEN_EXP, "iat": TOKEN_IAT, **claims}, + _PRIVATE_PEM, + algorithm="RS256", + ) + + +POLICY: Final = default_policy() + + +def _sessions( + ttl: timedelta = timedelta(hours=1), + master_key: str = MASTER_KEY, + issuer: str = "cp", + audience: str = SESSION_AUDIENCE, +) -> SessionTokenCodec: + return SessionTokenCodec(master_key=master_key, issuer=issuer, ttl=ttl, policy=POLICY, audience=audience) + + +def _viewer(ids: frozenset[str]) -> Principal: + return build_principal(POLICY, "alice", "a@x.test", frozenset({VIEWER_ROLE}), None, ids, "session") + + +def _oidc_verifier(settings: OidcSettings, key: str | Rejected) -> OidcTokenVerifier: + return OidcTokenVerifier(settings, StaticKeys(key), POLICY) + + +def test_session_round_trip_preserves_subject_roles_scopes_and_binding() -> None: + codec: Final = _sessions() + original: Final = _viewer(frozenset({"eu", "us"})) + verified: Final = codec.verify(codec.mint(original)) + assert isinstance(verified, Principal) + assert (verified.subject, verified.email, verified.roles) == ("alice", "a@x.test", original.roles) + assert verified.scopes == original.scopes + assert verified.dataplanes == original.dataplanes + assert verified.token_kind == "session" + + +def test_fleet_wide_session_round_trips_as_fleet_wide() -> None: + codec: Final = _sessions() + admin: Final = build_principal(POLICY, "root", None, frozenset({ADMIN_ROLE}), None, frozenset(), "session") + verified: Final = codec.verify(codec.mint(admin)) + assert isinstance(verified, Principal) and verified.dataplanes.all + + +def test_expired_session_is_rejected() -> None: + codec: Final = _sessions(ttl=timedelta(seconds=-1)) + assert isinstance(codec.verify(codec.mint(_viewer(frozenset()))), Rejected) + + +def test_session_signed_with_another_key_is_rejected() -> None: + other: Final = _sessions(master_key="sk-other") + assert isinstance(_sessions().verify(other.mint(_viewer(frozenset()))), Rejected) + + +def test_session_from_another_issuer_is_rejected() -> None: + other: Final = _sessions(issuer="not-cp") + assert isinstance(_sessions().verify(other.mint(_viewer(frozenset()))), Rejected) + + +def test_session_minted_for_another_audience_is_rejected() -> None: + """Same master key and issuer, but a token for another component never opens the control plane.""" + other: Final = _sessions(audience="litellm-proxy-ui") + assert isinstance(_sessions().verify(other.mint(_viewer(frozenset()))), Rejected) + + +def test_session_without_an_audience_claim_is_rejected() -> None: + no_aud: Final = jwt.encode( + {"iss": "cp", "sub": "alice", "roles": ["viewer"], "iat": TOKEN_IAT, "exp": TOKEN_EXP}, + MASTER_KEY, + algorithm="HS256", + ) + assert isinstance(_sessions().verify(no_aud), Rejected) + with_aud: Final = jwt.encode( + { + "iss": "cp", + "aud": SESSION_AUDIENCE, + "sub": "alice", + "roles": ["viewer"], + "iat": TOKEN_IAT, + "exp": TOKEN_EXP, + }, + MASTER_KEY, + algorithm="HS256", + ) + assert isinstance(_sessions().verify(with_aud), Principal) + + +def test_oidc_scope_claim_with_only_foreign_scopes_grants_no_control_plane_scope() -> None: + verifier: Final = _oidc_verifier(_oidc_settings(), _PUBLIC_PEM) + scoped: Final = verifier.verify(_oidc_token({"roles": ["admin"], "scope": "openid profile email"})) + assert isinstance(scoped, Principal) and scoped.scopes == frozenset() + unscoped: Final = verifier.verify(_oidc_token({"roles": ["admin"]})) + assert isinstance(unscoped, Principal) and unscoped.scopes == POLICY.scopes_for(ADMIN_ROLE) + + +def test_oidc_token_maps_claims_to_a_clamped_principal() -> None: + verifier: Final = _oidc_verifier(_oidc_settings(), _PUBLIC_PEM) + token: Final = _oidc_token( + { + "roles": ["viewer"], + "scope": "openid usage:read proxy:write", + "litellm_dataplanes": ["eu"], + "email": "u@x.test", + } + ) + principal: Final = verifier.verify(token) + assert isinstance(principal, Principal) + assert principal.roles == frozenset({VIEWER_ROLE}) + assert principal.scopes == frozenset({Scope.USAGE_READ}) + assert principal.dataplanes.ids == frozenset({"eu"}) and not principal.dataplanes.all + assert (principal.email, principal.token_kind) == ("u@x.test", "oidc") + + +@pytest.mark.parametrize("dataplanes", [1, ["eu", 1]]) +def test_oidc_token_with_malformed_dataplanes_claim_is_rejected(dataplanes: object) -> None: + verifier: Final = _oidc_verifier(_oidc_settings(), _PUBLIC_PEM) + assert isinstance(verifier.verify(_oidc_token({"roles": ["admin"], "litellm_dataplanes": dataplanes})), Rejected) + + +def test_oidc_token_without_dataplanes_claim_keeps_fleet_wide_access() -> None: + verifier: Final = _oidc_verifier(_oidc_settings(), _PUBLIC_PEM) + principal: Final = verifier.verify(_oidc_token({"roles": ["admin"]})) + assert isinstance(principal, Principal) and principal.dataplanes.all + + +def test_oidc_token_with_a_fixed_past_expiration_is_rejected() -> None: + verifier: Final = _oidc_verifier(_oidc_settings(), _PUBLIC_PEM) + assert isinstance(verifier.verify(_oidc_token({"roles": ["admin"], "exp": EXPIRED_TOKEN_EXP})), Rejected) + + +def test_oidc_roles_claim_name_is_configurable_and_space_delimited_roles_parse() -> None: + verifier: Final = _oidc_verifier(_oidc_settings(roles_claim="groups"), _PUBLIC_PEM) + principal: Final = verifier.verify(_oidc_token({"groups": "viewer other"})) + assert isinstance(principal, Principal) and principal.roles == frozenset({VIEWER_ROLE}) + + +@pytest.mark.parametrize( + "claims", + [ + {"roles": ["admin"], "aud": "someone-else"}, + {"roles": ["admin"], "iss": "https://other.test/"}, + {"roles": ["proxy_admin"]}, + {}, + ], +) +def test_oidc_token_with_wrong_audience_issuer_or_no_role_is_rejected(claims: dict[str, object]) -> None: + verifier: Final = _oidc_verifier(_oidc_settings(), _PUBLIC_PEM) + assert isinstance(verifier.verify(_oidc_token(claims)), Rejected) + + +def test_oidc_token_signed_by_an_unknown_key_is_rejected() -> None: + other_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + other_pem: Final = ( + other_key.public_key() + .public_bytes(serialization.Encoding.PEM, serialization.PublicFormat.SubjectPublicKeyInfo) + .decode() + ) + verifier: Final = _oidc_verifier(_oidc_settings(), other_pem) + assert isinstance(verifier.verify(_oidc_token({"roles": ["admin"]})), Rejected) + + +def test_oidc_token_is_rejected_when_the_jwks_lookup_fails() -> None: + verifier: Final = _oidc_verifier(_oidc_settings(), Rejected("jwks down")) + assert verifier.verify(_oidc_token({"roles": ["admin"]})) == Rejected("jwks down") + + +def test_oidc_discovery_is_lazy_cached_and_uses_the_discovered_jwks_uri() -> None: + configurations: Final = iter(({"jwks_uri": "https://idp.example.test/keys"},)) + clients: Final = iter((StaticJwksClient(),)) + + def fetch(url: str) -> Mapping[str, object]: + assert url == "https://idp.example.test/.well-known/openid-configuration" + return next(configurations) + + def client_factory(url: str) -> StaticJwksClient: + assert url == "https://idp.example.test/keys" + return next(clients) + + lookup: Final = PyJwkClientLookup(None, ISSUER, fetch=fetch, client_factory=client_factory) + assert lookup.signing_key("token") == _PUBLIC_PEM + assert lookup.signing_key("token") == _PUBLIC_PEM + + +def test_oidc_discovery_failures_are_retried_on_the_next_lookup() -> None: + clock: Final = _Clock() + configurations: Final = iter((OSError("discovery unavailable"), {"jwks_uri": "https://idp.example.test/keys"})) + + def fetch(_url: str) -> Mapping[str, object]: + outcome: Final = next(configurations) + if isinstance(outcome, Exception): + raise outcome + return outcome + + lookup: Final = PyJwkClientLookup( + None, + ISSUER, + fetch=fetch, + client_factory=lambda _url: StaticJwksClient(), + clock=clock, + ) + assert isinstance(lookup.signing_key("token"), Rejected) + clock.advance(30) + assert lookup.signing_key("token") == _PUBLIC_PEM + + +def test_concurrent_oidc_discovery_failures_are_single_flight_and_cool_down() -> None: + clock: Final = _Clock() + fetch_calls: Final = _FetchCounter() + fetch_started: Final = threading.Event() + release_fetch: Final = threading.Event() + + def fetch(_url: str) -> Mapping[str, object]: + call_number: Final = fetch_calls.record() + fetch_started.set() + if call_number == 1: + assert release_fetch.wait(timeout=5) + raise OSError("discovery unavailable") + return {"jwks_uri": "https://idp.example.test/keys"} + + lookup: Final = PyJwkClientLookup( + None, + ISSUER, + fetch=fetch, + client_factory=lambda _url: StaticJwksClient(), + clock=clock, + ) + oidc: Final = OidcTokenVerifier(_oidc_settings(), lookup, POLICY) + verifier: Final = TokenVerifier(MASTER_KEY, _sessions(), oidc, POLICY) + forged_token: Final = jwt.encode( + { + "iss": ISSUER, + "aud": AUDIENCE, + "sub": "forged", + "roles": ["admin"], + "iat": TOKEN_IAT, + "exp": TOKEN_EXP, + }, + "not-the-idp-key", + algorithm="HS256", + ) + with ThreadPoolExecutor(max_workers=8) as executor: + futures: Final = tuple(executor.submit(verifier.verify, forged_token) for _ in range(8)) + assert fetch_started.wait(timeout=5) + release_fetch.set() + assert all(isinstance(future.result(timeout=5), Rejected) for future in futures) + assert fetch_calls.count == 1 + assert isinstance(verifier.verify(forged_token), Rejected) + assert fetch_calls.count == 1 + clock.advance(30) + assert isinstance(verifier.verify(forged_token), Rejected) + assert fetch_calls.count == 2 + + +def test_explicit_jwks_url_takes_precedence_over_oidc_discovery() -> None: + lookup: Final = PyJwkClientLookup( + "https://idp.example.test/custom-keys", + ISSUER, + fetch=lambda _url: pytest.fail("discovery should not run when jwks_url is configured"), + client_factory=lambda url: ( + StaticJwksClient() if url == "https://idp.example.test/custom-keys" else pytest.fail() + ), + ) + assert lookup.signing_key("token") == _PUBLIC_PEM + + +@pytest.mark.parametrize( + "error", + [OSError("network unavailable"), TimeoutError("request timed out"), ValueError("malformed JWKS"), jwt.PyJWTError()], +) +def test_jwks_client_errors_are_rejected(error: Exception) -> None: + class FailingJwksClient: + def get_signing_key_from_jwt(self, _token: str) -> str: + raise error + + lookup: Final = PyJwkClientLookup( + "https://idp.example.test/keys", + ISSUER, + client_factory=lambda _url: FailingJwksClient(), + ) + assert isinstance(lookup.signing_key("token"), Rejected) + + +def test_hs256_token_claiming_the_oidc_issuer_cannot_pass_as_oidc() -> None: + verifier: Final = _oidc_verifier(_oidc_settings(), _PUBLIC_PEM) + forged: Final = jwt.encode( + { + "iss": ISSUER, + "aud": AUDIENCE, + "sub": "x", + "roles": ["admin"], + "iat": TOKEN_IAT, + "exp": TOKEN_EXP, + }, + "a" * 32, + algorithm="HS256", + ) + assert isinstance(verifier.verify(forged), Rejected) + + +def test_token_verifier_routes_master_key_session_and_oidc_credentials() -> None: + sessions: Final = _sessions() + verifier: Final = TokenVerifier(MASTER_KEY, sessions, _oidc_verifier(_oidc_settings(), _PUBLIC_PEM), POLICY) + assert verifier.verify(MASTER_KEY) == master_key_principal() + session: Final = verifier.verify(sessions.mint(_viewer(frozenset({"eu"})))) + assert isinstance(session, Principal) and session.token_kind == "session" + oidc: Final = verifier.verify(_oidc_token({"roles": ["viewer"]})) + assert isinstance(oidc, Principal) and oidc.token_kind == "oidc" + assert isinstance(verifier.verify("not-a-token"), Rejected) + assert isinstance(verifier.verify("sk-master-tes"), Rejected) + assert isinstance(verifier.verify("café"), Rejected) + + +def test_oidc_issued_token_is_not_accepted_as_a_session_when_oidc_is_off() -> None: + verifier: Final = TokenVerifier(MASTER_KEY, _sessions(), None, POLICY) + assert isinstance(verifier.verify(_oidc_token({"roles": ["admin"]})), Rejected) + + +def test_master_key_principal_is_a_fleet_wide_admin() -> None: + principal: Final = master_key_principal() + assert principal.roles == frozenset({ADMIN_ROLE}) and principal.dataplanes.all + assert principal.token_kind == "master_key" + assert principal.scopes == ALL_SCOPES + + +def test_master_key_keeps_every_scope_when_the_default_roles_are_switched_off() -> None: + custom_only: Final = build_policy({"finops": RoleDefinition(frozenset({Scope.USAGE_READ}), frozenset())}, False) + assert isinstance(custom_only, RolePolicy) and custom_only.scopes_for(ADMIN_ROLE) == frozenset() + verifier: Final = TokenVerifier(MASTER_KEY, _sessions(), None, custom_only) + verified: Final = verifier.verify(MASTER_KEY) + assert isinstance(verified, Principal) and verified.scopes == ALL_SCOPES diff --git a/tests/unit/proxy/control_plane/auth/test_login.py b/tests/unit/proxy/control_plane/auth/test_login.py new file mode 100644 index 00000000000..6543f8b6399 --- /dev/null +++ b/tests/unit/proxy/control_plane/auth/test_login.py @@ -0,0 +1,134 @@ +from datetime import timedelta +from typing import Final + +import pytest + +from litellm.proxy.control_plane.auth.jwt import SessionTokenCodec +from litellm.proxy.control_plane.auth.login import ( + LOCKOUT_FAILURES, + LOCKOUT_SECONDS, + InvalidCredentials, + LoginService, + LoginThrottle, + MAX_TRACKED_LOGIN_KEYS, + Throttled, +) +from litellm.proxy.control_plane.auth.rbac import default_policy +from litellm.types.proxy.control_plane.auth import SessionTokenResponse + + +class Clock: # mutable-ok: the test advances it + def __init__(self) -> None: + self.now = 1000.0 + + def __call__(self) -> float: + return self.now + + +def _service(clock: Clock, throttle: LoginThrottle | None = None) -> LoginService: + policy: Final = default_policy() + sessions: Final = SessionTokenCodec( + master_key="sk-master", issuer="cp", audience="cp-api", ttl=timedelta(minutes=5), policy=policy + ) + return LoginService( + sessions=sessions, + policy=policy, + ttl_seconds=300, + admin_username="admin", + admin_password="hunter2", + throttle=throttle if throttle is not None else LoginThrottle(clock), + ) + + +def test_account_locks_after_repeated_misses_and_unlocks_after_the_window() -> None: + clock: Final = Clock() + service: Final = _service(clock) + misses: Final = [service.login("admin", "wrong") for i in range(LOCKOUT_FAILURES)] + assert all(isinstance(miss, InvalidCredentials) for miss in misses) + locked: Final = service.login("admin", "hunter2") + assert locked == Throttled(retry_after_seconds=int(LOCKOUT_SECONDS)) + clock.now += LOCKOUT_SECONDS / 2 + assert service.login("admin", "hunter2") == Throttled(retry_after_seconds=int(LOCKOUT_SECONDS / 2)) + clock.now += LOCKOUT_SECONDS / 2 + 1 + assert isinstance(service.login("admin", "hunter2"), SessionTokenResponse) + + +def test_login_throttle_is_scoped_to_the_client_source() -> None: + clock: Final = Clock() + service: Final = _service(clock) + misses: Final = tuple(service.login("admin", "wrong", client_source="client-a") for _ in range(LOCKOUT_FAILURES)) + assert all(isinstance(miss, InvalidCredentials) for miss in misses) + assert isinstance(service.login("admin", "hunter2", client_source="client-a"), Throttled) + assert isinstance(service.login("admin", "hunter2", client_source="client-b"), SessionTokenResponse) + + +def test_unknown_usernames_do_not_grow_the_throttle_or_lock_out_the_admin() -> None: + clock: Final = Clock() + throttle: Final = LoginThrottle(clock) + service: Final = _service(clock, throttle) + for index in range(100): + assert isinstance(service.login(f"unknown-{index}", "wrong"), InvalidCredentials) + assert throttle.tracked() == 0 + assert isinstance(service.login("admin", "hunter2"), SessionTokenResponse) + + +def test_failure_recording_prunes_keys_whose_misses_expired() -> None: + clock: Final = Clock() + throttle: Final = LoginThrottle(clock) + throttle.record_failure(("old-source", "admin")) + clock.now += LOCKOUT_SECONDS + 1 + + throttle.record_failure(("new-source", "admin")) + + assert throttle.tracked() == 1 + + +def test_rotating_sources_stay_bounded_without_evicting_a_hot_lockout() -> None: + clock: Final = Clock() + throttle: Final = LoginThrottle(clock) + hot_key: Final = ("hot-source", "admin") + for _ in range(LOCKOUT_FAILURES - 1): + throttle.record_failure(hot_key) + for index in range(MAX_TRACKED_LOGIN_KEYS - 1): + clock.now += 0.001 + throttle.record_failure((f"source-{index}", "admin")) + + clock.now += 0.001 + throttle.record_failure(hot_key) + clock.now += 0.001 + throttle.record_failure(("source-over-cap", "admin")) + + assert throttle.tracked() == MAX_TRACKED_LOGIN_KEYS + assert throttle.locked_for(hot_key) > 0 + + +def test_a_success_clears_the_accounts_misses() -> None: + clock: Final = Clock() + service: Final = _service(clock) + for _ in range(LOCKOUT_FAILURES - 1): + service.login("admin", "wrong") + assert isinstance(service.login("admin", "hunter2"), SessionTokenResponse) + for _ in range(LOCKOUT_FAILURES - 1): + assert isinstance(service.login("admin", "wrong"), InvalidCredentials) + assert isinstance(service.login("admin", "hunter2"), SessionTokenResponse) + + +@pytest.mark.parametrize("username", ["admin", "nobody"]) +def test_no_configured_admin_rejects_everyone(username: str) -> None: + throttle: Final = LoginThrottle(Clock()) + sessions: Final = SessionTokenCodec( + master_key="sk-master", issuer="cp", audience="cp-api", ttl=timedelta(minutes=5), policy=default_policy() + ) + service: Final = LoginService(sessions, default_policy(), 300, None, None, throttle) + assert isinstance(service.login(username, "hunter2"), InvalidCredentials) + assert throttle.tracked() == 0 + + +@pytest.mark.parametrize("username", ["", "u" * 257]) +def test_login_request_bounds_the_username(username: str) -> None: + from pydantic import ValidationError + + from litellm.types.proxy.control_plane.auth import LoginRequest + + with pytest.raises(ValidationError): + LoginRequest(username=username, password="x") diff --git a/tests/unit/proxy/control_plane/auth/test_rbac.py b/tests/unit/proxy/control_plane/auth/test_rbac.py new file mode 100644 index 00000000000..eda37491d70 --- /dev/null +++ b/tests/unit/proxy/control_plane/auth/test_rbac.py @@ -0,0 +1,135 @@ +from typing import Final + +import pytest + +from litellm.proxy.control_plane.auth.rbac import ( + ADMIN_ROLE, + DEFAULT_ROLES, + EDITOR_ROLE, + USAGE_VIEWER_ROLE, + VIEWER_ROLE, + DataplaneBinding, + RoleDefinition, + RolePolicy, + Scope, + UnknownParent, + UnknownScope, + build_policy, + build_principal, + clamp_scopes, + default_policy, + parse_role_definition, + parse_scopes, +) + + +def _principal(policy: RolePolicy, roles: frozenset[str], ids: frozenset[str] = frozenset()): + return build_principal(policy, "u", None, roles, None, ids, "oidc") + + +@pytest.mark.parametrize("role", list(DEFAULT_ROLES)) +@pytest.mark.parametrize("scope", list(Scope)) +def test_every_default_role_scope_cell_matches_the_policy_table(role: str, scope: Scope) -> None: + assert _principal(default_policy(), frozenset({role})).has_scope(scope) is (scope in DEFAULT_ROLES[role]) + + +def test_default_tiers_only_admin_manages_access_only_editor_writes_and_usage_viewer_sees_only_usage() -> None: + assert DEFAULT_ROLES[ADMIN_ROLE] == frozenset(Scope) + assert {Scope.DATAPLANES_WRITE, Scope.ROLES_WRITE, Scope.AUDIT_READ}.isdisjoint(DEFAULT_ROLES[EDITOR_ROLE]) + assert Scope.PROXY_WRITE in DEFAULT_ROLES[EDITOR_ROLE] and Scope.PROXY_WRITE not in DEFAULT_ROLES[VIEWER_ROLE] + assert DEFAULT_ROLES[USAGE_VIEWER_ROLE] == frozenset({Scope.USAGE_READ}) + + +def test_roles_union_their_grants() -> None: + principal: Final = _principal(default_policy(), frozenset({USAGE_VIEWER_ROLE, VIEWER_ROLE})) + assert principal.scopes == DEFAULT_ROLES[VIEWER_ROLE] + + +def test_claimed_scopes_narrow_but_never_widen_the_role_grant() -> None: + granted: Final = DEFAULT_ROLES[VIEWER_ROLE] + assert clamp_scopes(granted, frozenset({Scope.USAGE_READ, Scope.PROXY_WRITE})) == frozenset({Scope.USAGE_READ}) + assert clamp_scopes(granted, None) == granted + + +def test_a_scope_claim_naming_only_unknown_scopes_grants_nothing() -> None: + """A present claim is always honored, so a token scoped to foreign scopes cannot fall back to the role grant.""" + assert clamp_scopes(DEFAULT_ROLES[ADMIN_ROLE], frozenset()) == frozenset() + assert clamp_scopes(DEFAULT_ROLES[ADMIN_ROLE], parse_scopes(("openid", "profile"))) == frozenset() + + +def test_custom_role_overrides_a_default_and_can_inherit() -> None: + policy: Final = build_policy( + { + "viewer": RoleDefinition(scopes=frozenset({Scope.FLEET_READ}), inherits=frozenset()), + "finops": RoleDefinition(scopes=frozenset({Scope.AUDIT_READ}), inherits=frozenset({"usage_viewer"})), + "lead": RoleDefinition(scopes=frozenset(), inherits=frozenset({"finops", "editor"})), + } + ) + assert isinstance(policy, RolePolicy) + assert policy.scopes_for("viewer") == frozenset({Scope.FLEET_READ}) + assert policy.scopes_for("finops") == frozenset({Scope.AUDIT_READ, Scope.USAGE_READ}) + assert policy.scopes_for("lead") == DEFAULT_ROLES[EDITOR_ROLE] | {Scope.AUDIT_READ, Scope.USAGE_READ} + assert policy.scopes_for("admin") == frozenset(Scope) + + +def test_defaults_can_be_switched_off_so_only_custom_roles_exist() -> None: + policy: Final = build_policy( + {"ops": RoleDefinition(frozenset({Scope.USAGE_READ}), frozenset())}, include_defaults=False + ) + assert isinstance(policy, RolePolicy) + assert policy.names == frozenset({"ops"}) + assert policy.known(("admin", "ops")) == frozenset({"ops"}) + + +def test_inheriting_an_unknown_role_is_an_error_value() -> None: + policy: Final = build_policy({"x": RoleDefinition(frozenset(), frozenset({"nope"}))}) + assert policy == UnknownParent(role="x", parent="nope") + + +def test_inheritance_cycles_terminate() -> None: + policy: Final = build_policy( + { + "a": RoleDefinition(frozenset({Scope.USAGE_READ}), frozenset({"b"})), + "b": RoleDefinition(frozenset({Scope.FLEET_READ}), frozenset({"a"})), + } + ) + assert isinstance(policy, RolePolicy) + assert policy.scopes_for("a") == policy.scopes_for("b") == frozenset({Scope.USAGE_READ, Scope.FLEET_READ}) + + +def test_role_definition_rejects_unknown_scopes() -> None: + assert parse_role_definition("r", ("usage:read", "nope"), ()) == UnknownScope(role="r", scope="nope") + parsed: Final = parse_role_definition("r", ("usage:read",), ("viewer",)) + assert parsed == RoleDefinition(scopes=frozenset({Scope.USAGE_READ}), inherits=frozenset({"viewer"})) + + +@pytest.mark.parametrize("ids", [frozenset(), frozenset({"*"}), frozenset({"*", "eu"})]) +def test_no_binding_or_star_is_fleet_wide(ids: frozenset[str]) -> None: + principal: Final = _principal(default_policy(), frozenset({VIEWER_ROLE}), ids) + assert principal.dataplanes == DataplaneBinding(all=True, ids=frozenset()) + assert principal.dataplanes.allows("anything") + + +def test_explicit_binding_limits_dataplanes_regardless_of_role() -> None: + principal: Final = _principal(default_policy(), frozenset({ADMIN_ROLE}), frozenset({"eu"})) + assert principal.dataplanes.allows("eu") and not principal.dataplanes.allows("us") + + +@pytest.mark.parametrize( + ("binding", "requested", "expected"), + [ + (DataplaneBinding(True, frozenset()), None, None), + (DataplaneBinding(True, frozenset()), frozenset({"a"}), frozenset({"a"})), + (DataplaneBinding(False, frozenset({"a", "b"})), None, frozenset({"a", "b"})), + (DataplaneBinding(False, frozenset({"a", "b"})), frozenset({"b", "c"}), frozenset({"b"})), + ], +) +def test_restrict_intersects_the_request_with_the_binding( + binding: DataplaneBinding, requested: frozenset[str] | None, expected: frozenset[str] | None +) -> None: + assert binding.restrict(requested) == expected + + +def test_unknown_scopes_are_dropped_not_raised() -> None: + assert parse_scopes(("usage:read", "openid", "profile")) == frozenset({Scope.USAGE_READ}) + assert parse_scopes(None) is None diff --git a/tests/unit/proxy/control_plane/common/__init__.py b/tests/unit/proxy/control_plane/common/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/proxy/control_plane/health/__init__.py b/tests/unit/proxy/control_plane/health/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/proxy/control_plane/test_app.py b/tests/unit/proxy/control_plane/test_app.py new file mode 100644 index 00000000000..89eef97d441 --- /dev/null +++ b/tests/unit/proxy/control_plane/test_app.py @@ -0,0 +1,259 @@ +from datetime import timedelta +from itertools import repeat +from typing import Final + +import jwt +import pytest +from fastapi.testclient import TestClient +from pydantic import SecretStr + +from litellm.proxy.control_plane.app import LicenseRequired, create_app +from litellm.proxy.control_plane.auth.dependencies import get_token_verifier +from litellm.proxy.control_plane.auth.jwt import PyJwkClientLookup, SessionTokenCodec, build_token_verifier +from litellm.proxy.control_plane.auth.login import LOCKOUT_FAILURES, LoginService, LoginThrottle +from litellm.proxy.control_plane.auth.rbac import VIEWER_ROLE, build_principal +from litellm.proxy.control_plane.auth.router import get_login_service +from litellm.proxy.control_plane.container import InvalidRolePolicy, build_container +from litellm.proxy.control_plane.settings import ControlPlaneSettings + +MASTER_KEY: Final = "01234567890123456789012345678901" + + +class Licensed: + def is_premium(self) -> bool: + return True + + +class Unlicensed: + def is_premium(self) -> bool: + return False + + +def _settings(**overrides: object) -> ControlPlaneSettings: + return ControlPlaneSettings.model_validate( + {"LITELLM_MASTER_KEY": MASTER_KEY, "UI_USERNAME": "admin", "UI_PASSWORD": "pw", **overrides} + ) + + +def _client(settings: ControlPlaneSettings | None = None) -> TestClient: + return TestClient(create_app(build_container(settings or _settings()), Licensed())) + + +def _bearer(token: str) -> dict[str, str]: + return {"Authorization": f"Bearer {token}"} + + +def test_app_refuses_to_build_without_a_valid_license() -> None: + with pytest.raises(LicenseRequired): + create_app(build_container(_settings()), Unlicensed()) + + +def test_health_routes_need_no_credentials() -> None: + client: Final = _client() + assert client.get("/healthz").json() == {"status": "ok"} + assert client.get("/readyz").json() == {"status": "ok"} + + +def test_unknown_path_returns_an_rfc9457_not_found_body() -> None: + response: Final = _client().get("/api/v1/not-a-route") + assert response.status_code == 404 + assert response.headers["content-type"] == "application/problem+json" + assert response.json() == { + "type": "urn:litellm:error:not-found", + "title": "Not Found", + "status": 404, + "detail": "Not Found", + } + + +def test_wrong_method_returns_an_rfc9457_error_and_preserves_allow() -> None: + response: Final = _client().post("/api/v1/auth/me") + assert response.status_code == 405 + assert response.headers["content-type"] == "application/problem+json" + assert "GET" in response.headers["allow"] + assert response.json() == { + "type": "urn:litellm:error:method-not-allowed", + "title": "Method Not Allowed", + "status": 405, + "detail": "Method Not Allowed", + } + + +def test_jwks_lookup_errors_make_auth_me_return_401_instead_of_500() -> None: + settings: Final = _settings(oidc_issuer="https://idp.example.test", oidc_audience="control-plane") + container: Final = build_container(settings) + + class FailingJwksClient: + def get_signing_key_from_jwt(self, _token: str) -> str: + raise OSError("JWKS server is unavailable") + + lookup: Final = PyJwkClientLookup( + None, + "https://idp.example.test", + fetch=lambda _url: {"jwks_uri": "https://idp.example.test/keys"}, + client_factory=lambda _url: FailingJwksClient(), + ) + verifier: Final = build_token_verifier(settings, container.policy, keys=lookup) + app: Final = create_app(container, Licensed()) + app.dependency_overrides[get_token_verifier] = lambda: verifier + token: Final = jwt.encode({"iss": "https://idp.example.test", "sub": "alice"}, "secret", algorithm="HS256") + response: Final = TestClient(app, raise_server_exceptions=False).get("/api/v1/auth/me", headers=_bearer(token)) + assert response.status_code == 401 + assert response.json()["type"] == "urn:litellm:error:unauthorized" + + +def test_missing_bearer_is_401_error_json_with_a_challenge() -> None: + response: Final = _client().get("/api/v1/auth/me") + assert response.status_code == 401 + assert response.headers["content-type"] == "application/problem+json" + assert response.headers["www-authenticate"].startswith("Bearer ") + assert response.json()["type"] == "urn:litellm:error:unauthorized" + + +def test_garbage_bearer_is_401() -> None: + response: Final = _client().get("/api/v1/auth/me", headers=_bearer("nope")) + assert response.status_code == 401 + assert response.json()["detail"] == "Invalid bearer credentials." + + +def test_master_key_is_a_fleet_wide_cp_admin() -> None: + body: Final = _client().get("/api/v1/auth/me", headers=_bearer(MASTER_KEY)).json() + assert body["roles"] == ["admin"] + assert body["dataplanes"] is None + assert body["token_kind"] == "master_key" + assert "dataplanes:write" in body["scopes"] + + +def test_login_mints_a_session_that_me_accepts() -> None: + client: Final = _client() + login: Final = client.post("/api/v1/auth/login", json={"username": "admin", "password": "pw"}) + assert login.status_code == 200, login.text + assert (login.headers["cache-control"], login.headers["pragma"]) == ("no-store", "no-cache") + token: Final = login.json() + assert token["token_type"] == "bearer" and token["expires_in"] == _settings().session_ttl_seconds + me: Final = client.get("/api/v1/auth/me", headers=_bearer(token["access_token"])).json() + assert (me["subject"], me["roles"], me["token_kind"]) == ("admin", ["admin"], "session") + + +def test_login_throttle_is_client_scoped_and_includes_retry_after() -> None: + settings: Final = _settings() + container: Final = build_container(settings) + clock: Final = iter(repeat(1000.0)) + password: Final = settings.ui_password + assert password is not None + login_service: Final = LoginService( + sessions=SessionTokenCodec( + MASTER_KEY, + settings.session_issuer, + timedelta(seconds=settings.session_ttl_seconds), + container.policy, + ), + policy=container.policy, + ttl_seconds=settings.session_ttl_seconds, + admin_username=settings.ui_username, + admin_password=password.get_secret_value(), + throttle=LoginThrottle(lambda: next(clock)), + ) + app: Final = create_app(container, Licensed()) + app.dependency_overrides[get_login_service] = lambda: login_service + with ( + TestClient(app, client=("198.51.100.1", 50000)) as client_a, + TestClient(app, client=("198.51.100.2", 50000)) as client_b, + ): + failures: Final = tuple( + client_a.post("/api/v1/auth/login", json={"username": "admin", "password": "wrong"}) + for _ in range(LOCKOUT_FAILURES) + ) + assert tuple(response.status_code for response in failures) == (401,) * LOCKOUT_FAILURES + locked: Final = client_a.post("/api/v1/auth/login", json={"username": "admin", "password": "pw"}) + success: Final = client_b.post("/api/v1/auth/login", json={"username": "admin", "password": "pw"}) + + assert locked.status_code == 429 + assert locked.headers["retry-after"] == "300" + assert success.status_code == 200, success.text + + +@pytest.mark.parametrize("body", [{"username": "admin", "password": "wrong"}, {"username": "root", "password": "pw"}]) +def test_wrong_credentials_are_401(body: dict[str, str]) -> None: + response: Final = _client().post("/api/v1/auth/login", json=body) + assert response.status_code == 401 + assert response.json()["type"] == "urn:litellm:error:unauthorized" + + +@pytest.mark.parametrize( + "body", + [ + {"username": "ádmin", "password": "pw"}, + {"username": "admin", "password": "päss"}, + ], +) +def test_non_ascii_login_credentials_are_401(body: dict[str, str]) -> None: + response: Final = _client().post("/api/v1/auth/login", json=body) + assert response.status_code == 401 + assert response.json()["type"] == "urn:litellm:error:unauthorized" + + +def test_non_ascii_bearer_is_401() -> None: + response: Final = _client().get("/api/v1/auth/me", headers={"Authorization": b"Bearer caf\xe9"}) + assert response.status_code == 401 + + +def test_login_is_disabled_when_no_admin_user_is_configured() -> None: + settings: Final = ControlPlaneSettings.model_validate({"LITELLM_MASTER_KEY": MASTER_KEY}) + response: Final = _client(settings).post("/api/v1/auth/login", json={"username": "", "password": ""}) + assert response.status_code == 422 + assert _client(settings).post("/api/v1/auth/login", json={"username": "a", "password": "b"}).status_code == 401 + + +def test_invalid_login_body_is_422_error_json() -> None: + response: Final = _client().post("/api/v1/auth/login", json={"username": "admin"}) + assert response.status_code == 422 + assert response.json()["type"] == "urn:litellm:error:invalid-request-body" + + +def test_oversized_control_plane_request_is_rejected_before_json_parsing() -> None: + response: Final = _client().post( + "/api/v1/auth/login", + content=b'{"username":"admin","password":"' + b"x" * (1100 * 1024) + b'"}', + headers={"content-type": "application/json"}, + ) + assert response.status_code == 413 + assert response.headers["content-type"] == "application/problem+json" + assert response.json()["type"] == "urn:litellm:error:request-too-large" + + +def test_session_scopes_are_reported_and_clamped_to_the_role() -> None: + settings: Final = _settings() + container: Final = build_container(settings) + codec: Final = SessionTokenCodec(MASTER_KEY, settings.session_issuer, timedelta(hours=1), container.policy) + viewer: Final = build_principal( + container.policy, "v", None, frozenset({VIEWER_ROLE}), None, frozenset({"eu"}), "session" + ) + body: Final = _client(settings).get("/api/v1/auth/me", headers=_bearer(codec.mint(viewer))).json() + assert body["dataplanes"] == ["eu"] + assert "proxy:write" not in body["scopes"] and "usage:read" in body["scopes"] + + +def test_master_key_setting_is_a_secret() -> None: + assert isinstance(_settings().master_key, SecretStr) + assert MASTER_KEY not in repr(_settings()) + + +def test_custom_roles_from_settings_are_honoured_on_me() -> None: + settings: Final = _settings(roles={"finops": {"scopes": ["usage:read"], "inherits": ["usage_viewer"]}}) + container: Final = build_container(settings) + codec: Final = SessionTokenCodec(MASTER_KEY, settings.session_issuer, timedelta(hours=1), container.policy) + finops: Final = build_principal(container.policy, "f", None, frozenset({"finops"}), None, frozenset(), "session") + client: Final = TestClient(create_app(container, Licensed())) + body: Final = client.get("/api/v1/auth/me", headers=_bearer(codec.mint(finops))).json() + assert (body["roles"], body["scopes"]) == (["finops"], ["usage:read"]) + + +def test_a_role_granting_an_unknown_scope_fails_startup() -> None: + with pytest.raises(InvalidRolePolicy, match="unknown scope 'nope'"): + build_container(_settings(roles={"x": {"scopes": ["nope"]}})) + + +def test_a_role_inheriting_an_unknown_role_fails_startup() -> None: + with pytest.raises(InvalidRolePolicy, match="inherits unknown role 'ghost'"): + build_container(_settings(roles={"x": {"inherits": ["ghost"]}})) diff --git a/tests/unit/proxy/control_plane/test_settings.py b/tests/unit/proxy/control_plane/test_settings.py new file mode 100644 index 00000000000..71ece6c8e56 --- /dev/null +++ b/tests/unit/proxy/control_plane/test_settings.py @@ -0,0 +1,53 @@ +from typing import Final + +import pytest +from pydantic import ValidationError + +from litellm.proxy.control_plane.settings import ControlPlaneSettings +from litellm.types.proxy.control_plane.auth import PASSWORD_MAX_LENGTH, USERNAME_MAX_LENGTH + +MASTER_KEY: Final = "01234567890123456789012345678901" + + +@pytest.mark.parametrize( + ("username", "password"), + [("u" * (USERNAME_MAX_LENGTH + 1), "p"), ("u", "p" * (PASSWORD_MAX_LENGTH + 1))], +) +def test_configured_admin_credentials_share_the_login_request_bounds(username: str, password: str) -> None: + with pytest.raises(ValidationError): + ControlPlaneSettings.model_validate( + {"LITELLM_MASTER_KEY": MASTER_KEY, "UI_USERNAME": username, "UI_PASSWORD": password} + ) + + +def test_password_login_without_default_roles_requires_a_custom_admin_role() -> None: + with pytest.raises(ValidationError, match="custom admin role"): + ControlPlaneSettings.model_validate( + { + "LITELLM_MASTER_KEY": MASTER_KEY, + "UI_USERNAME": "admin", + "UI_PASSWORD": "pw", + "include_default_roles": False, + } + ) + settings: Final = ControlPlaneSettings.model_validate( + { + "LITELLM_MASTER_KEY": MASTER_KEY, + "UI_USERNAME": "admin", + "UI_PASSWORD": "pw", + "include_default_roles": False, + "roles": {"admin": {"scopes": ["dataplanes:read"]}}, + } + ) + assert not settings.include_default_roles and "admin" in settings.roles + + +def test_credentials_at_the_bound_are_accepted() -> None: + settings: Final = ControlPlaneSettings.model_validate( + { + "LITELLM_MASTER_KEY": MASTER_KEY, + "UI_USERNAME": "u" * USERNAME_MAX_LENGTH, + "UI_PASSWORD": "p" * PASSWORD_MAX_LENGTH, + } + ) + assert settings.ui_username is not None and len(settings.ui_username) == USERNAME_MAX_LENGTH