diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 89c04208d65..4985da9e95d 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -19,7 +19,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `security/` - secret handling and log-leak protection - `router/` - routing and reliability behavior (fallbacks, cooldowns) - `load/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163); do not add a live load test that runs in the default collection. What remains here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend; marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, driven by `.github/workflows/weekly_load_anomaly.yml`) and markerless harness unit tests for the Locust/session-anomaly aggregation logic -- `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite +- `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate, JWT auth (RS256 tokens minted by the test-only issuer in `jwt_issuer.py`, whose JWKS the proxy's `JWT_PUBLIC_KEY_URL` points at; see CONTRIBUTING.md for the start command and config block), and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite - `gateway/` - proxy configuration only (`litellm-config.yml`); no tests - `claude_code/` - the Claude Code compatibility matrix: drives the real `claude` CLI (and HTTP probes) against a proxy for each feature x provider cell, reporting tagged-union outcomes via the `compat_result` fixture; ships its own driver/builder/publisher plus `_*_unit_tests/` trees. The HTTP probes ride the shared transport (`ProxyClient.count_tokens` / `ProxyClient.messages`); the CLI-driving path stays bespoke - `ui/` - the Admin UI browser suite: Playwright in TypeScript, driving the dashboard served by a live proxy on port 4000 (seeded postgres + mock LLM upstream; see its `run_e2e.sh`). It is a self-contained npm package with its own lockfile and does not use the Python harness, pytest markers, or the shared transport; the Python rules in this file (typed models, `Result` unions, basedpyright zero-error gate) do not apply inside it. Its only Python file, `fixtures/mock_llm_server/server.py`, is excluded from the e2e basedpyright gate via the root `pyrightconfig.json` diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index b270feb820e..ad414bb6647 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -27,14 +27,31 @@ The suites run against a live proxy, so bring one up first by running the litell 2. Bring up a Postgres and a Redis for the proxy to use. The repo-root `docker-compose.yml` already defines a Postgres on `5432`; a `docker run -p 6379:6379 redis:7` covers Redis. Point `DATABASE_URL` / `REDIS_HOST` / `REDIS_PORT` at whatever you run. Tests that read Redis directly default to the deployed shape (TLS + cluster mode) whenever `REDIS_HOST` is set, so for a local standalone Redis also set `REDIS_CLUSTER=false` and `REDIS_SSL=false` (plus `REDIS_PASSWORD` when your Redis requires auth) -3. Start the litellm proxy locally against your config and confirm it is live: +3. Start the test-only JWT issuer the `other/` suite mints tokens from, then the litellm proxy against your config, and confirm both are live. The issuer is a fake identity provider (`jwt_issuer.py`) with an open mint endpoint, so it binds loopback only; it generates one RSA key per process and serves it as a JWKS, and the proxy caches that JWKS for `public_key_ttl` (600s) without refetching on an unknown `kid`, so restart the proxy whenever you restart the issuer: ```bash set -a && source .env && set +a - litellm --config .yml --port 4000 + uv run python tests/e2e/jwt_issuer.py & + curl -fs http://127.0.0.1:4190/.well-known/jwks.json + JWT_PUBLIC_KEY_URL=http://127.0.0.1:4190/.well-known/jwks.json litellm --config .yml --port 4000 curl -fs http://localhost:4000/health/liveliness ``` + The issuer's port is `E2E_JWT_ISSUER_PORT` (default `4190`, the port in the URLs above), read by both the issuer process and the tests, so set it in one place if you change it. JWT auth is an enterprise feature, so the proxy also needs `LITELLM_LICENSE` in its environment, and its config needs the JWT block below. `enable_jwt_auth` only routes bearer tokens with three dot-separated segments into the JWT path, so `sk-` virtual keys and the master key keep working for every other suite. `proxy_batch_write_at` is lowered so the JWT spend-attribution test sees its row well inside the poll deadline: + + ```yaml + general_settings: + proxy_batch_write_at: 5 + enable_jwt_auth: true + litellm_jwtauth: + user_id_jwt_field: sub + user_email_jwt_field: email + team_ids_jwt_field: groups + user_id_upsert: true + ``` + + Leave `JWT_AUDIENCE` and `JWT_ISSUER` unset unless you also put matching `aud` / `iss` claims in the tokens the tests mint; the issuer sets `iss` to its own base URL + 4. Run a suite against it; the harness reads `LITELLM_PROXY_URL` (default `http://localhost:4000`): ```bash diff --git a/tests/e2e/coverage_registry/other.yaml b/tests/e2e/coverage_registry/other.yaml index 814ebae2e0b..8b9cb4d2adf 100644 --- a/tests/e2e/coverage_registry/other.yaml +++ b/tests/e2e/coverage_registry/other.yaml @@ -9,9 +9,12 @@ - {id: other.auth.llm_chat.not_bearer_scheme_denied, module: other, tier: P0, area: auth, assertions: [not_bearer_scheme_denied], source: "vendor testing strategy §11.1 / LIT-4778", rationale: "NotBearer scheme on chat is 401/403"} - {id: other.auth.realtime.missing_header_denied, module: other, tier: P1, area: auth, assertions: [missing_header_denied], source: "vendor testing strategy §9.19 / LIT-4778", rationale: "Realtime client-secret and calls routes reject requests without Authorization"} - {id: other.config.responses.metadata_redis_ttl_bounded, module: other, tier: P0, area: config, assertions: [ttl_bounded], source: "responses + redis cache", rationale: "Responses store+metadata must not leave TTL-unbounded Redis entries (LIT-1201)"} -- {id: other.auth.jwt.valid_token_allows, module: other, tier: P0, area: auth, assertions: [valid_token_allows], source: "handle_jwt.py:77-150", rationale: "Valid JWT with correct issuer + claims grants access"} -- {id: other.auth.jwt.expired_denied, module: other, tier: P0, area: auth, assertions: [expired_denied], source: "handle_jwt.py:125-135", rationale: "Expired JWT rejected even with valid signature"} -- {id: other.auth.jwt.invalid_signature_denied, module: other, tier: P0, area: auth, assertions: [invalid_signature_denied], source: "handle_jwt.py:145-150", rationale: "Bad/missing signature fails verification"} +- {id: other.auth.jwt.valid_token_allows, module: other, tier: P0, area: auth, assertions: [valid_token_allows], source: "handle_jwt.py:1217-1256 auth_jwt / user_api_key_auth.py:1365-1377", rationale: "An RS256 JWT signed by the configured JWKS whose groups claim names an existing team is accepted on /chat/completions"} +- {id: other.auth.jwt.spend_attributed_to_claims, module: other, tier: P0, area: auth, assertions: [spend_attributed_to_claims], source: "handle_jwt.py:2224 auth_builder / user_api_key_auth.py:1438-1474", rationale: "The spend log row for a JWT-authenticated call carries the team_id from the groups claim and the user_id from sub, not a virtual key's identity. Single-group claim only: the proxy picks the team from a set, so attribution over several groups is unordered"} +- {id: other.auth.jwt.expired_denied, module: other, tier: P0, area: auth, assertions: [expired_denied], source: "handle_jwt.py:1244-1250", rationale: "Expired JWT rejected 401 (Token Expired) even with valid signature; leeway is 0"} +- {id: other.auth.jwt.invalid_signature_denied, module: other, tier: P0, area: auth, assertions: [invalid_signature_denied], source: "handle_jwt.py:1158-1166 _decode_jwt_with_public_key", rationale: "A genuine token whose signature bytes were altered fails verification with 401"} +- {id: other.auth.jwt.unknown_team_denied, module: other, tier: P0, area: auth, assertions: [unknown_team_denied], source: "handle_jwt.py:1549-1632 find_team_with_model_access", rationale: "A verified JWT whose groups claim resolves to no existing team is denied with 403 naming the unresolved team, never silently admitted without a team. The proxy words it as a model-access denial, the same body an existing team without model access gets"} +- {id: other.auth.jwt.virtual_key_unaffected, module: other, tier: P0, area: auth, assertions: [virtual_key_unaffected], source: "handle_jwt.py:213 is_jwt / user_api_key_auth.py:1332-1333", rationale: "enable_jwt_auth only routes three-segment bearer tokens into the JWT branch, so sk- virtual keys keep working on the same proxy"} - {id: other.auth.model_access_group.wildcard_bare_name_allowed, module: other, tier: P0, area: auth, assertions: [wildcard_bare_name_allowed], source: "auth_checks.py:3232 / LIT-5813", fail_before_fix: proven, rationale: "A grant of a group holding a wildcard deployment covers the bare model names callers actually send, not only the provider-prefixed spelling"} - {id: other.auth.model_access_group.member_allowed, module: other, tier: P0, area: auth, assertions: [member_allowed], source: "auth_checks.py:3232", rationale: "A key whose allow-list is a model access group can call the deployments in that group"} - {id: other.auth.model_access_group.non_member_denied, module: other, tier: P0, area: auth, assertions: [non_member_denied], source: "auth_checks.py:3232", rationale: "That same grant reaches nothing outside the group, including provider models the group's wildcard does not cover"} diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index 09d17b9a0db..0e7f4c0ad2b 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -15,6 +15,7 @@ from typing import Final from dotenv import load_dotenv from fixture_mode import deterministic_marker, parse_fixture_mode +from jwt_issuer import jwt_issuer_url from provider_edge import provider_edge_api_base # Local runs keep provider / DataDog keys in tests/e2e/.env (see CONTRIBUTING.md). @@ -43,6 +44,9 @@ UI_BASE_URL = os.environ.get("E2E_UI_BASE_URL", PROXY_BASE_URL).rstrip("/") CHEAP_ANTHROPIC_MODEL = os.environ.get("E2E_CHEAP_ANTHROPIC_MODEL", "claude-haiku-4-5") CHEAP_OPENAI_MODEL = os.environ.get("E2E_CHEAP_OPENAI_MODEL", "gpt-5.5") +# Test-only JWT issuer (jwt_issuer.py); the port is E2E_JWT_ISSUER_PORT (see CONTRIBUTING.md). +JWT_ISSUER_URL = jwt_issuer_url() + LINEAR_MCP_URL = os.environ.get("E2E_LINEAR_MCP_URL", "https://mcp.linear.app/mcp") LINEAR_STORAGE_STATE = os.environ.get("E2E_LINEAR_STORAGE_STATE", "") diff --git a/tests/e2e/e2e_http.py b/tests/e2e/e2e_http.py index 9d5f1658e91..4c65c53b653 100644 --- a/tests/e2e/e2e_http.py +++ b/tests/e2e/e2e_http.py @@ -386,6 +386,27 @@ def get_external[R: BaseModel]( return _classify(resp, response_type) +def post_external[R: BaseModel]( + url: str, + *, + json: BaseModel, + response_type: type[R], + timeout: float = 30.0, +) -> Result[R]: + """POST an absolute URL outside the proxy (e.g. the e2e JWT issuer's mint + endpoint). Like get_external: no proxy base url, no proxy auth, and the same + tagged-union classification as every other call.""" + try: + resp = requests.post( + url, + json=json.model_dump(by_alias=True, exclude_none=True), + timeout=timeout, + ) + except requests.RequestException as exc: + return NetworkError(message=str(exc)) + return _classify(resp, response_type) + + def delete[R: BaseModel]( url: URL, *, diff --git a/tests/e2e/jwt_issuer.py b/tests/e2e/jwt_issuer.py new file mode 100644 index 00000000000..4006fcaa314 --- /dev/null +++ b/tests/e2e/jwt_issuer.py @@ -0,0 +1,230 @@ +"""Test-only fake identity provider for the e2e JWT suite. Never deploy it. + +Run it next to the proxy (`uv run python tests/e2e/jwt_issuer.py`). On start it +generates one RSA signing key and keeps it for the life of the process, serving +the public half at `GET /.well-known/jwks.json` and signing whatever claims are +POSTed to `/token`. The proxy's `JWT_PUBLIC_KEY_URL` points at the JWKS URL, and +tests mint RS256 tokens by POSTing claims, so the private key never leaves this +process and no test holds it. + +One key per process, rather than per pytest run, is what survives the proxy's +JWKS cache: the proxy caches the JWKS for `litellm_jwtauth.public_key_ttl` +(600s by default) and does not refetch on an unknown `kid`, so a key rotated +every run would be rejected until the cache expired. Restart the proxy whenever +you restart the issuer. + +The mint endpoint takes no credential: anyone who can reach it gets a token the +proxy trusts. It therefore binds 127.0.0.1 only, must never be exposed beyond +loopback, and must only ever be trusted by a proxy under test. +""" + +from __future__ import annotations + +import logging +import os +import threading +import time +import uuid +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import Final, Literal + +import jwt +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric import rsa +from jwt.utils import to_base64url_uint +from pydantic import BaseModel, RootModel, ValidationError + +JWT_ISSUER_PORT_ENV: Final = "E2E_JWT_ISSUER_PORT" +DEFAULT_JWT_ISSUER_PORT: Final = 4190 +LOOPBACK_HOST: Final = "127.0.0.1" +JWKS_PATH: Final = "/.well-known/jwks.json" +TOKEN_PATH: Final = "/token" +DEFAULT_TOKEN_LIFETIME_SECONDS: Final = 300 + +ClaimValue = str | int | float | bool | None | list[str] + + +def jwt_issuer_port() -> int: + raw: Final = os.environ.get(JWT_ISSUER_PORT_ENV, "").strip() + return int(raw) if raw else DEFAULT_JWT_ISSUER_PORT + + +def jwt_issuer_url() -> str: + return f"http://{LOOPBACK_HOST}:{jwt_issuer_port()}" + + +class RsaJwk(BaseModel): + kty: Literal["RSA"] = "RSA" + alg: Literal["RS256"] = "RS256" + use: Literal["sig"] = "sig" + kid: str + n: str + e: str + + +class JwksDocument(BaseModel): + keys: tuple[RsaJwk, ...] + + +class TokenRequest(RootModel[Mapping[str, ClaimValue]]): + """The JSON body of POST /token: the claims to sign, verbatim. Nested objects + are not supported; every value is a scalar or a list of strings.""" + + +class MintedToken(BaseModel): + token: str + + +class IssuerError(BaseModel): + error: str + + +@dataclass(frozen=True, slots=True) +class SigningKey: + kid: str + private_pem: str + jwk: RsaJwk + + +def generate_signing_key(kid: str | None = None) -> SigningKey: + private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + numbers: Final = private_key.public_key().public_numbers() + resolved_kid: Final = kid if kid is not None else uuid.uuid4().hex + return SigningKey( + kid=resolved_kid, + private_pem=private_key.private_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PrivateFormat.PKCS8, + encryption_algorithm=serialization.NoEncryption(), + ).decode(), + jwk=RsaJwk( + kid=resolved_kid, + n=to_base64url_uint(numbers.n).decode(), + e=to_base64url_uint(numbers.e).decode(), + ), + ) + + +def mint( + key: SigningKey, + claims: Mapping[str, ClaimValue], + *, + issuer: str, + now: int, + lifetime_seconds: int = DEFAULT_TOKEN_LIFETIME_SECONDS, +) -> str: + """Sign `claims` as a compact RS256 JWT carrying `key.kid` in its header. + `iss`, `iat`, and `exp` are filled in when absent and left alone when the + caller sets them, so a test can mint an already-expired token.""" + payload: Final[dict[str, ClaimValue]] = { + "iss": issuer, + "iat": now, + "exp": now + lifetime_seconds, + **claims, + } + return jwt.encode(payload, key.private_pem, algorithm="RS256", headers={"kid": key.kid}) + + +class _IssuerServer(ThreadingHTTPServer): + daemon_threads = True + + def __init__(self, bind: tuple[str, int], *, key: SigningKey, clock: Callable[[], int]) -> None: + super().__init__(bind, _IssuerHandler) + self.key: Final = key + self.clock: Final = clock + + @property + def url(self) -> str: + host, port = self.server_address[0], self.server_address[1] + return f"http://{host}:{port}" + + +class _IssuerHandler(BaseHTTPRequestHandler): + def _issuer(self) -> _IssuerServer: + issuer: Final = self.server + assert isinstance(issuer, _IssuerServer) + return issuer + + def do_GET(self) -> None: + if self.path != JWKS_PATH: + self._send(404, IssuerError(error=f"unknown path {self.path}; the JWKS is at {JWKS_PATH}")) + return + self._send(200, JwksDocument(keys=(self._issuer().key.jwk,))) + + def do_POST(self) -> None: + if self.path != TOKEN_PATH: + self._send(404, IssuerError(error=f"unknown path {self.path}; mint tokens at {TOKEN_PATH}")) + return + length: Final = int(self.headers.get("Content-Length", "0")) + try: + request: Final = TokenRequest.model_validate_json(self.rfile.read(length)) + except ValidationError as exc: + self._send(400, IssuerError(error=f"claims must be a JSON object of scalar or string-list values: {exc}")) + return + issuer: Final = self._issuer() + token: Final = mint(issuer.key, request.root, issuer=issuer.url, now=issuer.clock()) + self._send(200, MintedToken(token=token)) + + def _send(self, status: int, body: BaseModel) -> None: + payload: Final = body.model_dump_json().encode() + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + + +@dataclass(frozen=True, slots=True) +class RunningIssuer: + url: str + key: SigningKey + server: _IssuerServer + + @property + def jwks_url(self) -> str: + return f"{self.url}{JWKS_PATH}" + + def shutdown(self) -> None: + self.server.shutdown() + self.server.server_close() + + +def _wall_clock() -> int: + return int(time.time()) + + +def start_jwt_issuer( + *, + port: int = 0, + key: SigningKey | None = None, + clock: Callable[[], int] = _wall_clock, +) -> RunningIssuer: + """Serve the issuer on loopback in a daemon thread. `port=0` takes an + OS-assigned port for in-process tests; the CLI passes the documented one.""" + server: Final = _IssuerServer((LOOPBACK_HOST, port), key=key or generate_signing_key(), clock=clock) + thread: Final = threading.Thread(target=server.serve_forever, name="e2e-jwt-issuer", daemon=True) + thread.start() + return RunningIssuer(url=server.url, key=server.key, server=server) + + +def main() -> None: + logging.basicConfig(level=logging.INFO, format="%(message)s") + running: Final = start_jwt_issuer(port=jwt_issuer_port()) + logging.getLogger(__name__).info( + "e2e jwt issuer listening on %s jwks=%s mint=POST %s%s kid=%s (test-only, loopback, ctrl-c to stop)", + running.url, + running.jwks_url, + running.url, + TOKEN_PATH, + running.key.kid, + ) + try: + threading.Event().wait() + except KeyboardInterrupt: + running.shutdown() + + +if __name__ == "__main__": + main() diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 2c6c0e9bbd4..17f5418bd73 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -1078,6 +1078,18 @@ class UserListResponse(BaseModel): total: int +class JwtClaimsBody(BaseModel): + """Claims POSTed to the e2e JWT issuer's /token: exactly what the + `litellm_jwtauth` block in CONTRIBUTING.md reads (sub -> user_id, email -> + user_email, groups -> team ids), plus an explicit `exp` for the expired case; + the issuer fills in iss/iat/exp when they are left unset.""" + + sub: str + email: str + groups: Sequence[str] + exp: int | None = None + + class OrgNewBody(BaseModel): organization_alias: str models: list[str] = [] diff --git a/tests/e2e/other/other_client.py b/tests/e2e/other/other_client.py index 1aa83ac42c7..b51f3338f52 100644 --- a/tests/e2e/other/other_client.py +++ b/tests/e2e/other/other_client.py @@ -1,19 +1,27 @@ """Client for the `other` holding-pen suite: the auth gate (master key vs an -invalid key on an admin route) and the process-lifecycle health probes -(liveness, public readiness, authenticated readiness diagnostics). +invalid key on an admin route), JWT auth against the test-only issuer +(jwt_issuer.py), and the process-lifecycle health probes (liveness, public +readiness, authenticated readiness diagnostics). Holds the shared ProxyClient so `resources` / `scoped_key` still clean up, and adds only the routes these behaviors need. The health probes deliberately send no auth header (public routes), so they go through the transport with an empty -headers model rather than a bearer. +headers model rather than a bearer. Tokens are minted by POSTing claims to the +issuer, so no test ever holds a signing key. """ from __future__ import annotations from dataclasses import dataclass +from typing import Final -from e2e_http import NoBody, ProbeResult, Result +import pytest + +from e2e_config import JWT_ISSUER_URL +from e2e_http import NetworkError, NoBody, ProbeResult, Result, Success, post_external +from jwt_issuer import TOKEN_PATH, MintedToken from models import ( + JwtClaimsBody, ReadinessDetailsResponse, ReadinessResponse, UserListParams, @@ -25,6 +33,7 @@ from proxy_client import ProxyClient @dataclass(frozen=True, slots=True) class OtherClient: proxy: ProxyClient + jwt_issuer_url: str def liveness(self) -> ProbeResult: """GET /health/liveliness. Unauthenticated; the probe returns status + @@ -68,6 +77,21 @@ class OtherClient: response_type=UserListResponse, ) + def mint_jwt(self, claims: JwtClaimsBody) -> str: + """Have the test-only issuer sign `claims` into a compact RS256 JWT. A + missing issuer is a hard failure naming the start command, not a skip.""" + result: Final = post_external(f"{self.jwt_issuer_url}{TOKEN_PATH}", json=claims, response_type=MintedToken) + match result: + case Success(data=minted): + return minted.token + case NetworkError(message=message): + pytest.fail( + f"No live JWT issuer at {self.jwt_issuer_url}: {message}. Start it next to the proxy with " + "`uv run python tests/e2e/jwt_issuer.py` (see CONTRIBUTING.md)" + ) + case _: + raise AssertionError(result) + def build_client(proxy: ProxyClient) -> OtherClient: - return OtherClient(proxy=proxy) + return OtherClient(proxy=proxy, jwt_issuer_url=JWT_ISSUER_URL) diff --git a/tests/e2e/other/test_jwt_auth_e2e.py b/tests/e2e/other/test_jwt_auth_e2e.py new file mode 100644 index 00000000000..49a9350a3a2 --- /dev/null +++ b/tests/e2e/other/test_jwt_auth_e2e.py @@ -0,0 +1,129 @@ +"""Live e2e: RS256 JWTs minted by the test-only issuer (jwt_issuer.py) against a +proxy running with `enable_jwt_auth: true` and the `litellm_jwtauth` block from +CONTRIBUTING.md (sub -> user_id, email -> user_email, groups -> team ids, +user_id_upsert). + +Every case mints through the issuer, so the tests never hold a signing key: the +bad-signature case corrupts a genuine signature, the expired case asks the +issuer for a token whose `exp` is already in the past. Those identities get +their own freshly created team so a rejection can only be blamed on the token, +while the unknown-team case names a team that was never created. An acceptance +is proven twice, at the boundary (200 from a real provider) and in the spend log +the proxy attributes to the claims. The last case keeps a plain `sk-` virtual +key working on the same proxy, guarding against the flag turning JWT on for +everyone. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Final + +import pytest + +from e2e_config import CHEAP_OPENAI_MODEL, unique_marker +from e2e_http import UnauthorizedError, UnknownApiError, unwrap +from lifecycle import ResourceManager +from models import ChatBody, ChatMessage, JwtClaimsBody, TeamNewBody +from other_client import OtherClient + +pytestmark = pytest.mark.e2e + + +@dataclass(frozen=True, slots=True) +class JwtIdentity: + user_id: str + team_id: str + + def claims(self, *, exp: int | None = None) -> JwtClaimsBody: + return JwtClaimsBody(sub=self.user_id, email=f"{self.user_id}@example.com", groups=(self.team_id,), exp=exp) + + +@pytest.fixture +def identity(client: OtherClient, resources: ResourceManager) -> JwtIdentity: + marker: Final = unique_marker() + team_id: Final = client.proxy.create_team( + TeamNewBody(team_alias=f"e2e-jwt-{marker}", team_id=f"e2e-jwt-team-{marker}") + ) + resources.defer(lambda: client.proxy.delete_team(team_id)) + user_id: Final = f"e2e-jwt-user-{marker}" + resources.defer(lambda: client.proxy.delete_user(user_id)) + return JwtIdentity(user_id=user_id, team_id=team_id) + + +def _ping() -> ChatBody: + return ChatBody( + model=CHEAP_OPENAI_MODEL, + messages=[ChatMessage(role="user", content=f"Reply with the single word pong. {unique_marker()}")], + max_tokens=16, + ) + + +def _corrupt_signature(token: str) -> str: + header, payload, signature = token.split(".") + flipped: Final = "A" if signature[10] != "A" else "B" + return f"{header}.{payload}.{signature[:10]}{flipped}{signature[11:]}" + + +class TestJwtAuth: + @pytest.mark.covers("other.auth.jwt.valid_token_allows", "other.auth.jwt.spend_attributed_to_claims") + def test_valid_token_for_an_existing_team_is_accepted_and_attributed( + self, client: OtherClient, identity: JwtIdentity + ) -> None: + token: Final = client.mint_jwt(identity.claims()) + + response: Final = unwrap(client.proxy.chat(token, _ping())) + assert response.id is not None and response.choices, ( + f"chat under a valid JWT returned no completion: {response}" + ) + + rows: Final = client.proxy.poll_logs_for_request_id(response.id) + assert rows, f"no spend log row for request {response.id} within the poll deadline" + row: Final = rows[0] + assert row.team_id == identity.team_id, ( + f"spend row must carry the team from the JWT groups claim {identity.team_id!r}, got {row.team_id!r}" + ) + assert row.user == identity.user_id, ( + f"spend row must carry the user from the JWT sub claim {identity.user_id!r}, got {row.user!r}" + ) + + @pytest.mark.covers("other.auth.jwt.invalid_signature_denied") + def test_tampered_signature_is_rejected(self, client: OtherClient, identity: JwtIdentity) -> None: + tampered: Final = _corrupt_signature(client.mint_jwt(identity.claims())) + + result: Final = client.proxy.chat(tampered, _ping()) + assert isinstance(result, UnauthorizedError), ( + f"a JWT whose signature does not verify must be rejected with 401, got {result}" + ) + assert "signature verification failed" in result.body.lower(), ( + f"the 401 must come from signature verification, not another auth failure, got {result.body[:300]}" + ) + + @pytest.mark.covers("other.auth.jwt.expired_denied") + def test_expired_token_is_rejected(self, client: OtherClient, identity: JwtIdentity) -> None: + expired: Final = client.mint_jwt(identity.claims(exp=1)) + + result: Final = client.proxy.chat(expired, _ping()) + assert isinstance(result, UnauthorizedError), ( + f"an expired JWT must be rejected with 401 even though its signature verifies, got {result}" + ) + assert "expired" in result.body.lower(), f"the 401 must say the token expired, got {result.body[:300]}" + + @pytest.mark.covers("other.auth.jwt.unknown_team_denied") + def test_token_naming_a_team_that_does_not_exist_is_rejected(self, client: OtherClient) -> None: + marker: Final = unique_marker() + never_created: Final = JwtIdentity(user_id=f"e2e-jwt-user-{marker}", team_id=f"e2e-jwt-missing-team-{marker}") + token: Final = client.mint_jwt(never_created.claims()) + + result: Final = client.proxy.chat(token, _ping()) + assert isinstance(result, UnknownApiError) and result.status_code == 403, ( + f"a valid JWT whose groups name no existing team must be rejected with 403, got {result}" + ) + assert never_created.team_id in result.body, ( + f"the 403 must name the team it could not resolve ({never_created.team_id}), got {result.body[:300]}" + ) + + @pytest.mark.covers("other.auth.jwt.virtual_key_unaffected") + def test_plain_virtual_key_still_works_with_jwt_auth_enabled(self, client: OtherClient, scoped_key: str) -> None: + response: Final = unwrap(client.proxy.chat(scoped_key, _ping())) + assert response.choices, f"an sk- key must keep working on a proxy with enable_jwt_auth, got {response}" diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index cdc20e5299a..5292f56b324 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -21,6 +21,7 @@ from e2e_http import ( Result, StreamingResponse, Success, + UnknownApiError, is_ok, unwrap, ) @@ -65,6 +66,11 @@ from models import ( SpendLogsPage, SpendLogsPageParams, SpendLogsParams, + TeamDeleteBody, + TeamNewBody, + TeamNewResponse, + UserDeleteBody, + UserDeleteResponse, ) from e2e_config import ( CONTROL_PLANE_BASE_URL, @@ -409,6 +415,41 @@ class ProxyClient: if not is_ok(result): warnings.warn(f"delete_credential({credential_name!r}) failed: {result}", stacklevel=2) + def create_team(self, body: TeamNewBody) -> str: + return unwrap( + self.transport.post( + "/team/new", + headers=self.transport.master, + json=body, + response_type=TeamNewResponse, + ) + ).team_id + + def delete_team(self, team_id: str) -> None: + result = self.transport.post( + "/team/delete", + headers=self.transport.master, + json=TeamDeleteBody(team_ids=[team_id]), + response_type=NoBody, + ) + if not is_ok(result): + warnings.warn(f"delete_team({team_id!r}) failed: {result}", stacklevel=2) + + def delete_user(self, user_id: str) -> None: + """Best-effort teardown; a 404 is not a leak, since JWT tests defer this for + a user the proxy only upserts after a successful auth.""" + result = self.transport.post( + "/user/delete", + headers=self.transport.master, + json=UserDeleteBody(user_ids=[user_id]), + response_type=UserDeleteResponse, + ) + match result: + case Success() | UnknownApiError(status_code=404): + return + case _: + warnings.warn(f"delete_user({user_id!r}) failed: {result}", stacklevel=2) + # ---- LLM calls ------------------------------------------------------ def chat(self, key: str, body: ChatBody) -> Result[ChatResponse]: diff --git a/tests/e2e/test_jwt_issuer.py b/tests/e2e/test_jwt_issuer.py new file mode 100644 index 00000000000..7acaaf13949 --- /dev/null +++ b/tests/e2e/test_jwt_issuer.py @@ -0,0 +1,141 @@ +"""Harness coverage for the test-only JWT issuer (jwt_issuer.py). + +No proxy and no ``e2e`` marker. The issuer is booted in-process on an +OS-assigned port with a fixed clock and driven over HTTP through +``e2e_http.post_external`` / ``get_external``, the same transport the live +suite uses, so what is pinned here is the contract the live JWT tests lean on: +a token minted at ``/token`` verifies against the key served at +``/.well-known/jwks.json`` under the ``kid`` in its header, ``iss``/``iat``/``exp`` +are filled in only when the caller left them out, and malformed claim bodies or +unknown paths are refused instead of signed. +""" + +from __future__ import annotations + +import time +from collections.abc import Iterator +from typing import Final + +import jwt +import pytest +from pydantic import BaseModel, RootModel + +from e2e_http import UnknownApiError, get_external, post_external, unwrap +from jwt_issuer import ( + DEFAULT_TOKEN_LIFETIME_SECONDS, + JWKS_PATH, + JWT_ISSUER_PORT_ENV, + TOKEN_PATH, + JwksDocument, + MintedToken, + RunningIssuer, + jwt_issuer_port, + start_jwt_issuer, +) + +FROZEN_NOW: Final = int(time.time()) + + +class _Claims(BaseModel): + sub: str + groups: tuple[str, ...] = () + exp: int | None = None + + +class _DecodedClaims(BaseModel): + sub: str + iss: str + iat: int + exp: int + groups: tuple[str, ...] = () + + +class _NotAnObject(RootModel[tuple[str, ...]]): + pass + + +@pytest.fixture(scope="module") +def issuer() -> Iterator[RunningIssuer]: + running: Final = start_jwt_issuer(clock=lambda: FROZEN_NOW) + yield running + running.shutdown() + + +def _mint(issuer: RunningIssuer, claims: _Claims) -> str: + return unwrap(post_external(f"{issuer.url}{TOKEN_PATH}", json=claims, response_type=MintedToken)).token + + +def _served_jwks(issuer: RunningIssuer) -> JwksDocument: + return unwrap(get_external(issuer.jwks_url, response_type=JwksDocument)) + + +def _decode(token: str, jwks: JwksDocument, *, verify_exp: bool = True) -> _DecodedClaims: + key: Final = jwt.PyJWK.from_json(jwks.keys[0].model_dump_json()) + decoded: Final = jwt.decode(token, key, algorithms=["RS256"], options={"verify_exp": verify_exp}) + return _DecodedClaims.model_validate(decoded) + + +class TestJwtIssuer: + def test_minted_token_verifies_against_the_served_jwks(self, issuer: RunningIssuer) -> None: + token: Final = _mint(issuer, _Claims(sub="alice", groups=("team-a",))) + jwks: Final = _served_jwks(issuer) + + assert len(jwks.keys) == 1 + assert jwt.get_unverified_header(token)["kid"] == jwks.keys[0].kid + claims: Final = _decode(token, jwks) + assert claims.sub == "alice" + assert claims.groups == ("team-a",) + assert claims.iss == issuer.url + assert claims.iat == FROZEN_NOW + assert claims.exp == FROZEN_NOW + DEFAULT_TOKEN_LIFETIME_SECONDS + + def test_a_token_signed_by_another_key_does_not_verify(self, issuer: RunningIssuer) -> None: + other: Final = start_jwt_issuer(clock=lambda: FROZEN_NOW) + try: + foreign_token: Final = _mint(other, _Claims(sub="alice")) + finally: + other.shutdown() + + with pytest.raises(jwt.InvalidSignatureError): + _decode(foreign_token, _served_jwks(issuer)) + + def test_an_explicit_exp_is_signed_as_given(self, issuer: RunningIssuer) -> None: + expired_at: Final = FROZEN_NOW - 60 + token: Final = _mint(issuer, _Claims(sub="alice", exp=expired_at)) + jwks: Final = _served_jwks(issuer) + + assert _decode(token, jwks, verify_exp=False).exp == expired_at + with pytest.raises(jwt.ExpiredSignatureError): + _ = _decode(token, jwks) + + def test_non_object_claims_are_refused(self, issuer: RunningIssuer) -> None: + result: Final = post_external( + f"{issuer.url}{TOKEN_PATH}", json=_NotAnObject(("not", "claims")), response_type=MintedToken + ) + assert isinstance(result, UnknownApiError) + assert result.status_code == 400 + + @pytest.mark.parametrize( + ("method", "path"), + [("GET", TOKEN_PATH), ("POST", JWKS_PATH), ("GET", "/token/anything")], + ids=["get-token", "post-jwks", "get-other"], + ) + def test_unknown_routes_are_404(self, issuer: RunningIssuer, method: str, path: str) -> None: + url: Final = f"{issuer.url}{path}" + result: Final = ( + get_external(url, response_type=MintedToken) + if method == "GET" + else post_external(url, json=_Claims(sub="alice"), response_type=MintedToken) + ) + assert isinstance(result, UnknownApiError) + assert result.status_code == 404 + + +class TestIssuerPort: + def test_defaults_to_the_documented_port(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv(JWT_ISSUER_PORT_ENV, raising=False) + assert jwt_issuer_port() == 4190, "CONTRIBUTING.md hardcodes 4190 in JWT_PUBLIC_KEY_URL" + + def test_env_override_wins(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv(JWT_ISSUER_PORT_ENV, " 4321 ") + assert jwt_issuer_port() == 4321