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

This commit is contained in:
Max 2026-05-29 17:46:37 +01:00 • committed by GitHub
commit ce267b6fc5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 404 additions and 38 deletions

View file

@ -1,6 +1,6 @@
import json
import time
from typing import Any
from typing import Any, TypedDict
from pydantic import AnyHttpUrl
from starlette.authentication import AuthCredentials, AuthenticationBackend, SimpleUser
@ -19,6 +19,30 @@ class AuthenticatedUser(SimpleUser):
self.scopes = auth_info.scopes
class AuthorizationContext(TypedDict):
client_id: str
issuer: str | None
subject: str | None
def authorization_context(user: AuthenticatedUser) -> AuthorizationContext:
"""Identify the principal `user` represents, for transports to compare
against the principal that created a session. Components the token
verifier does not supply are `None`, so the comparison degrades to the
remaining components.
See `examples/servers/simple-auth/mcp_simple_auth/token_verifier.py` for
a verifier that populates `subject` and `claims` from an introspection
response."""
token = user.access_token
issuer = (token.claims or {}).get("iss")
return AuthorizationContext(
client_id=token.client_id,
issuer=str(issuer) if issuer is not None else None,
subject=token.subject,
)
class BearerAuthBackend(AuthenticationBackend):
"""
Authentication backend that validates Bearer tokens using a TokenVerifier.

View file

@ -52,6 +52,7 @@ from starlette.responses import Response
from starlette.types import Receive, Scope, Send
import mcp.types as types
from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser, AuthorizationContext, authorization_context
from mcp.server.transport_security import (
TransportSecurityMiddleware,
TransportSecuritySettings,
@ -75,6 +76,9 @@ class SseServerTransport:
_endpoint: str
_read_stream_writers: dict[UUID, MemoryObjectSendStream[SessionMessage | Exception]]
# Identity of the credential that created each session; requests for a
# session must present the same credential.
_session_owners: dict[UUID, AuthorizationContext]
_security: TransportSecurityMiddleware
def __init__(self, endpoint: str, security_settings: TransportSecuritySettings | None = None) -> None:
@ -115,6 +119,7 @@ class SseServerTransport:
self._endpoint = endpoint
self._read_stream_writers = {}
self._session_owners = {}
self._security = TransportSecurityMiddleware(security_settings)
logger.debug(f"SseServerTransport initialized with endpoint: {endpoint}")
@ -142,6 +147,9 @@ class SseServerTransport:
write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
session_id = uuid4()
user = scope.get("user")
if isinstance(user, AuthenticatedUser):
self._session_owners[session_id] = authorization_context(user)
self._read_stream_writers[session_id] = read_stream_writer
logger.debug(f"Created new session with ID: {session_id}")
@ -177,26 +185,34 @@ class SseServerTransport:
}
)
async with anyio.create_task_group() as tg:
try:
async with anyio.create_task_group() as tg:
async def response_wrapper(scope: Scope, receive: Receive, send: Send):
"""
The EventSourceResponse returning signals a client close / disconnect.
In this case we close our side of the streams to signal the client that
the connection has been closed.
"""
await EventSourceResponse(content=sse_stream_reader, data_sender_callable=sse_writer)(
scope, receive, send
)
await read_stream_writer.aclose()
await write_stream_reader.aclose()
logging.debug(f"Client session disconnected {session_id}")
async def response_wrapper(scope: Scope, receive: Receive, send: Send):
"""
The EventSourceResponse returning signals a client close / disconnect.
In this case we close our side of the streams to signal the client that
the connection has been closed.
"""
await EventSourceResponse(content=sse_stream_reader, data_sender_callable=sse_writer)(
scope, receive, send
)
await read_stream_writer.aclose()
await write_stream_reader.aclose()
await sse_stream_reader.aclose()
logging.debug(f"Client session disconnected {session_id}")
logger.debug("Starting SSE response task")
tg.start_soon(response_wrapper, scope, receive, send)
logger.debug("Starting SSE response task")
tg.start_soon(response_wrapper, scope, receive, send)
logger.debug("Yielding read and write streams")
yield (read_stream, write_stream)
logger.debug("Yielding read and write streams")
yield (read_stream, write_stream)
finally:
# The connection is gone: stop routing messages to this session
# and drop its entries so they do not accumulate for the lifetime
# of the transport.
self._read_stream_writers.pop(session_id, None)
self._session_owners.pop(session_id, None)
async def handle_post_message(self, scope: Scope, receive: Receive, send: Send) -> None: # pragma: no cover
logger.debug("Handling POST message")
@ -227,6 +243,15 @@ class SseServerTransport:
response = Response("Could not find session", status_code=404)
return await response(scope, receive, send)
user = scope.get("user")
requestor = authorization_context(user) if isinstance(user, AuthenticatedUser) else None
if requestor != self._session_owners.get(session_id):
# A session can only be used with the credential that created it.
# Respond exactly as if the session did not exist.
logger.warning("Rejecting message for session %s: credential does not match", session_id)
response = Response("Could not find session", status_code=404)
return await response(scope, receive, send)
body = await request.body()
logger.debug(f"Received JSON: {body}")

View file

@ -5,7 +5,6 @@ from __future__ import annotations
import contextlib
import logging
from collections.abc import AsyncIterator
from http import HTTPStatus
from typing import Any
from uuid import uuid4
@ -15,6 +14,7 @@ from starlette.requests import Request
from starlette.responses import Response
from starlette.types import Receive, Scope, Send
from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser, AuthorizationContext, authorization_context
from mcp.server.lowlevel.server import Server as MCPServer
from mcp.server.streamable_http import (
MCP_SESSION_ID_HEADER,
@ -88,6 +88,9 @@ class StreamableHTTPSessionManager:
# Session tracking (only used if not stateless)
self._session_creation_lock = anyio.Lock()
self._server_instances: dict[str, StreamableHTTPServerTransport] = {}
# Identity of the credential that created each session; requests for a
# session must present the same credential.
self._session_owners: dict[str, AuthorizationContext] = {}
# The task group will be set during lifespan
self._task_group = None
@ -135,6 +138,7 @@ class StreamableHTTPSessionManager:
self._task_group = None
# Clear any remaining server instances
self._server_instances.clear()
self._session_owners.clear()
async def handle_request(
self,
@ -227,12 +231,32 @@ class StreamableHTTPSessionManager:
request = Request(scope, receive)
request_mcp_session_id = request.headers.get(MCP_SESSION_ID_HEADER)
user = scope.get("user")
requestor = authorization_context(user) if isinstance(user, AuthenticatedUser) else None
# Existing session case
if request_mcp_session_id is not None and request_mcp_session_id in self._server_instances: # pragma: no cover
if request_mcp_session_id is not None and request_mcp_session_id in self._server_instances:
transport = self._server_instances[request_mcp_session_id]
if requestor != self._session_owners.get(request_mcp_session_id):
# A session can only be used with the credential that created
# it. Respond exactly as if the session did not exist.
logger.warning(
"Rejecting request for session %s: credential does not match the one that created the session",
request_mcp_session_id[:64],
)
body = JSONRPCError(
jsonrpc="2.0", id="server-error", error=ErrorData(code=INVALID_REQUEST, message="Session not found")
)
response = Response(
body.model_dump_json(by_alias=True, exclude_none=True),
status_code=404,
media_type="application/json",
)
await response(scope, receive, send)
return
logger.debug("Session already exists, handling request directly")
# Push back idle deadline on activity
if transport.idle_scope is not None and self.session_idle_timeout is not None:
if transport.idle_scope is not None and self.session_idle_timeout is not None: # pragma: no cover
transport.idle_scope.deadline = anyio.current_time() + self.session_idle_timeout
await transport.handle_request(scope, receive, send)
return
@ -251,6 +275,8 @@ class StreamableHTTPSessionManager:
)
assert http_transport.mcp_session_id is not None
if requestor is not None:
self._session_owners[http_transport.mcp_session_id] = requestor
self._server_instances[http_transport.mcp_session_id] = http_transport
logger.info(f"Created new transport with session ID: {new_session_id}")
@ -281,6 +307,7 @@ class StreamableHTTPSessionManager:
assert http_transport.mcp_session_id is not None
logger.info(f"Session {http_transport.mcp_session_id} idle timeout")
self._server_instances.pop(http_transport.mcp_session_id, None)
self._session_owners.pop(http_transport.mcp_session_id, None)
await http_transport.terminate()
except Exception:
logger.exception(f"Session {http_transport.mcp_session_id} crashed")
@ -296,6 +323,7 @@ class StreamableHTTPSessionManager:
"active instances."
)
del self._server_instances[http_transport.mcp_session_id]
self._session_owners.pop(http_transport.mcp_session_id, None)
# Assert task group is not None for type checking
assert self._task_group is not None
@ -306,19 +334,10 @@ class StreamableHTTPSessionManager:
await http_transport.handle_request(scope, receive, send)
else:
# Unknown or expired session ID - return 404 per MCP spec
# TODO: Align error code once spec clarifies
# See: https://github.com/modelcontextprotocol/python-sdk/issues/1821
error_response = JSONRPCError(
jsonrpc="2.0",
id="server-error",
error=ErrorData(
code=INVALID_REQUEST,
message="Session not found",
),
body = JSONRPCError(
jsonrpc="2.0", id="server-error", error=ErrorData(code=INVALID_REQUEST, message="Session not found")
)
response = Response(
content=error_response.model_dump_json(by_alias=True, exclude_none=True),
status_code=HTTPStatus.NOT_FOUND,
media_type="application/json",
body.model_dump_json(by_alias=True, exclude_none=True), status_code=404, media_type="application/json"
)
await response(scope, receive, send)