mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-12 03:48:56 +00:00
test(integration): audit cells for JWT and MCP algorithm allowlists (LIT-8429)
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yucheng
parent
a7c83a14e2
commit
93939f2a53
3 changed files with 907 additions and 14 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue