[v1.x] Bind transport sessions to the authenticated principal (#2719)
This commit is contained in:
Max
GitHub
parent
1abcca2408
commit
ce267b6fc5
2 changed files with 404 additions and 38 deletions
|
|
@ -1,9 +1,13 @@
|
|||
"""Tests for SSE server DNS rebinding protection."""
|
||||
"""Tests for SSE server request validation."""
|
||||
|
||||
import logging
|
||||
import multiprocessing
|
||||
import re
|
||||
import socket
|
||||
from collections.abc import Iterator
|
||||
from typing import Any
|
||||
|
||||
import anyio
|
||||
import httpx
|
||||
import pytest
|
||||
import uvicorn
|
||||
|
|
@ -11,8 +15,11 @@ 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
|
||||
|
||||
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
|
||||
|
|
@ -22,6 +29,23 @@ logger = logging.getLogger(__name__)
|
|||
SERVER_NAME = "test_sse_security_server"
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def reset_sse_starlette_exit_event() -> Iterator[None]:
|
||||
"""sse-starlette<2 caches a module-level anyio.Event on AppStatus; clear it
|
||||
around each test so it is never bound to a closed event loop. Clearing it
|
||||
afterwards matters too: later test modules fork uvicorn subprocesses on
|
||||
Linux and would otherwise inherit a stale event."""
|
||||
from sse_starlette.sse import AppStatus
|
||||
|
||||
def clear() -> None:
|
||||
if hasattr(AppStatus, "should_exit_event"): # pragma: no cover
|
||||
setattr(AppStatus, "should_exit_event", None)
|
||||
|
||||
clear()
|
||||
yield
|
||||
clear()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def server_port() -> int:
|
||||
with socket.socket() as s:
|
||||
|
|
@ -291,3 +315,112 @@ 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) -> dict[str, Any]:
|
||||
"""Build an ASGI scope for a request to the SSE transport."""
|
||||
scope: dict[str, Any] = {
|
||||
"type": "http",
|
||||
"method": method,
|
||||
"path": path,
|
||||
"root_path": "",
|
||||
"query_string": b"",
|
||||
"headers": [(b"content-type", b"application/json")],
|
||||
}
|
||||
if user is not None:
|
||||
scope["user"] = user
|
||||
return scope
|
||||
|
||||
|
||||
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 = _sse_scope("POST", "/messages/", user)
|
||||
scope["query_string"] = f"session_id={session_id}".encode()
|
||||
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)
|
||||
|
||||
await transport.handle_post_message(scope, receive, send)
|
||||
response_start = next(msg for msg in sent if msg["type"] == "http.response.start")
|
||||
return response_start["status"]
|
||||
|
||||
|
||||
_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:
|
||||
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
|
||||
|
|
|
|||
|
|
@ -6,9 +6,11 @@ from unittest.mock import AsyncMock, patch
|
|||
|
||||
import anyio
|
||||
import pytest
|
||||
from starlette.types import Message
|
||||
from starlette.types import Message, Scope
|
||||
|
||||
from mcp.server import streamable_http_manager
|
||||
from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser
|
||||
from mcp.server.auth.provider import AccessToken
|
||||
from mcp.server.lowlevel import Server
|
||||
from mcp.server.streamable_http import MCP_SESSION_ID_HEADER, StreamableHTTPServerTransport
|
||||
from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
|
||||
|
|
@ -390,3 +392,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
|
||||
|
|
|
|||
Loading…
Reference in a new issue