diff --git a/tests/integration/authorization/test_jwt_algorithm_allowlist.py b/tests/integration/authorization/test_jwt_algorithm_allowlist.py index 210d40a0d6a..93dc98c17c4 100644 --- a/tests/integration/authorization/test_jwt_algorithm_allowlist.py +++ b/tests/integration/authorization/test_jwt_algorithm_allowlist.py @@ -1,22 +1,65 @@ +import asyncio +import base64 import json +import os +import socket import time import uuid +from concurrent.futures import ThreadPoolExecutor from pathlib import Path from typing import Final +import httpx import jwt +import psutil import yaml -from cryptography.hazmat.primitives.asymmetric import ed25519, rsa +from anthropic import Anthropic +from cryptography.hazmat.primitives.asymmetric import ec, ed25519, rsa +from jwt.algorithms import ECAlgorithm, OKPAlgorithm, RSAAlgorithm +from openai import AsyncOpenAI, OpenAI from tests.integration._support.client import Gateway, eventually +from tests.integration._support.database import read_rows from tests.integration._support.process import owned_proxy_process from tests.integration._support.wire import Reply, Request, wire_server KEY_ID: Final = "integration-jwt-allowlist-key" +APPROVED_WARNING_PARTS: Final = ("EdDSA", "deprecated", "LITELLM_FIPS_MODE") def _jwks_reply(public_jwk: str) -> Reply: - return Reply(body=json.dumps({"keys": [{**json.loads(public_jwk), "kid": KEY_ID}]}).encode()) + return _jwks_reply_keys([{**json.loads(public_jwk), "kid": KEY_ID}]) + + +def _jwks_reply_keys(keys: list[dict[str, object]]) -> Reply: + return Reply(body=json.dumps({"keys": keys}).encode()) + + +def _eddsa_keypair() -> tuple[ed25519.Ed25519PrivateKey, dict[str, object]]: + private_key: Final = ed25519.Ed25519PrivateKey.generate() + return private_key, {**json.loads(OKPAlgorithm.to_jwk(private_key.public_key())), "kid": "ed"} + + +def _rsa_keypair() -> tuple[rsa.RSAPrivateKey, dict[str, object]]: + private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + return private_key, {**json.loads(RSAAlgorithm.to_jwk(private_key.public_key())), "kid": "rsa"} + + +def _ec_keypair() -> tuple[ec.EllipticCurvePrivateKey, dict[str, object]]: + private_key: Final = ec.generate_private_key(ec.SECP256R1()) + return private_key, {**json.loads(ECAlgorithm.to_jwk(private_key.public_key())), "kid": "ec"} + + +def _free_port() -> int: + with socket.socket() as reserve: + reserve.bind(("127.0.0.1", 0)) + return reserve.getsockname()[1] + + +def _observed_bodies() -> tuple[str, ...]: + response: Final = httpx.get(f"{os.environ['INTEGRATION_UPSTREAM_URL']}/__observations", timeout=15) + assert response.status_code == 200, response.text + return tuple(json.dumps(entry["body"]) for entry in response.json()["requests"]) def _jwt_auth_config(tmp_path: Path) -> Path: @@ -31,12 +74,12 @@ def _jwt_auth_config(tmp_path: Path) -> Path: return path -def _token(key: object, algorithm: str, subject: str) -> str: +def _token(key: object, algorithm: str, subject: str, kid: str = KEY_ID) -> str: return jwt.encode( {"sub": subject, "iat": int(time.time()), "exp": int(time.time()) + 300}, key, # pyright: ignore[reportArgumentType] # jwt.encode takes Any key material algorithm=algorithm, - headers={"kid": KEY_ID}, + headers={"kid": kid}, ) @@ -44,7 +87,7 @@ def test_eddsa_signed_token_is_accepted_and_logged_as_deprecated_outside_fips_mo gateway: Gateway, tmp_path: Path ) -> None: private_key: Final = ed25519.Ed25519PrivateKey.generate() - public_jwk: Final = jwt.algorithms.OKPAlgorithm.to_jwk(private_key.public_key()) + public_jwk: Final = OKPAlgorithm.to_jwk(private_key.public_key()) def respond(request: Request) -> Reply: assert request.method == "GET", request @@ -75,7 +118,7 @@ def test_eddsa_signed_token_is_accepted_and_logged_as_deprecated_outside_fips_mo def test_rs256_signed_token_is_accepted_without_deprecation_log(gateway: Gateway, tmp_path: Path) -> None: private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) - public_jwk: Final = jwt.algorithms.RSAAlgorithm.to_jwk(private_key.public_key()) + public_jwk: Final = RSAAlgorithm.to_jwk(private_key.public_key()) def respond(request: Request) -> Reply: assert request.method == "GET", request @@ -97,3 +140,548 @@ def test_rs256_signed_token_is_accepted_without_deprecation_log(gateway: Gateway ) assert response.status_code == 200, response.text assert "EdDSA" not in owned.log.read_text(), owned.log.read_text() + + +def _deprecation_lines(log_text: str) -> list[str]: + return [line for line in log_text.splitlines() if all(part in line for part in APPROVED_WARNING_PARTS)] + + +def _alg_lied_token(private_key: ed25519.Ed25519PrivateKey, kid: str, subject: str) -> str: + def segment(value: bytes) -> str: + return base64.urlsafe_b64encode(value).rstrip(b"=").decode() + + signing_input: Final = ( + f"{segment(json.dumps({'alg': 'RS256', 'typ': 'JWT', 'kid': kid}).encode())}." + f"{segment(json.dumps({'sub': subject, 'iat': int(time.time()), 'exp': int(time.time()) + 300}).encode())}" + ) + signature: Final = segment(private_key.sign(signing_input.encode())) + return f"{signing_input}.{signature}" + + +def test_eddsa_non_stream_request_reaches_upstream_and_logs_deprecation(gateway: Gateway, tmp_path: Path) -> None: + private_key, jwks_key = _eddsa_keypair() + + def respond(request: Request) -> Reply: + assert request.method == "GET", request + return _jwks_reply_keys([jwks_key]) + + with wire_server(respond) as jwks, gateway.scenario() as scenario: + subject: Final = f"integration-jwt-{uuid.uuid4().hex}" + token: Final = _token(private_key, "EdDSA", subject, kid="ed") + with owned_proxy_process( + gateway, tmp_path, {"JWT_PUBLIC_KEY_URL": jwks.url}, config=_jwt_auth_config(tmp_path), workers=2 + ) as owned: + model: Final = scenario.model() + scenario.cleanups.callback(scenario.delete_user, subject) + marker: Final = f"jwt-allowlist-a1-{uuid.uuid4().hex}" + _observed_bodies() + response: Final = owned.gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": marker}]}, + key=token, + ) + assert response.status_code == 200, response.text + deprecation: Final = eventually( + lambda: owned.log.read_text(), + lambda text: _deprecation_lines(text) != [], + seconds=30, + ) + assert _deprecation_lines(deprecation) != [] + observed: Final = _observed_bodies() + assert any(marker in body for body in observed), observed + + +def test_eddsa_streamed_request_through_openai_sdk_reaches_upstream_and_logs_deprecation( + gateway: Gateway, tmp_path: Path +) -> None: + private_key, jwks_key = _eddsa_keypair() + + def respond(request: Request) -> Reply: + assert request.method == "GET", request + return _jwks_reply_keys([jwks_key]) + + with wire_server(respond) as jwks, gateway.scenario() as scenario: + subject: Final = f"integration-jwt-{uuid.uuid4().hex}" + token: Final = _token(private_key, "EdDSA", subject, kid="ed") + with owned_proxy_process( + gateway, tmp_path, {"JWT_PUBLIC_KEY_URL": jwks.url}, config=_jwt_auth_config(tmp_path), workers=2 + ) as owned: + model: Final = scenario.model() + scenario.cleanups.callback(scenario.delete_user, subject) + marker: Final = f"jwt-allowlist-a3-{uuid.uuid4().hex}" + _observed_bodies() + sdk: Final = OpenAI(base_url=f"{owned.gateway.client.base_url}/v1", api_key=token, timeout=30) + chunks: Final = [ + chunk + for chunk in sdk.chat.completions.create( + model=model, messages=[{"role": "user", "content": marker}], stream=True + ) + ] + assert chunks, "stream produced no chunks" + deprecation: Final = eventually( + lambda: owned.log.read_text(), + lambda text: _deprecation_lines(text) != [], + seconds=30, + ) + assert _deprecation_lines(deprecation) != [] + observed: Final = _observed_bodies() + assert any(marker in body for body in observed), observed + + +def test_rs256_messages_request_through_anthropic_sdk_reaches_upstream(gateway: Gateway, tmp_path: Path) -> None: + private_key, jwks_key = _rsa_keypair() + completion: Final = json.dumps( + { + "id": "msg_allowlist_a4", + "type": "message", + "role": "assistant", + "model": "claude-3-5-haiku-20241022", + "content": [{"type": "text", "text": "a4-reply"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 5, "output_tokens": 2}, + } + ).encode() + + def jwks(request: Request) -> Reply: + assert request.method == "GET", request + return _jwks_reply_keys([jwks_key]) + + def provider(request: Request) -> Reply: + return Reply(body=completion) + + with wire_server(jwks) as keys_server, wire_server(provider) as provider_server, gateway.scenario() as scenario: + subject: Final = f"integration-jwt-{uuid.uuid4().hex}" + token: Final = _token(private_key, "RS256", subject, kid="rsa") + with owned_proxy_process( + gateway, + tmp_path, + {"JWT_PUBLIC_KEY_URL": keys_server.url}, + config=_jwt_auth_config(tmp_path), + workers=2, + ) as owned: + model: Final = scenario.model(model="anthropic/claude-3-5-haiku-latest", api_base=provider_server.url) + scenario.cleanups.callback(scenario.delete_user, subject) + sdk: Final = Anthropic( + base_url=str(owned.gateway.client.base_url).rstrip("/"), + auth_token=token, + timeout=30, + max_retries=0, + ) + message: Final = sdk.messages.create( + model=model, max_tokens=8, messages=[{"role": "user", "content": "a4"}] + ) + text: Final = message.content[0].text + assert "a4-reply" in text, message + received: Final = provider_server.drain() + assert received, "upstream never received the request" + + +def test_rs256_responses_request_through_openai_async_sdk_reaches_upstream(gateway: Gateway, tmp_path: Path) -> None: + private_key, jwks_key = _rsa_keypair() + response_object: Final = json.dumps( + { + "id": "resp_allowlist_a5", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4o-mini", + "output": [ + { + "type": "message", + "id": "msg_a5", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "a5-reply", "annotations": []}], + } + ], + "usage": {"input_tokens": 5, "output_tokens": 2, "total_tokens": 7}, + } + ).encode() + + def jwks(request: Request) -> Reply: + assert request.method == "GET", request + return _jwks_reply_keys([jwks_key]) + + def provider(request: Request) -> Reply: + return Reply(body=response_object) + + with wire_server(jwks) as keys_server, wire_server(provider) as provider_server, gateway.scenario() as scenario: + subject: Final = f"integration-jwt-{uuid.uuid4().hex}" + token: Final = _token(private_key, "RS256", subject, kid="rsa") + with owned_proxy_process( + gateway, + tmp_path, + {"JWT_PUBLIC_KEY_URL": keys_server.url}, + config=_jwt_auth_config(tmp_path), + workers=2, + ) as owned: + model: Final = scenario.model(api_base=provider_server.url) + scenario.cleanups.callback(scenario.delete_user, subject) + + async def call() -> str: + sdk: Final = AsyncOpenAI(base_url=f"{owned.gateway.client.base_url}/v1", api_key=token, timeout=30) + response: Final = await sdk.responses.create(model=model, input="a5") + return response.id + + identity: Final = asyncio.run(call()) + assert identity.startswith("resp_"), identity + received: Final = provider_server.drain() + assert received, "upstream never received the request" + + +def test_es256_signed_token_is_accepted(gateway: Gateway, tmp_path: Path) -> None: + private_key, jwks_key = _ec_keypair() + + def respond(request: Request) -> Reply: + assert request.method == "GET", request + return _jwks_reply_keys([jwks_key]) + + with wire_server(respond) as jwks, gateway.scenario() as scenario: + subject: Final = f"integration-jwt-{uuid.uuid4().hex}" + token: Final = jwt.encode( + {"sub": subject, "iat": int(time.time()), "exp": int(time.time()) + 300}, + private_key, + algorithm="ES256", + headers={"kid": "ec"}, + ) + with owned_proxy_process( + gateway, tmp_path, {"JWT_PUBLIC_KEY_URL": jwks.url}, config=_jwt_auth_config(tmp_path), workers=2 + ) as owned: + model: Final = scenario.model() + scenario.cleanups.callback(scenario.delete_user, subject) + response: Final = owned.gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "jwt alg allowlist es256"}]}, + key=token, + ) + assert response.status_code == 200, response.text + + +def test_hs256_token_against_oct_jwks_key_is_rejected(gateway: Gateway, tmp_path: Path) -> None: + secret: Final = b"integration-hs256-jwt-secret-0123456789abcdef" + oct_key: Final = { + "kty": "oct", + "kid": "sym", + "alg": "HS256", + "k": base64.urlsafe_b64encode(secret).rstrip(b"=").decode(), + } + + def respond(request: Request) -> Reply: + assert request.method == "GET", request + return _jwks_reply_keys([oct_key]) + + with wire_server(respond) as jwks, gateway.scenario() as scenario: + subject: Final = f"integration-jwt-{uuid.uuid4().hex}" + token: Final = jwt.encode( + {"sub": subject, "iat": int(time.time()), "exp": int(time.time()) + 300}, + secret, + algorithm="HS256", + headers={"kid": "sym"}, + ) + with owned_proxy_process( + gateway, tmp_path, {"JWT_PUBLIC_KEY_URL": jwks.url}, config=_jwt_auth_config(tmp_path), workers=2 + ) as owned: + model: Final = scenario.model() + response: Final = owned.gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "jwt alg allowlist hs256"}]}, + key=token, + ) + assert response.status_code == 401, response.text + assert "error" in response.text, response.text + + +def test_token_signed_eddsa_with_rs256_header_is_rejected(gateway: Gateway, tmp_path: Path) -> None: + private_key, jwks_key = _eddsa_keypair() + + def respond(request: Request) -> Reply: + assert request.method == "GET", request + return _jwks_reply_keys([jwks_key]) + + with wire_server(respond) as jwks, gateway.scenario() as scenario: + subject: Final = f"integration-jwt-{uuid.uuid4().hex}" + token: Final = _alg_lied_token(private_key, "ed", subject) + with owned_proxy_process( + gateway, tmp_path, {"JWT_PUBLIC_KEY_URL": jwks.url}, config=_jwt_auth_config(tmp_path), workers=2 + ) as owned: + model: Final = scenario.model() + response: Final = owned.gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "jwt alg allowlist lied"}]}, + key=token, + ) + assert response.status_code == 401, response.text + + +def test_request_without_authorization_header_is_rejected(gateway: Gateway, tmp_path: Path) -> None: + _, jwks_key = _rsa_keypair() + + def respond(request: Request) -> Reply: + assert request.method == "GET", request + return _jwks_reply_keys([jwks_key]) + + with wire_server(respond) as jwks, gateway.scenario() as scenario: + with owned_proxy_process( + gateway, tmp_path, {"JWT_PUBLIC_KEY_URL": jwks.url}, config=_jwt_auth_config(tmp_path), workers=2 + ) as owned: + model: Final = scenario.model() + response: Final = owned.gateway.client.post( + "/v1/chat/completions", + json={"model": model, "messages": [{"role": "user", "content": "jwt alg allowlist noauth"}]}, + ) + assert response.status_code == 401, response.text + + +def test_empty_jwks_document_rejects_tokens(gateway: Gateway, tmp_path: Path) -> None: + private_key, _jwks_key = _rsa_keypair() + + def respond(request: Request) -> Reply: + assert request.method == "GET", request + return _jwks_reply_keys([]) + + with wire_server(respond) as jwks, gateway.scenario() as scenario: + subject: Final = f"integration-jwt-{uuid.uuid4().hex}" + token: Final = jwt.encode( + {"sub": subject, "iat": int(time.time()), "exp": int(time.time()) + 300}, + private_key, + algorithm="RS256", + headers={"kid": "rsa"}, + ) + with owned_proxy_process( + gateway, tmp_path, {"JWT_PUBLIC_KEY_URL": jwks.url}, config=_jwt_auth_config(tmp_path), workers=2 + ) as owned: + model: Final = scenario.model() + response: Final = owned.gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "jwt alg allowlist empty"}]}, + key=token, + ) + assert response.status_code == 401, response.text + + +def test_five_eddsa_requests_warn_at_most_once_per_worker(gateway: Gateway, tmp_path: Path) -> None: + private_key, jwks_key = _eddsa_keypair() + + def respond(request: Request) -> Reply: + assert request.method == "GET", request + return _jwks_reply_keys([jwks_key]) + + with wire_server(respond) as jwks, gateway.scenario() as scenario: + subject: Final = f"integration-jwt-{uuid.uuid4().hex}" + token: Final = _token(private_key, "EdDSA", subject, kid="ed") + with owned_proxy_process( + gateway, tmp_path, {"JWT_PUBLIC_KEY_URL": jwks.url}, config=_jwt_auth_config(tmp_path), workers=2 + ) as owned: + model: Final = scenario.model() + scenario.cleanups.callback(scenario.delete_user, subject) + for _ in range(5): + response: Final = owned.gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "jwt alg allowlist burst"}]}, + key=token, + ) + assert response.status_code == 200, response.text + log_text: Final = eventually( + lambda: owned.log.read_text(), + lambda text: _deprecation_lines(text) != [], + seconds=30, + ) + warnings: Final = _deprecation_lines(log_text) + assert 1 <= len(warnings) <= 2, warnings + + +def test_master_key_request_still_works_on_jwt_enabled_proxy(gateway: Gateway, tmp_path: Path) -> None: + _, jwks_key = _rsa_keypair() + + def respond(request: Request) -> Reply: + assert request.method == "GET", request + return _jwks_reply_keys([jwks_key]) + + with wire_server(respond) as jwks, gateway.scenario() as scenario: + with owned_proxy_process( + gateway, tmp_path, {"JWT_PUBLIC_KEY_URL": jwks.url}, config=_jwt_auth_config(tmp_path), workers=2 + ) as owned: + model: Final = scenario.model() + response: Final = owned.gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "jwt alg allowlist master"}]}, + ) + assert response.status_code == 200, response.text + + +def _spend_row_counts(identifiers: set[str]) -> dict[str, int]: + if not identifiers: + return {} + rows: Final = read_rows( + 'SELECT request_id FROM "LiteLLM_SpendLogs" WHERE request_id = ANY(%s::text[])', + ("{" + ",".join(sorted(identifiers)) + "}",), + ) + counts: Final = {identifier: 0 for identifier in identifiers} + for row in rows: + identifier = str(row["request_id"]) + if identifier in counts: + counts[identifier] += 1 + return counts + + +def test_mixed_jwt_burst_survives_jwks_restart_and_writes_spend_rows(gateway: Gateway, tmp_path: Path) -> None: + rsa_private, rsa_jwks_key = _rsa_keypair() + ed_private, ed_jwks_key = _eddsa_keypair() + port: Final = _free_port() + + def respond(request: Request) -> Reply: + assert request.method == "GET", request + return _jwks_reply_keys([rsa_jwks_key, ed_jwks_key]) + + with gateway.scenario() as scenario: + rs_token: Final = jwt.encode( + {"sub": f"integration-jwt-{uuid.uuid4().hex}", "iat": int(time.time()), "exp": int(time.time()) + 900}, + rsa_private, + algorithm="RS256", + headers={"kid": "rsa"}, + ) + ed_token: Final = jwt.encode( + {"sub": f"integration-jwt-{uuid.uuid4().hex}", "iat": int(time.time()), "exp": int(time.time()) + 900}, + ed_private, + algorithm="EdDSA", + headers={"kid": "ed"}, + ) + with owned_proxy_process( + gateway, + tmp_path, + {"JWT_PUBLIC_KEY_URL": f"http://127.0.0.1:{port}/jwks"}, + config=_jwt_auth_config(tmp_path), + workers=2, + ) as owned: + model: Final = scenario.model() + scenario.cleanups.callback( + scenario.delete_user, str(jwt.decode(rs_token, options={"verify_signature": False})["sub"]) + ) + scenario.cleanups.callback( + scenario.delete_user, str(jwt.decode(ed_token, options={"verify_signature": False})["sub"]) + ) + + def chat(token: str, index: int, stream: bool = False) -> httpx.Response: + with owned.gateway.client.stream( + "POST", + "/v1/chat/completions", + json={ + "model": model, + "messages": [{"role": "user", "content": f"d1-{index}"}], + **({"stream": True} if stream else {}), + }, + headers={"Authorization": f"Bearer {token}"}, + ) as response: + return httpx.Response( + status_code=response.status_code, + headers=dict(response.headers), + content=response.read(), + ) + + with wire_server(respond, port=port): + warmup_rs: Final = chat(rs_token, 0) + warmup_ed: Final = chat(ed_token, 1) + assert warmup_rs.status_code == 200, warmup_rs.text + assert warmup_ed.status_code == 200, warmup_ed.text + + statuses: Final = [] + bodies: Final = [] + with ThreadPoolExecutor(max_workers=15) as pool: + futures = [ + pool.submit(chat, rs_token if index % 2 == 0 else ed_token, index, index % 5 == 0) + for index in range(30) + ] + for future in futures: + result = future.result(timeout=60) + statuses.append(result.status_code) + bodies.append(result.text) + assert all(status in (200, 401) for status in statuses), (statuses, bodies[:3]) + assert any(status == 200 for status in statuses), statuses + + with wire_server(respond, port=port): + fresh_rs: Final = chat(rs_token, 100) + fresh_ed: Final = chat(ed_token, 101) + assert fresh_rs.status_code == 200, fresh_rs.text + assert fresh_ed.status_code == 200, fresh_ed.text + + succeeded: Final = { + json.loads(text)["id"] + for status, text in zip(statuses, bodies) + if status == 200 and "id" in text and text.strip().startswith("{") + } + succeeded.update(json.loads(text)["id"] for text in (fresh_rs.text, fresh_ed.text) if "id" in text) + counts: Final = eventually( + lambda: _spend_row_counts(succeeded), + lambda value: value == {identifier: 1 for identifier in succeeded}, + seconds=70, + ) + assert all(count == 1 for count in counts.values()), counts + + +def test_burst_survives_killed_worker_and_stays_alive(gateway: Gateway, tmp_path: Path) -> None: + rsa_private, rsa_jwks_key = _rsa_keypair() + ed_private, ed_jwks_key = _eddsa_keypair() + + def respond(request: Request) -> Reply: + assert request.method == "GET", request + return _jwks_reply_keys([rsa_jwks_key, ed_jwks_key]) + + with wire_server(respond) as jwks, gateway.scenario() as scenario: + rs_token: Final = jwt.encode( + {"sub": f"integration-jwt-{uuid.uuid4().hex}", "iat": int(time.time()), "exp": int(time.time()) + 900}, + rsa_private, + algorithm="RS256", + headers={"kid": "rsa"}, + ) + ed_token: Final = jwt.encode( + {"sub": f"integration-jwt-{uuid.uuid4().hex}", "iat": int(time.time()), "exp": int(time.time()) + 900}, + ed_private, + algorithm="EdDSA", + headers={"kid": "ed"}, + ) + with owned_proxy_process( + gateway, tmp_path, {"JWT_PUBLIC_KEY_URL": jwks.url}, config=_jwt_auth_config(tmp_path), workers=2 + ) as owned: + model: Final = scenario.model() + scenario.cleanups.callback( + scenario.delete_user, str(jwt.decode(rs_token, options={"verify_signature": False})["sub"]) + ) + scenario.cleanups.callback( + scenario.delete_user, str(jwt.decode(ed_token, options={"verify_signature": False})["sub"]) + ) + + def chat(token: str, index: int) -> httpx.Response: + return owned.gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": f"d2-{index}"}]}, + key=token, + ) + + with ThreadPoolExecutor(max_workers=12) as pool: + futures = [pool.submit(chat, rs_token if index % 2 == 0 else ed_token, index) for index in range(24)] + children: Final = psutil.Process(owned.process.pid).children(recursive=True) + assert children, "owned proxy reported no worker children" + children[0].kill() + results: Final = [] + for future in futures: + try: + result = future.result(timeout=60) + results.append(result.status_code) + except httpx.TransportError: + results.append(0) + assert all(status in (200, 401, 0) for status in results), results + assert any(status == 200 for status in results), results + liveliness: Final = eventually( + lambda: owned.gateway.client.get("/health/liveliness"), + lambda response: response.status_code == 200, + seconds=45, + ) + assert liveliness.status_code == 200, liveliness.text diff --git a/tests/integration/mcp/test_mcp_client_assertion_signing_alg.py b/tests/integration/mcp/test_mcp_client_assertion_signing_alg.py index c2c32f9866b..d00c51d1156 100644 --- a/tests/integration/mcp/test_mcp_client_assertion_signing_alg.py +++ b/tests/integration/mcp/test_mcp_client_assertion_signing_alg.py @@ -1,14 +1,18 @@ import json import os import uuid +from pathlib import Path from typing import Final +import httpx import psycopg +import pytest from cryptography.hazmat.primitives import serialization from cryptography.hazmat.primitives.asymmetric import rsa from integration._support.client import Gateway, eventually from integration._support.database import read_rows from integration._support.mcp import forget_mcp, mcp_peer, register_mcp +from integration._support.process import owned_proxy_process def _pem_private_key() -> str: @@ -110,11 +114,164 @@ def test_server_row_with_stale_client_assertion_signing_alg_still_loads(gateway: register_mcp(scenario, peer, "staletrigger" + uuid.uuid4().hex[:8]) listed: Final = eventually( lambda: gateway.request("GET", "/v1/mcp/server"), - lambda response: response.status_code == 200 - and any(server.get("server_id") == identity for server in response.json()), + lambda response: ( + response.status_code == 200 + and any(server.get("server_id") == identity for server in response.json()) + ), seconds=30, ) assert listed.status_code == 200, listed.text finally: with psycopg.connect(os.environ["DATABASE_URL"]) as connection: connection.execute('DELETE FROM "LiteLLM_MCPServerTable" WHERE server_id = %s', (identity,)) + + +APPROVED_ALGS: Final = ("RS256", "RS384", "RS512", "PS256", "PS384", "PS512", "ES256", "ES384", "ES512") + + +def _post_server(gateway: Gateway, name: str, credentials: object, omit_alg: bool = False) -> httpx.Response: + blob: Final = {} if omit_alg else {"client_assertion_signing_alg": credentials} + return gateway.request( + "POST", + "/v1/mcp/server", + { + "server_name": name, + "alias": name, + "url": "https://mcp.integration.invalid", + "transport": "http", + "credentials": blob, + }, + ) + + +def test_post_with_lowercase_hs256_is_rejected_naming_the_field(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + created: Final = _post_server(gateway, "b1hs" + uuid.uuid4().hex[:8], "hs256") + if created.status_code == 201: + scenario.cleanups.callback(forget_mcp, gateway, str(created.json()["server_id"])) + assert created.status_code == 422, created.text + assert "client_assertion_signing_alg" in created.text, created.text + for algorithm in APPROVED_ALGS: + assert algorithm in created.text, created.text + + +def test_put_with_eddsa_on_existing_server_is_rejected(gateway: Gateway) -> None: + with mcp_peer() as peer, gateway.scenario() as scenario: + identity: Final = register_mcp(scenario, peer, "b2ed" + uuid.uuid4().hex[:8]) + edited: Final = gateway.request( + "PUT", + "/v1/mcp/server", + {"server_id": identity, "credentials": {"client_assertion_signing_alg": "EdDSA"}}, + ) + assert edited.status_code == 422, edited.text + assert "client_assertion_signing_alg" in edited.text, edited.text + + +@pytest.mark.parametrize("algorithm", APPROVED_ALGS) +def test_each_approved_client_assertion_signing_alg_is_accepted_and_stored(gateway: Gateway, algorithm: str) -> None: + with gateway.scenario() as scenario: + name: Final = "b3" + algorithm.lower() + uuid.uuid4().hex[:6] + created: Final = _post_server(gateway, name, algorithm) + assert created.status_code == 201, created.text + identity: Final = str(created.json()["server_id"]) + scenario.cleanups.callback(forget_mcp, gateway, identity) + listed: Final = gateway.request("GET", "/v1/mcp/server") + assert listed.status_code == 200, listed.text + assert any(server.get("server_id") == identity for server in listed.json()), listed.text + assert _stored_client_assertion_signing_alg(identity) == algorithm + + +@pytest.mark.parametrize( + ("label", "value", "expected"), + ( + ("empty", "", 422), + ("five-kb", "x" * 5120, 422), + ("integer", 7, 422), + ("list", ["RS256"], 422), + ("null", None, 201), + ), +) +def test_client_assertion_signing_alg_payload_variants( + gateway: Gateway, label: str, value: object, expected: int +) -> None: + with gateway.scenario() as scenario: + created: Final = _post_server(gateway, f"b{label}" + uuid.uuid4().hex[:8], value) + if created.status_code == 201: + scenario.cleanups.callback(forget_mcp, gateway, str(created.json()["server_id"])) + assert created.status_code == expected, created.text + if expected == 422 and label in ("empty", "five-kb"): + assert "client_assertion_signing_alg" in created.text, created.text + + +def test_server_post_without_credentials_alg_key_is_accepted(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + created: Final = _post_server(gateway, "b9absent" + uuid.uuid4().hex[:8], None, omit_alg=True) + assert created.status_code == 201, created.text + scenario.cleanups.callback(forget_mcp, gateway, str(created.json()["server_id"])) + + +def test_unauthenticated_server_post_with_hs256_is_rejected(gateway: Gateway) -> None: + response: Final = gateway.client.post( + "/v1/mcp/server", + json={ + "server_name": "b10unauth" + uuid.uuid4().hex[:8], + "credentials": {"client_assertion_signing_alg": "hs256"}, + }, + ) + assert response.status_code == 401, response.text + + +def _seeded_row_loads_with_fallback_warning(gateway: Gateway, tmp_path: Path, algorithm: str) -> None: + identity: Final = "dualread" + uuid.uuid4().hex[:8] + name: Final = "dualread" + uuid.uuid4().hex[:8] + with psycopg.connect(os.environ["DATABASE_URL"]) as connection: + connection.execute( + 'INSERT INTO "LiteLLM_MCPServerTable" (server_id, server_name, url, transport, credentials,' + " created_at, updated_at) VALUES (%s, %s, %s, %s, %s::jsonb, NOW(), NOW())", + ( + identity, + name, + "https://mcp.integration.invalid", + "http", + json.dumps({"client_assertion_signing_alg": algorithm}), + ), + ) + try: + with owned_proxy_process(gateway, tmp_path, {}) as owned: + listed: Final = eventually( + lambda: owned.gateway.request("GET", "/v1/mcp/server"), + lambda response: ( + response.status_code == 200 + and any(server.get("server_id") == identity for server in response.json()) + ), + seconds=30, + ) + assert listed.status_code == 200, listed.text + log_text: Final = eventually( + lambda: owned.log.read_text(), + lambda text: name in text and "not an approved algorithm" in text and "using RS256" in text, + seconds=30, + ) + assert name in log_text and "not an approved algorithm" in log_text and "using RS256" in log_text + finally: + with psycopg.connect(os.environ["DATABASE_URL"]) as connection: + connection.execute('DELETE FROM "LiteLLM_MCPServerTable" WHERE server_id = %s', (identity,)) + + +def test_seeded_hs256_row_loads_with_rs256_fallback_warning(gateway: Gateway, tmp_path: Path) -> None: + _seeded_row_loads_with_fallback_warning(gateway, tmp_path, "HS256") + + +def test_seeded_eddsa_row_loads_with_rs256_fallback_warning(gateway: Gateway, tmp_path: Path) -> None: + _seeded_row_loads_with_fallback_warning(gateway, tmp_path, "EdDSA") + + +def test_chat_still_works_after_alg_rejection(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + created: Final = _post_server(gateway, "b13chat" + uuid.uuid4().hex[:8], "hs256") + assert created.status_code in (201, 422), created.text + if created.status_code == 201: + scenario.cleanups.callback(forget_mcp, gateway, str(created.json()["server_id"])) + model: Final = scenario.model() + body: Final = gateway.chat(model) + assert body["id"], body diff --git a/tests/integration/mcp/test_mcp_jwt_signer_jwks_allowlist.py b/tests/integration/mcp/test_mcp_jwt_signer_jwks_allowlist.py index 8af03fe0663..2d1d48a7a95 100644 --- a/tests/integration/mcp/test_mcp_jwt_signer_jwks_allowlist.py +++ b/tests/integration/mcp/test_mcp_jwt_signer_jwks_allowlist.py @@ -1,22 +1,31 @@ import base64 import json +import os +import signal +import socket +import subprocess +import sys import time import uuid -from collections.abc import Callable, Iterator, Mapping +from collections.abc import Callable, Generator, Mapping +from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager +from pathlib import Path from typing import Final import httpx import jwt import pytest from cryptography.hazmat.primitives.asymmetric import ed25519, rsa -from integration._support.client import Gateway +from integration._support.client import Gateway, eventually from integration._support.mcp import mcp_peer, register_mcp, tool_calls, tool_names +from integration._support.process import owned_proxy_process from integration._support.wire import Reply, Request, wire_server +from jwt.algorithms import OKPAlgorithm, RSAAlgorithm @contextmanager -def _jwt_signer_guardrail(gateway: Gateway, discovery_uri: str) -> Iterator[None]: +def _jwt_signer_guardrail(gateway: Gateway, discovery_uri: str) -> Generator[None]: created: Final = gateway.client.post( "/guardrails", headers={"x-litellm-api-key": gateway.key}, @@ -42,7 +51,7 @@ def _jwt_signer_guardrail(gateway: Gateway, discovery_uri: str) -> Iterator[None @contextmanager -def _idp_server(jwks_keys: list[Mapping[str, object]]) -> Iterator[str]: +def _idp_server(jwks_keys: list[Mapping[str, object]]) -> Generator[str]: holder: Final = {"url": ""} def respond(request: Request) -> Reply: @@ -76,7 +85,7 @@ def _hs256_key_and_token() -> tuple[dict[str, object], str]: def _eddsa_key_and_token() -> tuple[dict[str, object], str]: private_key: Final = ed25519.Ed25519PrivateKey.generate() - public_jwk: Final = json.loads(jwt.algorithms.OKPAlgorithm.to_jwk(private_key.public_key())) + public_jwk: Final = json.loads(OKPAlgorithm.to_jwk(private_key.public_key())) key: Final = {**public_jwk, "alg": "EdDSA", "kid": "ed"} token: Final = jwt.encode(_claims(), private_key, algorithm="EdDSA", headers={"kid": "ed"}) return key, token @@ -84,7 +93,7 @@ def _eddsa_key_and_token() -> tuple[dict[str, object], str]: def _rs256_key_and_token() -> tuple[dict[str, object], str]: private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) - public_jwk: Final = json.loads(jwt.algorithms.RSAAlgorithm.to_jwk(private_key.public_key())) + public_jwk: Final = json.loads(RSAAlgorithm.to_jwk(private_key.public_key())) key: Final = {**public_jwk, "alg": "RS256", "kid": "rsa"} token: Final = jwt.encode(_claims(), private_key, algorithm="RS256", headers={"kid": "rsa"}) return key, token @@ -131,3 +140,142 @@ def test_jwks_rs256_key_still_verifies_the_incoming_token(gateway: Gateway) -> N assert response.status_code == 200, response.text assert response.json()["content"][0]["text"] == "9", response.text assert len(tool_calls(peer.drain())) == 1, "accepted call never reached the peer" + + +def _rs256_keypair() -> tuple[rsa.RSAPrivateKey, dict[str, object]]: + private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + return private_key, json.loads(RSAAlgorithm.to_jwk(private_key.public_key())) + + +def _rs256_key_and_token_no_alg() -> tuple[dict[str, object], str]: + private_key, public_jwk = _rs256_keypair() + key: Final = {**public_jwk, "kid": "rsa"} + token: Final = jwt.encode(_claims(), private_key, algorithm="RS256", headers={"kid": "rsa"}) + return key, token + + +def _signer_call(gateway: Gateway, key: str, identity: str, name: str, bearer: str) -> tuple[int, str]: + try: + response: Final = _call_with_bearer(gateway, key, identity, name, {"a": 4, "b": 5}, bearer) + return response.status_code, response.text + except httpx.HTTPError as error: + return 0, repr(error) + + +def test_jwks_rs256_key_without_alg_field_still_verifies_the_token(gateway: Gateway) -> None: + jwks_key, token = _rs256_key_and_token_no_alg() + with _idp_server([jwks_key]) as discovery_uri, _jwt_signer_guardrail(gateway, discovery_uri): + with mcp_peer() as peer, gateway.scenario() as scenario: + identity: Final = register_mcp(scenario, peer, "jwksc4" + uuid.uuid4().hex[:8]) + key: Final = scenario.key(object_permission={"mcp_servers": [identity]}) + name: Final = tool_names(gateway, key, identity)["add"] + peer.drain() + response: Final = _call_with_bearer(gateway, key, identity, name, {"a": 4, "b": 5}, token) + assert response.status_code == 200, response.text + assert response.json()["content"][0]["text"] == "9", response.text + assert len(tool_calls(peer.drain())) == 1, "accepted call never reached the peer" + + +def test_rs256_token_verifies_when_jwks_also_carries_a_non_approved_key(gateway: Gateway) -> None: + okp_key, _eddsa_token = _eddsa_key_and_token() + jwks_key, token = _rs256_key_and_token() + with _idp_server([okp_key, jwks_key]) as discovery_uri, _jwt_signer_guardrail(gateway, discovery_uri): + with mcp_peer() as peer, gateway.scenario() as scenario: + identity: Final = register_mcp(scenario, peer, "jwksc5" + uuid.uuid4().hex[:8]) + key: Final = scenario.key(object_permission={"mcp_servers": [identity]}) + name: Final = tool_names(gateway, key, identity)["add"] + peer.drain() + response: Final = _call_with_bearer(gateway, key, identity, name, {"a": 4, "b": 5}, token) + assert response.status_code == 200, response.text + assert response.json()["content"][0]["text"] == "9", response.text + assert len(tool_calls(peer.drain())) == 1, "accepted call never reached the peer" + + +@pytest.mark.parametrize( + ("label", "jwks_reply"), + ( + ("jwks-404", Reply(status=404, body=b'{"error": "missing"}')), + ("jwks-not-json", Reply(body=b"this is not json")), + ), + ids=("jwks-404", "jwks-not-json"), +) +def test_unusable_jwks_document_rejects_the_incoming_token(gateway: Gateway, label: str, jwks_reply: Reply) -> None: + _, token = _rs256_key_and_token() + + def respond(request: Request) -> Reply: + if request.target == "/.well-known/openid-configuration": + return Reply(body=json.dumps({"jwks_uri": holder["url"] + "/.well-known/jwks.json"}).encode()) + return jwks_reply + + holder: Final = {"url": ""} + with wire_server(respond) as idp: + holder["url"] = idp.url + with _jwt_signer_guardrail(gateway, idp.url + "/.well-known/openid-configuration"): + with mcp_peer() as peer, gateway.scenario() as scenario: + identity: Final = register_mcp(scenario, peer, "jwks" + uuid.uuid4().hex[:8]) + key: Final = scenario.key(object_permission={"mcp_servers": [identity]}) + name: Final = tool_names(gateway, key, identity)["add"] + peer.drain() + status, text = _signer_call(gateway, key, identity, name, token) + assert status == 401, (status, text) + assert "incoming token verification failed" in text, text + assert tool_calls(peer.drain()) == (), "rejected call reached the peer" + alive: Final = gateway.client.get("/health/liveliness") + assert alive.status_code == 200, alive.text + + +def test_signer_jwks_pause_rejects_then_recovers(gateway: Gateway, tmp_path: Path) -> None: + private_key, public_jwk = _rs256_keypair() + jwks_dir: Final = tmp_path / "jwks" / ".well-known" + jwks_dir.mkdir(parents=True) + port: Final = _reserve_port() + (jwks_dir / "jwks.json").write_text(json.dumps({"keys": [{**public_jwk, "kid": "rsa", "alg": "RS256"}]})) + (jwks_dir / "openid-configuration").write_text( + json.dumps({"jwks_uri": f"http://127.0.0.1:{port}/.well-known/jwks.json"}) + ) + token: Final = jwt.encode(_claims(), private_key, algorithm="RS256", headers={"kid": "rsa"}) + server: Final = subprocess.Popen( + [sys.executable, "-m", "http.server", str(port), "--bind", "127.0.0.1", "--directory", str(tmp_path / "jwks")], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + try: + with owned_proxy_process(gateway, tmp_path, {}) as owned: + discovery: Final = f"http://127.0.0.1:{port}/.well-known/openid-configuration" + with _jwt_signer_guardrail(owned.gateway, discovery): + with mcp_peer() as peer, owned.gateway.scenario() as scenario: + identity: Final = register_mcp(scenario, peer, "jwksd3" + uuid.uuid4().hex[:8]) + key: Final = scenario.key(object_permission={"mcp_servers": [identity]}) + name: Final = tool_names(owned.gateway, key, identity)["add"] + peer.drain() + os.kill(server.pid, signal.SIGSTOP) + outcomes: Final = [] + with ThreadPoolExecutor(max_workers=8) as pool: + futures = [ + pool.submit(_signer_call, owned.gateway, key, identity, name, token) for _ in range(8) + ] + for future in futures: + outcomes.append(future.result(timeout=80)) + assert all(status != 200 for status, _text in outcomes), outcomes + assert any(status in (0, 401, 500) for status, _text in outcomes), outcomes + os.kill(server.pid, signal.SIGCONT) + recovered: Final = eventually( + lambda: _signer_call(owned.gateway, key, identity, name, token), + lambda outcome: outcome[0] == 200, + seconds=70, + ) + assert recovered[0] == 200, recovered + assert tool_calls(peer.drain()) != (), "resumed call never reached the peer" + finally: + try: + os.kill(server.pid, signal.SIGCONT) + except ProcessLookupError: + pass + server.terminate() + server.wait(timeout=10) + + +def _reserve_port() -> int: + with socket.socket() as reserve: + reserve.bind(("127.0.0.1", 0)) + return reserve.getsockname()[1]