Max
173 lines
6.5 KiB
Python
173 lines
6.5 KiB
Python
"""Tests for WebSocket server request validation."""
|
|
|
|
# pyright: reportDeprecated=false
|
|
|
|
import logging
|
|
import multiprocessing
|
|
import socket
|
|
import warnings
|
|
|
|
import pytest
|
|
import uvicorn
|
|
from starlette.applications import Starlette
|
|
from starlette.routing import WebSocketRoute
|
|
from starlette.types import Message, Scope
|
|
from starlette.websockets import WebSocket
|
|
from websockets.asyncio.client import connect
|
|
from websockets.exceptions import InvalidStatus
|
|
from websockets.typing import Subprotocol
|
|
|
|
from mcp.server import Server
|
|
from mcp.server.transport_security import TransportSecuritySettings
|
|
from mcp.server.websocket import websocket_server
|
|
from tests.test_helpers import wait_for_server
|
|
|
|
logger = logging.getLogger(__name__)
|
|
SERVER_NAME = "test_ws_security_server"
|
|
|
|
# This suite intentionally exercises the deprecated WebSocket transport.
|
|
pytestmark = pytest.mark.filterwarnings(
|
|
"ignore:The WebSocket (client|server) transport is deprecated:DeprecationWarning"
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def server_port() -> int:
|
|
with socket.socket() as s:
|
|
s.bind(("127.0.0.1", 0))
|
|
return s.getsockname()[1]
|
|
|
|
|
|
def run_server_with_settings(port: int, security_settings: TransportSecuritySettings | None = None): # pragma: no cover
|
|
"""Run a WebSocket MCP server with the given security settings."""
|
|
warnings.filterwarnings("ignore", category=DeprecationWarning)
|
|
server = Server(SERVER_NAME)
|
|
|
|
async def handle_ws(websocket: WebSocket) -> None:
|
|
try:
|
|
async with websocket_server(
|
|
websocket.scope, websocket.receive, websocket.send, security_settings=security_settings
|
|
) as streams:
|
|
await server.run(streams[0], streams[1], server.create_initialization_options())
|
|
except ValueError as exc:
|
|
logger.debug(f"WebSocket connection failed validation: {exc}")
|
|
|
|
app = Starlette(routes=[WebSocketRoute("/ws", endpoint=handle_ws)])
|
|
uvicorn.run(app, host="127.0.0.1", port=port, log_level="error")
|
|
|
|
|
|
def start_server_process(port: int, security_settings: TransportSecuritySettings | None = None):
|
|
"""Start the server in a subprocess and wait until it accepts connections."""
|
|
process = multiprocessing.Process(target=run_server_with_settings, args=(port, security_settings))
|
|
process.start()
|
|
wait_for_server(port)
|
|
return process
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_ws_security_default_settings(server_port: int) -> None:
|
|
"""With no security settings the WebSocket transport accepts any Origin (matches SSE/StreamableHTTP default)."""
|
|
process = start_server_process(server_port)
|
|
try:
|
|
async with connect(
|
|
f"ws://127.0.0.1:{server_port}/ws",
|
|
subprotocols=[Subprotocol("mcp")],
|
|
additional_headers={"Origin": "http://evil.com"},
|
|
) as ws:
|
|
assert ws.subprotocol == "mcp"
|
|
finally:
|
|
process.terminate()
|
|
process.join()
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_ws_security_invalid_origin_header(server_port: int) -> None:
|
|
"""An Origin not in allowed_origins is rejected before the handshake completes."""
|
|
settings = TransportSecuritySettings(
|
|
enable_dns_rebinding_protection=True, allowed_hosts=["127.0.0.1:*"], allowed_origins=["http://localhost:*"]
|
|
)
|
|
process = start_server_process(server_port, settings)
|
|
try:
|
|
with pytest.raises(InvalidStatus) as exc_info:
|
|
async with connect(
|
|
f"ws://127.0.0.1:{server_port}/ws",
|
|
subprotocols=[Subprotocol("mcp")],
|
|
additional_headers={"Origin": "http://evil.com"},
|
|
):
|
|
pytest.fail("handshake should have been rejected") # pragma: no cover
|
|
assert exc_info.value.response.status_code == 403
|
|
finally:
|
|
process.terminate()
|
|
process.join()
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_ws_security_invalid_host_header(server_port: int) -> None:
|
|
"""A Host not in allowed_hosts is rejected before the handshake completes."""
|
|
settings = TransportSecuritySettings(enable_dns_rebinding_protection=True, allowed_hosts=["example.com"])
|
|
process = start_server_process(server_port, settings)
|
|
try:
|
|
with pytest.raises(InvalidStatus) as exc_info:
|
|
async with connect(f"ws://127.0.0.1:{server_port}/ws", subprotocols=[Subprotocol("mcp")]):
|
|
pytest.fail("handshake should have been rejected") # pragma: no cover
|
|
assert exc_info.value.response.status_code == 403
|
|
finally:
|
|
process.terminate()
|
|
process.join()
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_ws_security_allowed_origin(server_port: int) -> None:
|
|
"""An Origin matching allowed_origins is accepted."""
|
|
settings = TransportSecuritySettings(
|
|
enable_dns_rebinding_protection=True, allowed_hosts=["127.0.0.1:*"], allowed_origins=["http://localhost:*"]
|
|
)
|
|
process = start_server_process(server_port, settings)
|
|
try:
|
|
async with connect(
|
|
f"ws://127.0.0.1:{server_port}/ws",
|
|
subprotocols=[Subprotocol("mcp")],
|
|
additional_headers={"Origin": "http://localhost:8080"},
|
|
) as ws:
|
|
assert ws.subprotocol == "mcp"
|
|
finally:
|
|
process.terminate()
|
|
process.join()
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_ws_security_disabled(server_port: int) -> None:
|
|
"""Explicitly disabling protection accepts any Origin."""
|
|
settings = TransportSecuritySettings(enable_dns_rebinding_protection=False)
|
|
process = start_server_process(server_port, settings)
|
|
try:
|
|
async with connect(
|
|
f"ws://127.0.0.1:{server_port}/ws",
|
|
subprotocols=[Subprotocol("mcp")],
|
|
additional_headers={"Origin": "http://evil.com"},
|
|
) as ws:
|
|
assert ws.subprotocol == "mcp"
|
|
finally:
|
|
process.terminate()
|
|
process.join()
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_ws_security_rejects_before_accept() -> None:
|
|
"""A failing validation closes the connection before the handshake is accepted."""
|
|
settings = TransportSecuritySettings(enable_dns_rebinding_protection=True, allowed_hosts=["example.com"])
|
|
sent: list[Message] = []
|
|
|
|
async def receive() -> Message:
|
|
raise NotImplementedError
|
|
|
|
async def send(message: Message) -> None:
|
|
sent.append(message)
|
|
|
|
scope: Scope = {"type": "websocket", "headers": [(b"host", b"evil.com")]}
|
|
with pytest.raises(ValueError, match="Request validation failed"):
|
|
async with websocket_server(scope, receive, send, security_settings=settings):
|
|
pytest.fail("should not yield streams") # pragma: no cover
|
|
|
|
assert [m["type"] for m in sent] == ["websocket.close"]
|