[v1.x] Bind transport sessions to the authenticated principal (#2719)

tests/server/test_sse_security.py

@@ -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