diff --git a/tests/integration/authorization/test_jwt_algorithm_allowlist.py b/tests/integration/authorization/test_jwt_algorithm_allowlist.py new file mode 100644 index 00000000000..210d40a0d6a --- /dev/null +++ b/tests/integration/authorization/test_jwt_algorithm_allowlist.py @@ -0,0 +1,99 @@ +import json +import time +import uuid +from pathlib import Path +from typing import Final + +import jwt +import yaml +from cryptography.hazmat.primitives.asymmetric import ed25519, rsa + +from tests.integration._support.client import Gateway, eventually +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" + + +def _jwks_reply(public_jwk: str) -> Reply: + return Reply(body=json.dumps({"keys": [{**json.loads(public_jwk), "kid": KEY_ID}]}).encode()) + + +def _jwt_auth_config(tmp_path: Path) -> Path: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["general_settings"] = { + **config["general_settings"], + "enable_jwt_auth": True, + "litellm_jwtauth": {"user_id_jwt_field": "sub", "user_id_upsert": True}, + } + path: Final = tmp_path / "jwt_allowlist.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +def _token(key: object, algorithm: str, subject: str) -> 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}, + ) + + +def test_eddsa_signed_token_is_accepted_and_logged_as_deprecated_outside_fips_mode( + gateway: Gateway, tmp_path: Path +) -> None: + private_key: Final = ed25519.Ed25519PrivateKey.generate() + public_jwk: Final = jwt.algorithms.OKPAlgorithm.to_jwk(private_key.public_key()) + + def respond(request: Request) -> Reply: + assert request.method == "GET", request + return _jwks_reply(public_jwk) + + 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) + with owned_proxy_process( + gateway, tmp_path, {"JWT_PUBLIC_KEY_URL": jwks.url}, config=_jwt_auth_config(tmp_path) + ) 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 control"}]}, + key=token, + ) + assert response.status_code == 200, response.text + deprecation: Final = eventually( + lambda: owned.log.read_text(), + lambda text: "EdDSA" in text and "deprecated" in text and "LITELLM_FIPS_MODE" in text, + seconds=30, + ) + assert "EdDSA" in deprecation and "deprecated" in deprecation and "LITELLM_FIPS_MODE" in deprecation + + +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()) + + def respond(request: Request) -> Reply: + assert request.method == "GET", request + return _jwks_reply(public_jwk) + + with wire_server(respond) as jwks, gateway.scenario() as scenario: + subject: Final = f"integration-jwt-{uuid.uuid4().hex}" + token: Final = _token(private_key, "RS256", subject) + with owned_proxy_process( + gateway, tmp_path, {"JWT_PUBLIC_KEY_URL": jwks.url}, config=_jwt_auth_config(tmp_path) + ) 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 control"}]}, + key=token, + ) + assert response.status_code == 200, response.text + assert "EdDSA" not in owned.log.read_text(), owned.log.read_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 new file mode 100644 index 00000000000..883c335b71d --- /dev/null +++ b/tests/integration/mcp/test_mcp_client_assertion_signing_alg.py @@ -0,0 +1,87 @@ +import json +import uuid +from typing import Final + +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric import rsa +from integration._support.client import Gateway +from integration._support.database import read_rows +from integration._support.mcp import forget_mcp, mcp_peer, register_mcp + + +def _pem_private_key() -> str: + return ( + rsa.generate_private_key(public_exponent=65537, key_size=2048) + .private_bytes( + serialization.Encoding.PEM, + serialization.PrivateFormat.PKCS8, + serialization.NoEncryption(), + ) + .decode() + ) + + +def _stored_client_assertion_signing_alg(identity: str) -> str: + rows: Final = read_rows('SELECT credentials FROM "LiteLLM_MCPServerTable" WHERE server_id = %s', (identity,)) + assert len(rows) == 1, rows + credentials: Final = rows[0]["credentials"] + if credentials is None: + return "RS256" + blob: Final = credentials if isinstance(credentials, dict) else json.loads(credentials) + return str(blob.get("client_assertion_signing_alg") or "RS256") + + +def test_non_approved_client_assertion_signing_alg_is_rejected_on_create_and_update(gateway: Gateway) -> None: + with mcp_peer() as peer, gateway.scenario() as scenario: + alias: Final = "alg" + uuid.uuid4().hex[:8] + created: Final = gateway.request( + "POST", + "/v1/mcp/server", + { + "server_name": alias, + "alias": alias, + **peer.registration(), + "token_exchange_endpoint": "https://idp.integration.invalid/oauth2/token", + "credentials": { + "client_private_key": _pem_private_key(), + "client_assertion_signing_alg": "HS256", + }, + }, + ) + if created.status_code == 201: + scenario.cleanups.callback(forget_mcp, gateway, str(created.json()["server_id"])) + assert created.status_code in (400, 422), created.text + assert "client_assertion_signing_alg" in created.text, created.text + + identity: Final = register_mcp(scenario, peer, alias + "v") + edited: Final = gateway.request( + "PUT", + "/v1/mcp/server", + {"server_id": identity, "credentials": {"client_assertion_signing_alg": "EdDSA"}}, + ) + assert edited.status_code in (400, 422), edited.text + assert "client_assertion_signing_alg" in edited.text, edited.text + assert _stored_client_assertion_signing_alg(identity) == "RS256" + + +def test_approved_client_assertion_signing_alg_round_trips(gateway: Gateway) -> None: + with mcp_peer() as peer, gateway.scenario() as scenario: + alias: Final = "algok" + uuid.uuid4().hex[:8] + created: Final = gateway.request( + "POST", + "/v1/mcp/server", + { + "server_name": alias, + "alias": alias, + **peer.registration(), + "token_exchange_endpoint": "https://idp.integration.invalid/oauth2/token", + "credentials": { + "client_private_key": _pem_private_key(), + "client_assertion_signing_alg": "ES256", + }, + }, + ) + assert created.status_code == 201, created.text + identity: Final = str(created.json()["server_id"]) + scenario.cleanups.callback(forget_mcp, gateway, identity) + assert _stored_client_assertion_signing_alg(identity) == "ES256" diff --git a/tests/integration/mcp/test_mcp_jwt_signer_jwks_allowlist.py b/tests/integration/mcp/test_mcp_jwt_signer_jwks_allowlist.py new file mode 100644 index 00000000000..8af03fe0663 --- /dev/null +++ b/tests/integration/mcp/test_mcp_jwt_signer_jwks_allowlist.py @@ -0,0 +1,133 @@ +import base64 +import json +import time +import uuid +from collections.abc import Callable, Iterator, Mapping +from contextlib import contextmanager +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.mcp import mcp_peer, register_mcp, tool_calls, tool_names +from integration._support.wire import Reply, Request, wire_server + + +@contextmanager +def _jwt_signer_guardrail(gateway: Gateway, discovery_uri: str) -> Iterator[None]: + created: Final = gateway.client.post( + "/guardrails", + headers={"x-litellm-api-key": gateway.key}, + json={ + "guardrail": { + "guardrail_name": "signer" + uuid.uuid4().hex[:8], + "litellm_params": { + "guardrail": "mcp_jwt_signer", + "mode": "pre_mcp_call", + "default_on": True, + "access_token_discovery_uri": discovery_uri, + }, + } + }, + ) + assert created.status_code == 200, created.text + identity: Final = str(created.json()["guardrail_id"]) + try: + yield + finally: + deleted: Final = gateway.client.delete(f"/guardrails/{identity}", headers={"x-litellm-api-key": gateway.key}) + assert deleted.status_code == 200, deleted.text + + +@contextmanager +def _idp_server(jwks_keys: list[Mapping[str, object]]) -> Iterator[str]: + holder: Final = {"url": ""} + + 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()) + if request.target == "/.well-known/jwks.json": + return Reply(body=json.dumps({"keys": list(jwks_keys)}).encode()) + return Reply(status=404) + + with wire_server(respond) as idp: + holder["url"] = idp.url + yield idp.url + "/.well-known/openid-configuration" + + +def _claims() -> dict[str, object]: + now: Final = int(time.time()) + return {"sub": "integration-mcp-user", "iat": now, "exp": now + 300} + + +def _hs256_key_and_token() -> tuple[dict[str, object], str]: + secret: Final = b"integration-hs256-client-secret-0123456789abcdef" + key: Final = { + "kty": "oct", + "alg": "HS256", + "kid": "sym", + "k": base64.urlsafe_b64encode(secret).rstrip(b"=").decode(), + } + token: Final = jwt.encode(_claims(), secret, algorithm="HS256", headers={"kid": "sym"}) + return key, token + + +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())) + key: Final = {**public_jwk, "alg": "EdDSA", "kid": "ed"} + token: Final = jwt.encode(_claims(), private_key, algorithm="EdDSA", headers={"kid": "ed"}) + return key, token + + +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())) + key: Final = {**public_jwk, "alg": "RS256", "kid": "rsa"} + token: Final = jwt.encode(_claims(), private_key, algorithm="RS256", headers={"kid": "rsa"}) + return key, token + + +def _call_with_bearer( + gateway: Gateway, key: str, identity: str, name: str, arguments: dict[str, object], bearer: str +) -> httpx.Response: + return gateway.client.post( + "/mcp-rest/tools/call", + headers={"x-litellm-api-key": key, "Authorization": f"Bearer {bearer}"}, + json={"server_id": identity, "name": name, "arguments": arguments}, + ) + + +@pytest.mark.parametrize("key_and_token", (_hs256_key_and_token, _eddsa_key_and_token), ids=("oct-HS256", "OKP-EdDSA")) +def test_jwks_key_with_non_approved_alg_cannot_verify_the_incoming_token( + gateway: Gateway, key_and_token: Callable[[], tuple[dict[str, object], str]] +) -> None: + jwks_key, token = key_and_token() + with _idp_server([jwks_key]) as discovery_uri, _jwt_signer_guardrail(gateway, discovery_uri): + with mcp_peer() as peer, gateway.scenario() as scenario: + alias: Final = "jwksalg" + uuid.uuid4().hex[:8] + identity: Final = register_mcp(scenario, peer, alias) + 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 == 401, response.text + assert "incoming token verification failed" in response.text, response.text + assert tool_calls(peer.drain()) == (), "rejected call reached the peer" + + +def test_jwks_rs256_key_still_verifies_the_incoming_token(gateway: Gateway) -> None: + jwks_key, token = _rs256_key_and_token() + with _idp_server([jwks_key]) as discovery_uri, _jwt_signer_guardrail(gateway, discovery_uri): + with mcp_peer() as peer, gateway.scenario() as scenario: + alias: Final = "jwksok" + uuid.uuid4().hex[:8] + identity: Final = register_mcp(scenario, peer, alias) + 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"