test(auth): integration cells for JWT algorithm allowlists (LIT-8429)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

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()

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"

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"