Bind transport sessions to the authenticated principal (#2718)
This commit is contained in:
Max
GitHub
parent
24725633f1
commit
616476f692
3 changed files with 647 additions and 42 deletions
|
|
@ -1,27 +1,44 @@
|
|||
"""Tests for SSE server DNS rebinding protection."""
|
||||
"""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
|
||||
from starlette.routing import Mount, Route
|
||||
from starlette.types import Message, Receive, Scope, Send
|
||||
|
||||
from mcp.server import Server
|
||||
from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser
|
||||
from mcp.server.auth.provider import AccessToken
|
||||
from mcp.server.sse import SseServerTransport
|
||||
from mcp.server.transport_security import TransportSecuritySettings
|
||||
from mcp.types import Tool
|
||||
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
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
SERVER_NAME = "test_sse_security_server"
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def reset_sse_starlette_exit_event() -> None:
|
||||
"""sse-starlette<2 caches a module-level anyio.Event on AppStatus; reset it
|
||||
between tests so it is not bound to a previous test's event loop."""
|
||||
app_status = getattr(sse_starlette.sse, "AppStatus", None)
|
||||
if app_status is not None and hasattr(app_status, "should_exit_event"): # pragma: lax no cover
|
||||
app_status.should_exit_event = None
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def server_port() -> int:
|
||||
with socket.socket() as s:
|
||||
|
|
@ -291,3 +308,270 @@ async def test_sse_security_post_valid_content_type(server_port: int):
|
|||
finally:
|
||||
process.terminate()
|
||||
process.join()
|
||||
|
||||
|
||||
def _authenticated_user(client_id: str, subject: str | None = None, issuer: str | None = None) -> AuthenticatedUser:
|
||||
"""Build the scope["user"] value that AuthenticationMiddleware would set for this principal."""
|
||||
claims = {"iss": issuer} if issuer is not None else None
|
||||
return AuthenticatedUser(AccessToken(token="token", client_id=client_id, scopes=[], subject=subject, claims=claims))
|
||||
|
||||
|
||||
def _sse_scope(
|
||||
method: str, path: str, user: AuthenticatedUser | None, *, query_string: bytes = b"", body: bytes = b""
|
||||
) -> tuple[Scope, Receive, Send, list[Message]]:
|
||||
"""Build an ASGI scope/receive/send triple for a request to the SSE transport."""
|
||||
scope: Scope = {
|
||||
"type": "http",
|
||||
"method": method,
|
||||
"path": path,
|
||||
"root_path": "",
|
||||
"query_string": query_string,
|
||||
"headers": [(b"content-type", b"application/json")],
|
||||
}
|
||||
if user is not None:
|
||||
scope["user"] = user
|
||||
sent: list[Message] = []
|
||||
|
||||
async def receive() -> Message:
|
||||
return {"type": "http.request", "body": body, "more_body": False}
|
||||
|
||||
async def send(message: Message) -> None:
|
||||
sent.append(message)
|
||||
|
||||
return scope, receive, send, sent
|
||||
|
||||
|
||||
def _response_status(sent: list[Message]) -> int:
|
||||
response_start = next(msg for msg in sent if msg["type"] == "http.response.start")
|
||||
return response_start["status"]
|
||||
|
||||
|
||||
async def _post_message(transport: SseServerTransport, session_id: str, user: AuthenticatedUser | None) -> int:
|
||||
"""POST a message to an SSE session as `user` and return the response status."""
|
||||
body = b'{"jsonrpc": "2.0", "id": 1, "method": "ping", "params": null}'
|
||||
scope, receive, send, sent = _sse_scope(
|
||||
"POST", "/messages/", user, query_string=f"session_id={session_id}".encode(), body=body
|
||||
)
|
||||
await transport.handle_post_message(scope, receive, send)
|
||||
return _response_status(sent)
|
||||
|
||||
|
||||
_Principal = tuple[str] | tuple[str, str] | tuple[str, str, str]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
@pytest.mark.parametrize(
|
||||
("creator", "sender", "expected"),
|
||||
[
|
||||
pytest.param(("client-a",), ("client-b",), 404, id="different-client"),
|
||||
pytest.param(("client-a",), None, 404, id="unauthenticated-sender"),
|
||||
pytest.param(("client-a", "alice"), ("client-a", "bob"), 404, id="same-client-different-subject"),
|
||||
pytest.param(("client-a", "alice"), ("client-a",), 404, id="same-client-no-subject"),
|
||||
pytest.param(
|
||||
("client-a", "alice", "https://i1"), ("client-a", "alice", "https://i2"), 404, id="different-issuer"
|
||||
),
|
||||
pytest.param(None, ("client-a",), 404, id="unauthenticated-creator"),
|
||||
pytest.param(("client-a",), ("client-a",), 202, id="same-client"),
|
||||
pytest.param(("client-a", "alice"), ("client-a", "alice"), 202, id="same-client-and-subject"),
|
||||
pytest.param(None, None, 202, id="both-unauthenticated"),
|
||||
],
|
||||
)
|
||||
async def test_sse_post_requires_the_credential_that_created_the_session(
|
||||
creator: _Principal | None,
|
||||
sender: _Principal | None,
|
||||
expected: int,
|
||||
):
|
||||
"""The session endpoint URL issued to one authenticated principal must not
|
||||
accept messages from a request authenticated as a different one."""
|
||||
transport = SseServerTransport("/messages/")
|
||||
session_id_received = anyio.Event()
|
||||
session_ids: list[str] = []
|
||||
client_disconnected = anyio.Event()
|
||||
|
||||
async def get_send(message: Message) -> None:
|
||||
# The first body chunk is the SSE event announcing the session URI to POST messages to.
|
||||
if message["type"] == "http.response.body" and not session_ids:
|
||||
match = re.search(rb"session_id=([0-9a-f]{32})", message.get("body", b""))
|
||||
assert match is not None, f"expected the endpoint event first, got {message!r}"
|
||||
session_ids.append(match.group(1).decode())
|
||||
session_id_received.set()
|
||||
|
||||
async def get_receive() -> Message:
|
||||
# The SSE client stays connected until the test signals otherwise.
|
||||
await client_disconnected.wait()
|
||||
return {"type": "http.disconnect"}
|
||||
|
||||
creator_user = _authenticated_user(*creator) if creator is not None else None
|
||||
sender_user = _authenticated_user(*sender) if sender is not None else None
|
||||
|
||||
async def hold_sse_connection() -> None:
|
||||
"""Establish the SSE session as `creator` and keep it open, as a server would."""
|
||||
scope, _, _, _ = _sse_scope("GET", "/sse", creator_user)
|
||||
with anyio.fail_after(5):
|
||||
async with transport.connect_sse(scope, get_receive, get_send) as (read_stream, write_stream):
|
||||
async with read_stream, write_stream: # pragma: no branch
|
||||
async for _ in read_stream:
|
||||
pass
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
tg.start_soon(hold_sse_connection)
|
||||
with anyio.fail_after(5):
|
||||
await session_id_received.wait()
|
||||
|
||||
assert await _post_message(transport, session_ids[0], sender_user) == expected
|
||||
|
||||
client_disconnected.set()
|
||||
|
||||
# Once the connection is gone the session is no longer routable.
|
||||
assert await _post_message(transport, session_ids[0], creator_user) == 404
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_sse_connect_rejects_a_non_http_scope():
|
||||
"""connect_sse refuses ASGI scopes that are not HTTP requests."""
|
||||
transport = SseServerTransport("/messages/")
|
||||
with pytest.raises(ValueError):
|
||||
async with transport.connect_sse({"type": "websocket"}, _no_receive, _no_send):
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_sse_connect_rejects_a_disallowed_host():
|
||||
"""connect_sse rejects requests whose Host header fails the configured security check."""
|
||||
settings = TransportSecuritySettings(allowed_hosts=["allowed.example.com"])
|
||||
transport = SseServerTransport("/messages/", security_settings=settings)
|
||||
scope, receive, send, sent = _sse_scope("GET", "/sse", None)
|
||||
scope["headers"] = [(b"host", b"disallowed.example.com")]
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
async with transport.connect_sse(scope, receive, send):
|
||||
raise NotImplementedError
|
||||
assert _response_status(sent) == 421
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_sse_post_without_a_session_id_returns_400():
|
||||
"""POSTs to the messages endpoint must include a session_id query parameter."""
|
||||
transport = SseServerTransport("/messages/")
|
||||
scope, receive, send, sent = _sse_scope("POST", "/messages/", None)
|
||||
|
||||
await transport.handle_post_message(scope, receive, send)
|
||||
assert _response_status(sent) == 400
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_sse_post_with_a_malformed_session_id_returns_400():
|
||||
"""A session_id that is not 32 hex characters is rejected before any session lookup."""
|
||||
transport = SseServerTransport("/messages/")
|
||||
scope, receive, send, sent = _sse_scope("POST", "/messages/", None, query_string=b"session_id=not-hex")
|
||||
|
||||
await transport.handle_post_message(scope, receive, send)
|
||||
assert _response_status(sent) == 400
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_sse_post_with_a_disallowed_host_is_rejected_before_session_lookup():
|
||||
"""The transport security check on POST runs before any session-ID handling."""
|
||||
settings = TransportSecuritySettings(allowed_hosts=["allowed.example.com"])
|
||||
transport = SseServerTransport("/messages/", security_settings=settings)
|
||||
scope, receive, send, sent = _sse_scope("POST", "/messages/", None)
|
||||
scope["headers"] = [(b"host", b"disallowed.example.com"), (b"content-type", b"application/json")]
|
||||
|
||||
await transport.handle_post_message(scope, receive, send)
|
||||
assert _response_status(sent) == 421
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_sse_round_trip_delivers_posted_messages_and_streams_responses():
|
||||
"""A POSTed JSON-RPC message reaches the server's read stream, and a message
|
||||
written to the server's write stream is sent to the client as an SSE event."""
|
||||
transport = SseServerTransport("/messages/")
|
||||
session = _SseSession(transport)
|
||||
|
||||
with anyio.fail_after(5):
|
||||
async with anyio.create_task_group() as tg:
|
||||
tg.start_soon(session.hold)
|
||||
await session.ready.wait()
|
||||
|
||||
# POST a parse-failing body: client gets 400, server's read stream receives the error.
|
||||
scope, receive, send, sent = _sse_scope(
|
||||
"POST", "/messages/", None, query_string=f"session_id={session.session_id}".encode(), body=b"not json"
|
||||
)
|
||||
await transport.handle_post_message(scope, receive, send)
|
||||
assert _response_status(sent) == 400
|
||||
assert isinstance(await session.next_read_item(), Exception)
|
||||
|
||||
# POST a valid message: client gets 202, server's read stream receives it.
|
||||
assert await _post_message(transport, session.session_id, None) == 202
|
||||
received = await session.next_read_item()
|
||||
assert isinstance(received, SessionMessage)
|
||||
assert isinstance(received.message, JSONRPCRequest)
|
||||
assert received.message.method == "ping"
|
||||
|
||||
# Server writes a response: it appears as an SSE `message` event on the GET stream.
|
||||
outgoing = JSONRPCResponse(jsonrpc="2.0", id=1, result={})
|
||||
await session.write_stream.send(SessionMessage(outgoing))
|
||||
chunk = await session.next_body_chunk()
|
||||
assert b"event: message" in chunk
|
||||
assert outgoing.model_dump_json(by_alias=True, exclude_unset=True).encode() in chunk
|
||||
|
||||
session.disconnect()
|
||||
|
||||
|
||||
class _SseSession:
|
||||
"""Drive an in-process SSE GET connection and surface what the server reads and the client receives.
|
||||
|
||||
`hold` runs the connection in a background task and consumes the server-side read stream
|
||||
into a buffer so that `handle_post_message` (which writes to that stream with a zero-capacity
|
||||
channel) never blocks the test body.
|
||||
"""
|
||||
|
||||
def __init__(self, transport: SseServerTransport) -> None:
|
||||
self.transport = transport
|
||||
self.ready = anyio.Event()
|
||||
self._disconnected = anyio.Event()
|
||||
self._body_send, self._body_recv = anyio.create_memory_object_stream[bytes](16)
|
||||
self._read_send, self._read_recv = anyio.create_memory_object_stream[SessionMessage | Exception](16)
|
||||
self.session_id = ""
|
||||
self.write_stream: WriteStream[SessionMessage]
|
||||
|
||||
async def hold(self) -> None:
|
||||
scope, _, _, _ = _sse_scope("GET", "/sse", None)
|
||||
async with self.transport.connect_sse(scope, self._receive, self._send) as (read, write):
|
||||
self.write_stream = write
|
||||
async with read, write, self._body_send, self._body_recv, self._read_send, self._read_recv:
|
||||
async for item in read:
|
||||
await self._read_send.send(item)
|
||||
|
||||
def disconnect(self) -> None:
|
||||
self._disconnected.set()
|
||||
|
||||
async def next_read_item(self) -> SessionMessage | Exception:
|
||||
return await self._read_recv.receive()
|
||||
|
||||
async def next_body_chunk(self) -> bytes:
|
||||
return await self._body_recv.receive()
|
||||
|
||||
async def _receive(self) -> Message:
|
||||
await self._disconnected.wait()
|
||||
return {"type": "http.disconnect"}
|
||||
|
||||
async def _send(self, message: Message) -> None:
|
||||
if message["type"] != "http.response.body":
|
||||
return
|
||||
body: bytes = message.get("body", b"")
|
||||
if not self.session_id:
|
||||
match = re.search(rb"session_id=([0-9a-f]{32})", body)
|
||||
assert match is not None, f"expected the endpoint event first, got {message!r}"
|
||||
self.session_id = match.group(1).decode()
|
||||
self.ready.set()
|
||||
else:
|
||||
await self._body_send.send(body)
|
||||
|
||||
|
||||
async def _no_receive() -> Message:
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
async def _no_send(message: Message) -> None:
|
||||
raise NotImplementedError
|
||||
|
|
|
|||
|
|
@ -8,11 +8,13 @@ from unittest.mock import AsyncMock, patch
|
|||
import anyio
|
||||
import httpx
|
||||
import pytest
|
||||
from starlette.types import Message
|
||||
from starlette.types import Message, Scope
|
||||
|
||||
from mcp import Client
|
||||
from mcp.client.streamable_http import streamable_http_client
|
||||
from mcp.server import Server, ServerRequestContext, streamable_http_manager
|
||||
from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser
|
||||
from mcp.server.auth.provider import AccessToken
|
||||
from mcp.server.streamable_http import MCP_SESSION_ID_HEADER, StreamableHTTPServerTransport
|
||||
from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
|
||||
from mcp.types import INVALID_REQUEST, ListToolsResult, PaginatedRequestParams
|
||||
|
|
@ -413,3 +415,166 @@ def test_session_idle_timeout_rejects_non_positive():
|
|||
def test_session_idle_timeout_rejects_stateless():
|
||||
with pytest.raises(RuntimeError, match="not supported in stateless"):
|
||||
StreamableHTTPSessionManager(app=Server("test"), session_idle_timeout=30, stateless=True)
|
||||
|
||||
|
||||
def _user(client_id: str, subject: str | None = None, issuer: str | None = None) -> AuthenticatedUser:
|
||||
"""Build the scope["user"] value that AuthenticationMiddleware would set for this principal."""
|
||||
claims = {"iss": issuer} if issuer is not None else None
|
||||
return AuthenticatedUser(AccessToken(token="token", client_id=client_id, scopes=[], subject=subject, claims=claims))
|
||||
|
||||
|
||||
def _request_scope(
|
||||
*, session_id: str | None = None, user: AuthenticatedUser | None = None, method: str = "POST"
|
||||
) -> Scope:
|
||||
"""Build an ASGI scope for a request to the MCP endpoint."""
|
||||
headers = [
|
||||
(b"content-type", b"application/json"),
|
||||
(b"accept", b"application/json, text/event-stream"),
|
||||
]
|
||||
if session_id is not None:
|
||||
headers.append((b"mcp-session-id", session_id.encode()))
|
||||
scope: Scope = {
|
||||
"type": "http",
|
||||
"method": method,
|
||||
"path": "/mcp",
|
||||
"headers": headers,
|
||||
}
|
||||
if user is not None:
|
||||
scope["user"] = user
|
||||
return scope
|
||||
|
||||
|
||||
async def _open_session(manager: StreamableHTTPSessionManager, user: AuthenticatedUser | None) -> str:
|
||||
"""Create a new session as `user` and return its session ID."""
|
||||
sent_messages: list[Message] = []
|
||||
|
||||
async def mock_send(message: Message) -> None:
|
||||
sent_messages.append(message)
|
||||
|
||||
async def mock_receive() -> Message:
|
||||
return {"type": "http.request", "body": b"", "more_body": False}
|
||||
|
||||
await manager.handle_request(_request_scope(user=user), mock_receive, mock_send)
|
||||
|
||||
response_start = next(msg for msg in sent_messages if msg["type"] == "http.response.start")
|
||||
headers = dict(response_start.get("headers", []))
|
||||
return headers[MCP_SESSION_ID_HEADER.encode()].decode()
|
||||
|
||||
|
||||
async def _request_session(
|
||||
manager: StreamableHTTPSessionManager, session_id: str, user: AuthenticatedUser | None, method: str = "POST"
|
||||
) -> int:
|
||||
"""Send a request for an existing session as `user` and return the response status."""
|
||||
sent_messages: list[Message] = []
|
||||
|
||||
async def mock_send(message: Message) -> None:
|
||||
sent_messages.append(message)
|
||||
|
||||
async def mock_receive() -> Message:
|
||||
return {"type": "http.request", "body": b"", "more_body": False}
|
||||
|
||||
await manager.handle_request(
|
||||
_request_scope(session_id=session_id, user=user, method=method), mock_receive, mock_send
|
||||
)
|
||||
|
||||
response_start = next(msg for msg in sent_messages if msg["type"] == "http.response.start")
|
||||
return response_start["status"]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def manager_with_live_session():
|
||||
"""A running manager around a real `Server`. Sessions remain registered until
|
||||
`manager.run()` exits because `Server.run` blocks waiting for an initialize message."""
|
||||
manager = StreamableHTTPSessionManager(app=Server("test-session-credentials"))
|
||||
async with manager.run():
|
||||
yield manager
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_session_accepts_requests_from_the_credential_that_created_it(
|
||||
manager_with_live_session: StreamableHTTPSessionManager,
|
||||
) -> None:
|
||||
"""Requests presenting the same credential as the one that created the session are served."""
|
||||
manager = manager_with_live_session
|
||||
session_id = await _open_session(manager, _user("client-a"))
|
||||
|
||||
status = await _request_session(manager, session_id, _user("client-a"))
|
||||
|
||||
# The request passes the manager's credential check and reaches the
|
||||
# session's transport, instead of being answered with 404 by the manager.
|
||||
assert status != 404
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
@pytest.mark.parametrize("method", ["POST", "GET", "DELETE"])
|
||||
async def test_session_rejects_requests_from_a_different_credential(
|
||||
manager_with_live_session: StreamableHTTPSessionManager, method: str
|
||||
) -> None:
|
||||
"""A session created by one credential cannot be used with another credential, whatever the method."""
|
||||
manager = manager_with_live_session
|
||||
session_id = await _open_session(manager, _user("client-a"))
|
||||
|
||||
assert await _request_session(manager, session_id, _user("client-b"), method) == 404
|
||||
# The session is still registered and still serves its creator.
|
||||
assert await _request_session(manager, session_id, _user("client-a")) != 404
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_session_rejects_requests_from_a_different_subject_of_the_same_client(
|
||||
manager_with_live_session: StreamableHTTPSessionManager,
|
||||
) -> None:
|
||||
"""Two end-users that share an OAuth client cannot use each other's sessions."""
|
||||
manager = manager_with_live_session
|
||||
session_id = await _open_session(manager, _user("client-a", subject="alice"))
|
||||
|
||||
assert await _request_session(manager, session_id, _user("client-a", subject="bob")) == 404
|
||||
assert await _request_session(manager, session_id, _user("client-a", subject=None)) == 404
|
||||
assert await _request_session(manager, session_id, _user("client-a", subject="alice")) != 404
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_session_rejects_requests_with_the_same_subject_from_a_different_issuer(
|
||||
manager_with_live_session: StreamableHTTPSessionManager,
|
||||
) -> None:
|
||||
"""A subject is unique only per issuer, so a colliding subject from a different issuer is not the same principal."""
|
||||
manager = manager_with_live_session
|
||||
creator = _user("client-a", subject="alice", issuer="https://issuer.one")
|
||||
session_id = await _open_session(manager, creator)
|
||||
|
||||
other_issuer = _user("client-a", subject="alice", issuer="https://issuer.two")
|
||||
assert await _request_session(manager, session_id, other_issuer) == 404
|
||||
assert await _request_session(manager, session_id, _user("client-a", subject="alice")) == 404
|
||||
assert await _request_session(manager, session_id, creator) != 404
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_session_rejects_unauthenticated_requests_for_an_authenticated_session(
|
||||
manager_with_live_session: StreamableHTTPSessionManager,
|
||||
) -> None:
|
||||
"""A session created with a credential cannot be used without one."""
|
||||
manager = manager_with_live_session
|
||||
session_id = await _open_session(manager, _user("client-a"))
|
||||
|
||||
assert await _request_session(manager, session_id, None) == 404
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_session_rejects_authenticated_requests_for_an_anonymous_session(
|
||||
manager_with_live_session: StreamableHTTPSessionManager,
|
||||
) -> None:
|
||||
"""A session created without a credential cannot be used with one."""
|
||||
manager = manager_with_live_session
|
||||
session_id = await _open_session(manager, None)
|
||||
|
||||
assert await _request_session(manager, session_id, _user("client-a")) == 404
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_anonymous_session_accepts_anonymous_requests(
|
||||
manager_with_live_session: StreamableHTTPSessionManager,
|
||||
) -> None:
|
||||
"""Servers without authentication keep working: no credential on either side."""
|
||||
manager = manager_with_live_session
|
||||
session_id = await _open_session(manager, None)
|
||||
|
||||
assert await _request_session(manager, session_id, None) != 404
|
||||
|
|
|
|||
88
tests/server/test_transport_security.py
Normal file
88
tests/server/test_transport_security.py
Normal file
|
|
@ -0,0 +1,88 @@
|
|||
"""Tests for the transport-security request validation middleware."""
|
||||
|
||||
import pytest
|
||||
from starlette.requests import Request
|
||||
|
||||
from mcp.server.transport_security import TransportSecurityMiddleware, TransportSecuritySettings
|
||||
|
||||
|
||||
def _request(host: str | None, origin: str | None, content_type: str | None = "application/json") -> Request:
|
||||
headers: list[tuple[bytes, bytes]] = []
|
||||
if content_type is not None:
|
||||
headers.append((b"content-type", content_type.encode()))
|
||||
if host is not None:
|
||||
headers.append((b"host", host.encode()))
|
||||
if origin is not None:
|
||||
headers.append((b"origin", origin.encode()))
|
||||
return Request({"type": "http", "method": "GET", "headers": headers})
|
||||
|
||||
|
||||
SETTINGS = TransportSecuritySettings(
|
||||
enable_dns_rebinding_protection=True,
|
||||
allowed_hosts=["good.example", "wild.example:*"],
|
||||
allowed_origins=["http://good.example", "http://wild.example:*"],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
@pytest.mark.parametrize(
|
||||
("host", "origin", "expected"),
|
||||
[
|
||||
pytest.param(None, None, 421, id="missing-host"),
|
||||
pytest.param("evil.example", None, 421, id="host-no-match"),
|
||||
pytest.param("evil.example:9000", None, 421, id="host-wildcard-base-mismatch"),
|
||||
pytest.param("good.example", None, None, id="host-exact-no-origin"),
|
||||
pytest.param("wild.example:9000", None, None, id="host-wildcard-match"),
|
||||
pytest.param("good.example", "http://evil.example", 403, id="origin-no-match"),
|
||||
pytest.param("good.example", "http://evil.example:9000", 403, id="origin-wildcard-base-mismatch"),
|
||||
pytest.param("good.example", "http://good.example", None, id="origin-exact"),
|
||||
pytest.param("good.example", "http://wild.example:9000", None, id="origin-wildcard-match"),
|
||||
],
|
||||
)
|
||||
async def test_validate_request_checks_host_then_origin(
|
||||
host: str | None, origin: str | None, expected: int | None
|
||||
) -> None:
|
||||
"""Host is checked first, then Origin; exact and wildcard-port allowlist entries are honoured."""
|
||||
middleware = TransportSecurityMiddleware(SETTINGS)
|
||||
response = await middleware.validate_request(_request(host, origin))
|
||||
assert (None if response is None else response.status_code) == expected
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_validate_request_skips_host_and_origin_when_protection_is_disabled() -> None:
|
||||
"""With DNS-rebinding protection off, any Host/Origin is accepted."""
|
||||
middleware = TransportSecurityMiddleware(TransportSecuritySettings(enable_dns_rebinding_protection=False))
|
||||
assert await middleware.validate_request(_request("evil.example", "http://evil.example")) is None
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_validate_request_defaults_to_protection_disabled() -> None:
|
||||
"""Constructing the middleware without settings leaves DNS-rebinding protection off."""
|
||||
middleware = TransportSecurityMiddleware()
|
||||
assert await middleware.validate_request(_request("evil.example", "http://evil.example")) is None
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
@pytest.mark.parametrize(
|
||||
("content_type", "expected"),
|
||||
[
|
||||
pytest.param("application/json", None, id="json"),
|
||||
pytest.param("application/json; charset=utf-8", None, id="json-with-charset"),
|
||||
pytest.param("APPLICATION/JSON", None, id="case-insensitive"),
|
||||
pytest.param("text/plain", 400, id="wrong-type"),
|
||||
pytest.param(None, 400, id="missing"),
|
||||
],
|
||||
)
|
||||
async def test_validate_request_checks_content_type_on_post(content_type: str | None, expected: int | None) -> None:
|
||||
"""POST requests must carry an application/json Content-Type, regardless of DNS-rebinding settings."""
|
||||
middleware = TransportSecurityMiddleware()
|
||||
response = await middleware.validate_request(_request("any", None, content_type=content_type), is_post=True)
|
||||
assert (None if response is None else response.status_code) == expected
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_validate_request_ignores_content_type_on_get() -> None:
|
||||
"""Content-Type is only enforced for POST requests."""
|
||||
middleware = TransportSecurityMiddleware(SETTINGS)
|
||||
response = await middleware.validate_request(_request("good.example", None, content_type="text/plain"))
|
||||
assert response is None
|
||||
Loading…
Reference in a new issue