Apply the request body limit to the SSE and OAuth endpoints (#3336)
This commit is contained in:
Max
GitHub
parent
0cee6249ba
commit
57394b0548
5 changed files with 420 additions and 189 deletions
|
|
@ -16,6 +16,7 @@ from starlette.applications import Starlette
|
|||
from mcp.server.auth.provider import AuthorizeError, RegistrationError, TokenError
|
||||
from mcp.server.auth.routes import create_auth_routes
|
||||
from mcp.server.auth.settings import ClientRegistrationOptions, RevocationOptions
|
||||
from mcp.server.transport_security import DEFAULT_MAX_REQUEST_BODY_SIZE
|
||||
from tests.server.mcpserver.auth.test_auth_integration import MockOAuthProvider
|
||||
|
||||
|
||||
|
|
@ -288,3 +289,61 @@ async def test_token_error_handling_refresh_token(
|
|||
data = refresh_response.json()
|
||||
assert data["error"] == "invalid_scope"
|
||||
assert data["error_description"] == "The requested scope is invalid"
|
||||
|
||||
|
||||
_FORM = "application/x-www-form-urlencoded"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
@pytest.mark.parametrize(
|
||||
("method", "path", "content_type"),
|
||||
[
|
||||
("POST", "/token", _FORM),
|
||||
("POST", "/revoke", _FORM),
|
||||
("POST", "/register", "application/json"),
|
||||
("POST", "/authorize", _FORM),
|
||||
# The other methods these routes accept reach the same body-reading handlers.
|
||||
("OPTIONS", "/token", _FORM),
|
||||
("OPTIONS", "/revoke", _FORM),
|
||||
("OPTIONS", "/register", "application/json"),
|
||||
("HEAD", "/authorize", _FORM),
|
||||
],
|
||||
)
|
||||
async def test_oversized_request_body_returns_413(
|
||||
client: httpx2.AsyncClient, method: str, path: str, content_type: str
|
||||
):
|
||||
"""Each endpoint that reads a request body rejects one over 4 MiB before parsing it, whatever the method."""
|
||||
response = await client.request(
|
||||
method, path, content=b"x" * (DEFAULT_MAX_REQUEST_BODY_SIZE + 1), headers={"Content-Type": content_type}
|
||||
)
|
||||
assert response.status_code == 413
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_request_body_within_the_limit_is_still_parsed(client: httpx2.AsyncClient):
|
||||
"""A small body is passed through to the handler intact: the form is parsed and its fields validated."""
|
||||
response = await client.post("/token", data={"grant_type": "authorization_code"})
|
||||
assert response.status_code == 401
|
||||
assert response.json() == {"error": "invalid_client", "error_description": "Missing client_id"}
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_cors_preflight_is_still_answered(client: httpx2.AsyncClient):
|
||||
"""A CORS preflight to a body-limited endpoint is answered by the CORS layer as before."""
|
||||
response = await client.options(
|
||||
"/token", headers={"Origin": "https://client.example.com", "Access-Control-Request-Method": "POST"}
|
||||
)
|
||||
assert response.status_code == 200
|
||||
assert response.headers["access-control-allow-origin"] == "*"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_oversized_cross_origin_request_gets_413_with_cors_headers(client: httpx2.AsyncClient):
|
||||
"""The 413 is produced inside the CORS layer, so a browser client can still read it."""
|
||||
response = await client.post(
|
||||
"/token",
|
||||
content=b"x" * (DEFAULT_MAX_REQUEST_BODY_SIZE + 1),
|
||||
headers={"Content-Type": _FORM, "Origin": "https://client.example.com"},
|
||||
)
|
||||
assert response.status_code == 413
|
||||
assert response.headers["access-control-allow-origin"] == "*"
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ from typing import Any
|
|||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import anyio
|
||||
import httpx2
|
||||
import pytest
|
||||
from inline_snapshot import snapshot
|
||||
from mcp_types import (
|
||||
|
|
@ -1812,6 +1813,19 @@ def test_streamable_http_no_redirect() -> None:
|
|||
assert streamable_routes[0].path == "/mcp", "Streamable route path should be /mcp"
|
||||
|
||||
|
||||
async def test_sse_app_applies_the_configured_request_body_limit() -> None:
|
||||
"""`sse_app(max_request_body_size=...)` rejects larger POSTs to the message endpoint with HTTP 413."""
|
||||
app = MCPServer("test").sse_app(max_request_body_size=8, host="0.0.0.0")
|
||||
transport = httpx2.ASGITransport(app=app)
|
||||
async with httpx2.AsyncClient(transport=transport, base_url="http://localhost") as http:
|
||||
response = await http.post(
|
||||
"/messages/?session_id=12345678123456781234567812345678",
|
||||
content=b"123456789",
|
||||
headers={"Content-Type": "application/json"},
|
||||
)
|
||||
assert response.status_code == 413
|
||||
|
||||
|
||||
async def test_report_progress_delegates_to_session_report_progress():
|
||||
"""Context.report_progress delegates to ServerSession.report_progress unconditionally.
|
||||
|
||||
|
|
|
|||
|
|
@ -18,7 +18,7 @@ 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.server.transport_security import DEFAULT_MAX_REQUEST_BODY_SIZE, TransportSecuritySettings
|
||||
from mcp.shared._stream_protocols import WriteStream
|
||||
from mcp.shared.message import SessionMessage
|
||||
from tests.interaction.transports import StreamingASGITransport
|
||||
|
|
@ -204,9 +204,18 @@ def _authenticated_user(client_id: str, subject: str | None = None, issuer: str
|
|||
|
||||
|
||||
def _sse_scope(
|
||||
method: str, path: str, user: AuthenticatedUser | None, *, query_string: bytes = b"", body: bytes = b""
|
||||
method: str,
|
||||
path: str,
|
||||
user: AuthenticatedUser | None,
|
||||
*,
|
||||
query_string: bytes = b"",
|
||||
body: bytes | list[bytes] = b"",
|
||||
) -> tuple[Scope, Receive, Send, list[Message]]:
|
||||
"""Build an ASGI scope/receive/send triple for a request to the SSE transport."""
|
||||
"""Build an ASGI scope/receive/send triple for a request to the SSE transport.
|
||||
|
||||
`body` may be a list of chunks to deliver the request body over several `http.request` messages;
|
||||
no Content-Length header is set either way.
|
||||
"""
|
||||
scope: Scope = {
|
||||
"type": "http",
|
||||
"method": method,
|
||||
|
|
@ -218,9 +227,11 @@ def _sse_scope(
|
|||
if user is not None:
|
||||
scope["user"] = user
|
||||
sent: list[Message] = []
|
||||
chunks = list(body) if isinstance(body, list) else [body]
|
||||
|
||||
async def receive() -> Message:
|
||||
return {"type": "http.request", "body": body, "more_body": False}
|
||||
chunk = chunks.pop(0)
|
||||
return {"type": "http.request", "body": chunk, "more_body": bool(chunks)}
|
||||
|
||||
async def send(message: Message) -> None:
|
||||
sent.append(message)
|
||||
|
|
@ -233,6 +244,10 @@ def _response_status(sent: list[Message]) -> int:
|
|||
return response_start["status"]
|
||||
|
||||
|
||||
def _response_body(sent: list[Message]) -> bytes:
|
||||
return b"".join(msg.get("body", b"") for msg in sent if msg["type"] == "http.response.body")
|
||||
|
||||
|
||||
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}'
|
||||
|
|
@ -368,6 +383,80 @@ async def test_sse_post_with_a_disallowed_host_is_rejected_before_session_lookup
|
|||
assert _response_status(sent) == 421
|
||||
|
||||
|
||||
# A well-formed session ID that no live session owns.
|
||||
_UNKNOWN_SESSION = b"session_id=12345678123456781234567812345678"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_sse_post_body_over_the_limit_returns_413():
|
||||
"""A POST body larger than max_request_body_size is answered with 413 before any session handling."""
|
||||
transport = SseServerTransport("/messages/", max_request_body_size=8)
|
||||
scope, receive, send, sent = _sse_scope(
|
||||
"POST", "/messages/", None, query_string=_UNKNOWN_SESSION, body=b"123456789"
|
||||
)
|
||||
|
||||
await transport.handle_post_message(scope, receive, send)
|
||||
assert _response_status(sent) == 413
|
||||
assert _response_body(sent) == b"Request body too large"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_sse_post_body_limit_defaults_to_four_mib():
|
||||
"""Without an explicit limit, a body one byte over 4 MiB (and no Content-Length) is answered with 413."""
|
||||
transport = SseServerTransport("/messages/")
|
||||
body = b"x" * (DEFAULT_MAX_REQUEST_BODY_SIZE + 1)
|
||||
scope, receive, send, sent = _sse_scope("POST", "/messages/", None, query_string=_UNKNOWN_SESSION, body=body)
|
||||
|
||||
await transport.handle_post_message(scope, receive, send)
|
||||
assert _response_status(sent) == 413
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_sse_post_streamed_body_over_the_limit_returns_413():
|
||||
"""The limit counts bytes across body chunks, not just a declared Content-Length."""
|
||||
transport = SseServerTransport("/messages/", max_request_body_size=8)
|
||||
scope, receive, send, sent = _sse_scope(
|
||||
"POST", "/messages/", None, query_string=_UNKNOWN_SESSION, body=[b"1234", b"56789"]
|
||||
)
|
||||
|
||||
await transport.handle_post_message(scope, receive, send)
|
||||
assert _response_status(sent) == 413
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_sse_post_within_the_limit_reaches_session_lookup():
|
||||
"""A body within the limit is passed on intact: an unknown session still gets its 404."""
|
||||
transport = SseServerTransport("/messages/", max_request_body_size=64)
|
||||
scope, receive, send, sent = _sse_scope(
|
||||
"POST", "/messages/", None, query_string=_UNKNOWN_SESSION, body=[b'{"jsonrpc": ', b'"2.0"}']
|
||||
)
|
||||
|
||||
await transport.handle_post_message(scope, receive, send)
|
||||
assert _response_status(sent) == 404
|
||||
assert _response_body(sent) == b"Could not find session"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
@pytest.mark.parametrize("method", ["GET", "PUT"])
|
||||
async def test_sse_message_endpoint_answers_405_to_non_post(method: str):
|
||||
"""The message endpoint only accepts POST; other methods get 405 with an Allow header."""
|
||||
transport = SseServerTransport("/messages/")
|
||||
scope, receive, send, sent = _sse_scope(method, "/messages/", None, query_string=_UNKNOWN_SESSION, body=b"{}")
|
||||
|
||||
await transport.handle_post_message(scope, receive, send)
|
||||
assert _response_status(sent) == 405
|
||||
response_start = next(msg for msg in sent if msg["type"] == "http.response.start")
|
||||
assert (b"allow", b"POST") in response_start["headers"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("max_request_body_size", [0, -1])
|
||||
def test_sse_transport_rejects_a_non_positive_body_limit(max_request_body_size: int):
|
||||
"""The body limit must be a positive number of bytes, matching StreamableHTTPSessionManager."""
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
SseServerTransport("/messages/", max_request_body_size=max_request_body_size)
|
||||
assert str(exc_info.value) == "max_request_body_size must be a positive number of bytes"
|
||||
|
||||
|
||||
@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
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ import anyio
|
|||
import httpx2
|
||||
import pytest
|
||||
from mcp_types import INVALID_REQUEST, ListToolsResult, PaginatedRequestParams
|
||||
from starlette.types import Message, Receive, Scope, Send
|
||||
from starlette.types import Message, Scope
|
||||
|
||||
from mcp import Client
|
||||
from mcp.client.streamable_http import streamable_http_client
|
||||
|
|
@ -18,11 +18,7 @@ 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 (
|
||||
DEFAULT_MAX_REQUEST_BODY_SIZE,
|
||||
RequestBodyLimitMiddleware,
|
||||
StreamableHTTPSessionManager,
|
||||
)
|
||||
from mcp.server.streamable_http_manager import DEFAULT_MAX_REQUEST_BODY_SIZE, StreamableHTTPSessionManager
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
|
|
@ -146,79 +142,6 @@ async def test_oversized_streamed_body_is_rejected_before_session_creation(
|
|||
assert response_start["status"] == 413
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_client_disconnect_while_streaming_request_body_is_replayed() -> None:
|
||||
"""SDK-defined: raw ASGI is required to prove a disconnect before body completion reaches the transport."""
|
||||
disconnect: Message = {"type": "http.disconnect"}
|
||||
request_messages: Iterator[Message] = iter(
|
||||
[{"type": "http.request", "body": b"1234", "more_body": True}, disconnect]
|
||||
)
|
||||
received_messages: list[Message] = []
|
||||
|
||||
async def receive() -> Message:
|
||||
return next(request_messages)
|
||||
|
||||
async def app(scope: Scope, receive: Receive, send: Send) -> None:
|
||||
received_messages.append(await receive())
|
||||
received_messages.append(await receive())
|
||||
|
||||
scope: Scope = {"type": "http", "method": "POST", "path": "/mcp", "headers": []}
|
||||
middleware = RequestBodyLimitMiddleware(app, max_body_size=8)
|
||||
|
||||
await middleware(scope, receive, AsyncMock())
|
||||
|
||||
assert received_messages == [
|
||||
{"type": "http.request", "body": b"1234", "more_body": True},
|
||||
disconnect,
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_client_disconnect_before_request_body_is_replayed() -> None:
|
||||
"""SDK-defined: raw ASGI proves a disconnect before the first body message reaches the transport."""
|
||||
disconnect: Message = {"type": "http.disconnect"}
|
||||
received_messages: list[Message] = []
|
||||
|
||||
async def receive() -> Message:
|
||||
return disconnect
|
||||
|
||||
async def app(scope: Scope, receive: Receive, send: Send) -> None:
|
||||
received_messages.append(await receive())
|
||||
|
||||
scope: Scope = {"type": "http", "method": "POST", "path": "/mcp", "headers": []}
|
||||
middleware = RequestBodyLimitMiddleware(app, max_body_size=8)
|
||||
|
||||
await middleware(scope, receive, AsyncMock())
|
||||
|
||||
assert received_messages == [disconnect]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_request_body_chunks_are_replayed_as_one_message() -> None:
|
||||
"""SDK-defined: raw ASGI proves chunk overhead is discarded before the body reaches the transport."""
|
||||
request_messages: Iterator[Message] = iter(
|
||||
[
|
||||
{"type": "http.request", "body": b"12", "more_body": True},
|
||||
{"type": "http.request", "body": b"34", "more_body": True},
|
||||
{"type": "http.request", "body": b"56", "more_body": False},
|
||||
]
|
||||
)
|
||||
received_messages: list[Message] = []
|
||||
|
||||
async def receive() -> Message:
|
||||
return next(request_messages)
|
||||
|
||||
async def app(scope: Scope, receive: Receive, send: Send) -> None:
|
||||
received_messages.append(await receive())
|
||||
|
||||
scope: Scope = {"type": "http", "method": "POST", "path": "/mcp", "headers": []}
|
||||
middleware = RequestBodyLimitMiddleware(app, max_body_size=8)
|
||||
|
||||
await middleware(scope, receive, AsyncMock())
|
||||
|
||||
assert received_messages == [{"type": "http.request", "body": b"123456", "more_body": False}]
|
||||
|
||||
|
||||
def test_request_body_limit_defaults_to_four_mib() -> None:
|
||||
"""SDK-defined: Streamable HTTP request bodies are limited to 4 MiB by default."""
|
||||
manager = StreamableHTTPSessionManager(app=Server("test-default-size-limit"))
|
||||
|
|
|
|||
|
|
@ -1,9 +1,17 @@
|
|||
"""Tests for the transport-security request validation middleware."""
|
||||
"""Tests for the request checks shared by the HTTP server transports."""
|
||||
|
||||
from collections.abc import Iterator
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
from starlette.requests import Request
|
||||
from starlette.types import Message, Receive, Scope, Send
|
||||
|
||||
from mcp.server.transport_security import TransportSecurityMiddleware, TransportSecuritySettings
|
||||
from mcp.server.transport_security import (
|
||||
RequestBodyLimitMiddleware,
|
||||
TransportSecurityMiddleware,
|
||||
TransportSecuritySettings,
|
||||
)
|
||||
|
||||
|
||||
def _request(host: str | None, origin: str | None, content_type: str | None = "application/json") -> Request:
|
||||
|
|
@ -86,3 +94,111 @@ async def test_validate_request_ignores_content_type_on_get() -> None:
|
|||
middleware = TransportSecurityMiddleware(SETTINGS)
|
||||
response = await middleware.validate_request(_request("good.example", None, content_type="text/plain"))
|
||||
assert response is None
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_client_disconnect_while_streaming_request_body_is_replayed() -> None:
|
||||
"""SDK-defined: raw ASGI is required to prove a disconnect before body completion reaches the transport."""
|
||||
disconnect: Message = {"type": "http.disconnect"}
|
||||
request_messages: Iterator[Message] = iter(
|
||||
[{"type": "http.request", "body": b"1234", "more_body": True}, disconnect]
|
||||
)
|
||||
received_messages: list[Message] = []
|
||||
|
||||
async def receive() -> Message:
|
||||
return next(request_messages)
|
||||
|
||||
async def app(scope: Scope, receive: Receive, send: Send) -> None:
|
||||
received_messages.append(await receive())
|
||||
received_messages.append(await receive())
|
||||
|
||||
scope: Scope = {"type": "http", "method": "POST", "path": "/mcp", "headers": []}
|
||||
middleware = RequestBodyLimitMiddleware(app, max_body_size=8)
|
||||
|
||||
await middleware(scope, receive, AsyncMock())
|
||||
|
||||
assert received_messages == [
|
||||
{"type": "http.request", "body": b"1234", "more_body": True},
|
||||
disconnect,
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_client_disconnect_before_request_body_is_replayed() -> None:
|
||||
"""SDK-defined: raw ASGI proves a disconnect before the first body message reaches the transport."""
|
||||
disconnect: Message = {"type": "http.disconnect"}
|
||||
received_messages: list[Message] = []
|
||||
|
||||
async def receive() -> Message:
|
||||
return disconnect
|
||||
|
||||
async def app(scope: Scope, receive: Receive, send: Send) -> None:
|
||||
received_messages.append(await receive())
|
||||
|
||||
scope: Scope = {"type": "http", "method": "POST", "path": "/mcp", "headers": []}
|
||||
middleware = RequestBodyLimitMiddleware(app, max_body_size=8)
|
||||
|
||||
await middleware(scope, receive, AsyncMock())
|
||||
|
||||
assert received_messages == [disconnect]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_request_body_chunks_are_replayed_as_one_message() -> None:
|
||||
"""SDK-defined: raw ASGI proves chunk overhead is discarded before the body reaches the transport."""
|
||||
request_messages: Iterator[Message] = iter(
|
||||
[
|
||||
{"type": "http.request", "body": b"12", "more_body": True},
|
||||
{"type": "http.request", "body": b"34", "more_body": True},
|
||||
{"type": "http.request", "body": b"56", "more_body": False},
|
||||
]
|
||||
)
|
||||
received_messages: list[Message] = []
|
||||
|
||||
async def receive() -> Message:
|
||||
return next(request_messages)
|
||||
|
||||
async def app(scope: Scope, receive: Receive, send: Send) -> None:
|
||||
received_messages.append(await receive())
|
||||
|
||||
scope: Scope = {"type": "http", "method": "POST", "path": "/mcp", "headers": []}
|
||||
middleware = RequestBodyLimitMiddleware(app, max_body_size=8)
|
||||
|
||||
await middleware(scope, receive, AsyncMock())
|
||||
|
||||
assert received_messages == [{"type": "http.request", "body": b"123456", "more_body": False}]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
@pytest.mark.parametrize("method", ["GET", "PUT", "OPTIONS", "HEAD", "DELETE"])
|
||||
async def test_request_body_limit_applies_to_every_method(method: str) -> None:
|
||||
"""SDK-defined: the limit is a property of the request body, not of the method that carries it."""
|
||||
app = AsyncMock()
|
||||
sent_messages: list[Message] = []
|
||||
receive = AsyncMock(return_value={"type": "http.request", "body": b"123456789", "more_body": False})
|
||||
|
||||
async def send(message: Message) -> None:
|
||||
sent_messages.append(message)
|
||||
|
||||
scope: Scope = {"type": "http", "method": method, "path": "/mcp", "headers": []}
|
||||
middleware = RequestBodyLimitMiddleware(app, max_body_size=8)
|
||||
|
||||
await middleware(scope, receive, send)
|
||||
|
||||
assert [message["status"] for message in sent_messages if message["type"] == "http.response.start"] == [413]
|
||||
app.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_request_body_limit_leaves_non_http_scopes_alone() -> None:
|
||||
"""SDK-defined: only HTTP requests carry a body to limit; other ASGI scopes go straight to the app."""
|
||||
app = AsyncMock()
|
||||
receive = AsyncMock()
|
||||
send = AsyncMock()
|
||||
scope: Scope = {"type": "lifespan"}
|
||||
middleware = RequestBodyLimitMiddleware(app, max_body_size=8)
|
||||
|
||||
await middleware(scope, receive, send)
|
||||
|
||||
app.assert_awaited_once_with(scope, receive, send)
|
||||
receive.assert_not_awaited()
|
||||
|
|
|
|||
Loading…
Reference in a new issue