Run transport security tests in process instead of over sockets (#2764)
This commit is contained in:
Max
GitHub
parent
ed6adeee7e
commit
b3025f93b4
1 changed files with 212 additions and 454 deletions
|
|
@ -1,15 +1,12 @@
|
|||
"""Tests for SSE server request validation."""
|
||||
|
||||
import logging
|
||||
import multiprocessing
|
||||
import re
|
||||
import socket
|
||||
|
||||
import anyio
|
||||
import httpx
|
||||
import pytest
|
||||
import sse_starlette.sse
|
||||
import uvicorn
|
||||
from starlette.applications import Starlette
|
||||
from starlette.requests import Request
|
||||
from starlette.responses import Response
|
||||
|
|
@ -23,12 +20,16 @@ from mcp.server.sse import SseServerTransport
|
|||
from mcp.server.transport_security import TransportSecuritySettings
|
||||
from mcp.shared._stream_protocols import WriteStream
|
||||
from mcp.shared.message import SessionMessage
|
||||
from mcp.types import JSONRPCRequest, JSONRPCResponse, Tool
|
||||
from tests.test_helpers import wait_for_server
|
||||
from mcp.types import JSONRPCRequest, JSONRPCResponse
|
||||
from tests.interaction.transports import StreamingASGITransport
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
SERVER_NAME = "test_sse_security_server"
|
||||
|
||||
# The in-process app is mounted at this origin purely so URLs are well-formed and the default
|
||||
# Host header is a localhost form; nothing listens here.
|
||||
BASE_URL = "http://127.0.0.1:8000"
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def reset_sse_starlette_exit_event() -> None:
|
||||
|
|
@ -39,275 +40,161 @@ def reset_sse_starlette_exit_event() -> None:
|
|||
app_status.should_exit_event = None
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def server_port() -> int:
|
||||
with socket.socket() as s:
|
||||
s.bind(("127.0.0.1", 0))
|
||||
return s.getsockname()[1]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def server_url(server_port: int) -> str: # pragma: no cover
|
||||
return f"http://127.0.0.1:{server_port}"
|
||||
|
||||
|
||||
class SecurityTestServer(Server): # pragma: no cover
|
||||
def __init__(self):
|
||||
super().__init__(SERVER_NAME)
|
||||
|
||||
async def on_list_tools(self) -> list[Tool]:
|
||||
return []
|
||||
|
||||
|
||||
def run_server_with_settings(port: int, security_settings: TransportSecuritySettings | None = None): # pragma: no cover
|
||||
"""Run the SSE server with specified security settings."""
|
||||
app = SecurityTestServer()
|
||||
def sse_security_client(security_settings: TransportSecuritySettings | None = None) -> httpx.AsyncClient:
|
||||
"""An httpx client whose requests are served in process by an SSE app with the given settings."""
|
||||
server = Server(SERVER_NAME)
|
||||
sse_transport = SseServerTransport("/messages/", security_settings)
|
||||
|
||||
async def handle_sse(request: Request):
|
||||
async def handle_sse(request: Request) -> Response:
|
||||
try:
|
||||
async with sse_transport.connect_sse(request.scope, request.receive, request._send) as streams:
|
||||
if streams:
|
||||
await app.run(streams[0], streams[1], app.create_initialization_options())
|
||||
async with sse_transport.connect_sse(request.scope, request.receive, request._send) as (read, write):
|
||||
await server.run(read, write, server.create_initialization_options())
|
||||
except ValueError as e:
|
||||
# Validation error was already handled inside connect_sse
|
||||
# Validation error was already handled inside connect_sse, which sent the rejection
|
||||
# response itself; its non-empty body checkpoints, so the test reads the rejection
|
||||
# status before the trailing Response() below sends a second response start.
|
||||
logger.debug(f"SSE connection failed validation: {e}")
|
||||
return Response()
|
||||
|
||||
routes = [
|
||||
Route("/sse", endpoint=handle_sse),
|
||||
Mount("/messages/", app=sse_transport.handle_post_message),
|
||||
]
|
||||
|
||||
starlette_app = Starlette(routes=routes)
|
||||
uvicorn.run(starlette_app, host="127.0.0.1", port=port, log_level="error")
|
||||
|
||||
|
||||
def start_server_process(port: int, security_settings: TransportSecuritySettings | None = None):
|
||||
"""Start server in a separate process."""
|
||||
process = multiprocessing.Process(target=run_server_with_settings, args=(port, security_settings))
|
||||
process.start()
|
||||
# Wait for server to be ready to accept connections
|
||||
wait_for_server(port)
|
||||
return process
|
||||
app = Starlette(
|
||||
routes=[
|
||||
Route("/sse", endpoint=handle_sse),
|
||||
Mount("/messages/", app=sse_transport.handle_post_message),
|
||||
]
|
||||
)
|
||||
# The SSE GET runs until it observes a disconnect, so the bridge must let the application
|
||||
# drain on close rather than cancelling it.
|
||||
transport = StreamingASGITransport(app, cancel_on_close=False)
|
||||
return httpx.AsyncClient(transport=transport, base_url=BASE_URL)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_sse_security_default_settings(server_port: int):
|
||||
"""Test SSE with default security settings (protection disabled)."""
|
||||
process = start_server_process(server_port)
|
||||
async def test_sse_security_default_settings() -> None:
|
||||
"""With default security settings (protection disabled), any Host and Origin connect."""
|
||||
headers = {"Host": "evil.com", "Origin": "http://evil.com"}
|
||||
|
||||
try:
|
||||
headers = {"Host": "evil.com", "Origin": "http://evil.com"}
|
||||
|
||||
async with httpx.AsyncClient(timeout=5.0) as client:
|
||||
async with client.stream("GET", f"http://127.0.0.1:{server_port}/sse", headers=headers) as response:
|
||||
assert response.status_code == 200
|
||||
finally:
|
||||
process.terminate()
|
||||
process.join()
|
||||
async with sse_security_client() as client:
|
||||
async with client.stream("GET", "/sse", headers=headers) as response:
|
||||
assert response.status_code == 200
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_sse_security_invalid_host_header(server_port: int):
|
||||
"""Test SSE with invalid Host header."""
|
||||
# Enable security by providing settings with an empty allowed_hosts list
|
||||
async def test_sse_security_invalid_host_header() -> None:
|
||||
"""A Host header outside allowed_hosts is rejected with 421."""
|
||||
security_settings = TransportSecuritySettings(enable_dns_rebinding_protection=True, allowed_hosts=["example.com"])
|
||||
process = start_server_process(server_port, security_settings)
|
||||
|
||||
try:
|
||||
# Test with invalid host header
|
||||
headers = {"Host": "evil.com"}
|
||||
|
||||
async with httpx.AsyncClient() as client:
|
||||
response = await client.get(f"http://127.0.0.1:{server_port}/sse", headers=headers)
|
||||
assert response.status_code == 421
|
||||
assert response.text == "Invalid Host header"
|
||||
|
||||
finally:
|
||||
process.terminate()
|
||||
process.join()
|
||||
async with sse_security_client(security_settings) as client:
|
||||
response = await client.get("/sse", headers={"Host": "evil.com"})
|
||||
assert response.status_code == 421
|
||||
assert response.text == "Invalid Host header"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_sse_security_invalid_origin_header(server_port: int):
|
||||
"""Test SSE with invalid Origin header."""
|
||||
# Configure security to allow the host but restrict origins
|
||||
async def test_sse_security_invalid_origin_header() -> None:
|
||||
"""An Origin header outside allowed_origins is rejected with 403."""
|
||||
security_settings = TransportSecuritySettings(
|
||||
enable_dns_rebinding_protection=True, allowed_hosts=["127.0.0.1:*"], allowed_origins=["http://localhost:*"]
|
||||
)
|
||||
process = start_server_process(server_port, security_settings)
|
||||
|
||||
try:
|
||||
# Test with invalid origin header
|
||||
headers = {"Origin": "http://evil.com"}
|
||||
|
||||
async with httpx.AsyncClient() as client:
|
||||
response = await client.get(f"http://127.0.0.1:{server_port}/sse", headers=headers)
|
||||
assert response.status_code == 403
|
||||
assert response.text == "Invalid Origin header"
|
||||
|
||||
finally:
|
||||
process.terminate()
|
||||
process.join()
|
||||
async with sse_security_client(security_settings) as client:
|
||||
response = await client.get("/sse", headers={"Origin": "http://evil.com"})
|
||||
assert response.status_code == 403
|
||||
assert response.text == "Invalid Origin header"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_sse_security_post_invalid_content_type(server_port: int):
|
||||
"""Test POST endpoint with invalid Content-Type header."""
|
||||
# Configure security to allow the host
|
||||
async def test_sse_security_post_invalid_content_type() -> None:
|
||||
"""A POST whose Content-Type is not application/json (or is missing) is rejected with 400."""
|
||||
security_settings = TransportSecuritySettings(
|
||||
enable_dns_rebinding_protection=True, allowed_hosts=["127.0.0.1:*"], allowed_origins=["http://127.0.0.1:*"]
|
||||
)
|
||||
process = start_server_process(server_port, security_settings)
|
||||
fake_session_id = "12345678123456781234567812345678"
|
||||
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=5.0) as client:
|
||||
# Test POST with invalid content type
|
||||
fake_session_id = "12345678123456781234567812345678"
|
||||
response = await client.post(
|
||||
f"http://127.0.0.1:{server_port}/messages/?session_id={fake_session_id}",
|
||||
headers={"Content-Type": "text/plain"},
|
||||
content="test",
|
||||
)
|
||||
assert response.status_code == 400
|
||||
assert response.text == "Invalid Content-Type header"
|
||||
async with sse_security_client(security_settings) as client:
|
||||
response = await client.post(
|
||||
f"/messages/?session_id={fake_session_id}",
|
||||
headers={"Content-Type": "text/plain"},
|
||||
content="test",
|
||||
)
|
||||
assert response.status_code == 400
|
||||
assert response.text == "Invalid Content-Type header"
|
||||
|
||||
# Test POST with missing content type
|
||||
response = await client.post(
|
||||
f"http://127.0.0.1:{server_port}/messages/?session_id={fake_session_id}", content="test"
|
||||
)
|
||||
assert response.status_code == 400
|
||||
assert response.text == "Invalid Content-Type header"
|
||||
|
||||
finally:
|
||||
process.terminate()
|
||||
process.join()
|
||||
response = await client.post(f"/messages/?session_id={fake_session_id}", content="test")
|
||||
assert response.status_code == 400
|
||||
assert response.text == "Invalid Content-Type header"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_sse_security_disabled(server_port: int):
|
||||
"""Test SSE with security disabled."""
|
||||
async def test_sse_security_disabled() -> None:
|
||||
"""With protection explicitly disabled, a disallowed Host still connects."""
|
||||
settings = TransportSecuritySettings(enable_dns_rebinding_protection=False)
|
||||
process = start_server_process(server_port, settings)
|
||||
|
||||
try:
|
||||
# Test with invalid host header - should still work
|
||||
headers = {"Host": "evil.com"}
|
||||
|
||||
async with httpx.AsyncClient(timeout=5.0) as client:
|
||||
# For SSE endpoints, we need to use stream to avoid timeout
|
||||
async with client.stream("GET", f"http://127.0.0.1:{server_port}/sse", headers=headers) as response:
|
||||
# Should connect successfully even with invalid host
|
||||
assert response.status_code == 200
|
||||
|
||||
finally:
|
||||
process.terminate()
|
||||
process.join()
|
||||
async with sse_security_client(settings) as client:
|
||||
async with client.stream("GET", "/sse", headers={"Host": "evil.com"}) as response:
|
||||
assert response.status_code == 200
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_sse_security_custom_allowed_hosts(server_port: int):
|
||||
"""Test SSE with custom allowed hosts."""
|
||||
async def test_sse_security_custom_allowed_hosts() -> None:
|
||||
"""A custom entry in allowed_hosts connects; hosts outside the list are still rejected."""
|
||||
settings = TransportSecuritySettings(
|
||||
enable_dns_rebinding_protection=True,
|
||||
allowed_hosts=["localhost", "127.0.0.1", "custom.host"],
|
||||
allowed_origins=["http://localhost", "http://127.0.0.1", "http://custom.host"],
|
||||
)
|
||||
process = start_server_process(server_port, settings)
|
||||
|
||||
try:
|
||||
# Test with custom allowed host
|
||||
headers = {"Host": "custom.host"}
|
||||
async with sse_security_client(settings) as client:
|
||||
async with client.stream("GET", "/sse", headers={"Host": "custom.host"}) as response:
|
||||
assert response.status_code == 200
|
||||
|
||||
async with httpx.AsyncClient(timeout=5.0) as client:
|
||||
# For SSE endpoints, we need to use stream to avoid timeout
|
||||
async with client.stream("GET", f"http://127.0.0.1:{server_port}/sse", headers=headers) as response:
|
||||
# Should connect successfully with custom host
|
||||
assert response.status_code == 200
|
||||
|
||||
# Test with non-allowed host
|
||||
headers = {"Host": "evil.com"}
|
||||
|
||||
async with httpx.AsyncClient() as client:
|
||||
response = await client.get(f"http://127.0.0.1:{server_port}/sse", headers=headers)
|
||||
assert response.status_code == 421
|
||||
assert response.text == "Invalid Host header"
|
||||
|
||||
finally:
|
||||
process.terminate()
|
||||
process.join()
|
||||
response = await client.get("/sse", headers={"Host": "evil.com"})
|
||||
assert response.status_code == 421
|
||||
assert response.text == "Invalid Host header"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_sse_security_wildcard_ports(server_port: int):
|
||||
"""Test SSE with wildcard port patterns."""
|
||||
async def test_sse_security_wildcard_ports() -> None:
|
||||
"""A `host:*` pattern accepts that host with any port, for Host and Origin alike."""
|
||||
settings = TransportSecuritySettings(
|
||||
enable_dns_rebinding_protection=True,
|
||||
allowed_hosts=["localhost:*", "127.0.0.1:*"],
|
||||
allowed_origins=["http://localhost:*", "http://127.0.0.1:*"],
|
||||
)
|
||||
process = start_server_process(server_port, settings)
|
||||
|
||||
try:
|
||||
# Test with various port numbers
|
||||
async with sse_security_client(settings) as client:
|
||||
for test_port in [8080, 3000, 9999]:
|
||||
headers = {"Host": f"localhost:{test_port}"}
|
||||
async with client.stream("GET", "/sse", headers={"Host": f"localhost:{test_port}"}) as response:
|
||||
assert response.status_code == 200
|
||||
|
||||
async with httpx.AsyncClient(timeout=5.0) as client:
|
||||
# For SSE endpoints, we need to use stream to avoid timeout
|
||||
async with client.stream("GET", f"http://127.0.0.1:{server_port}/sse", headers=headers) as response:
|
||||
# Should connect successfully with any port
|
||||
assert response.status_code == 200
|
||||
|
||||
headers = {"Origin": f"http://localhost:{test_port}"}
|
||||
|
||||
async with httpx.AsyncClient(timeout=5.0) as client:
|
||||
# For SSE endpoints, we need to use stream to avoid timeout
|
||||
async with client.stream("GET", f"http://127.0.0.1:{server_port}/sse", headers=headers) as response:
|
||||
# Should connect successfully with any port
|
||||
assert response.status_code == 200
|
||||
|
||||
finally:
|
||||
process.terminate()
|
||||
process.join()
|
||||
async with client.stream("GET", "/sse", headers={"Origin": f"http://localhost:{test_port}"}) as response:
|
||||
assert response.status_code == 200
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_sse_security_post_valid_content_type(server_port: int):
|
||||
"""Test POST endpoint with valid Content-Type headers."""
|
||||
# Configure security to allow the host
|
||||
async def test_sse_security_post_valid_content_type() -> None:
|
||||
"""Every application/json Content-Type variant passes validation (reaching the session lookup)."""
|
||||
security_settings = TransportSecuritySettings(
|
||||
enable_dns_rebinding_protection=True, allowed_hosts=["127.0.0.1:*"], allowed_origins=["http://127.0.0.1:*"]
|
||||
)
|
||||
process = start_server_process(server_port, security_settings)
|
||||
valid_content_types = [
|
||||
"application/json",
|
||||
"application/json; charset=utf-8",
|
||||
"application/json;charset=utf-8",
|
||||
"APPLICATION/JSON", # Case insensitive
|
||||
]
|
||||
# A well-formed session ID that no live session owns.
|
||||
fake_session_id = "12345678123456781234567812345678"
|
||||
|
||||
try:
|
||||
async with httpx.AsyncClient() as client:
|
||||
# Test with various valid content types
|
||||
valid_content_types = [
|
||||
"application/json",
|
||||
"application/json; charset=utf-8",
|
||||
"application/json;charset=utf-8",
|
||||
"APPLICATION/JSON", # Case insensitive
|
||||
]
|
||||
|
||||
for content_type in valid_content_types:
|
||||
# Use a valid UUID format (even though session won't exist)
|
||||
fake_session_id = "12345678123456781234567812345678"
|
||||
response = await client.post(
|
||||
f"http://127.0.0.1:{server_port}/messages/?session_id={fake_session_id}",
|
||||
headers={"Content-Type": content_type},
|
||||
json={"test": "data"},
|
||||
)
|
||||
# Will get 404 because session doesn't exist, but that's OK
|
||||
# We're testing that it passes the content-type check
|
||||
assert response.status_code == 404
|
||||
assert response.text == "Could not find session"
|
||||
|
||||
finally:
|
||||
process.terminate()
|
||||
process.join()
|
||||
async with sse_security_client(security_settings) as client:
|
||||
for content_type in valid_content_types:
|
||||
response = await client.post(
|
||||
f"/messages/?session_id={fake_session_id}",
|
||||
headers={"Content-Type": content_type},
|
||||
json={"test": "data"},
|
||||
)
|
||||
# 404 proves the request passed the content-type check and reached the session lookup.
|
||||
assert response.status_code == 404
|
||||
assert response.text == "Could not find session"
|
||||
|
||||
|
||||
def _authenticated_user(client_id: str, subject: str | None = None, issuer: str | None = None) -> AuthenticatedUser:
|
||||
|
|
|
|||
Loading…
Reference in a new issue