modelcontextprotocol/python-sdk
Watch
1
0
Fork
You've already forked python-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Bind transport sessions to the authenticated principal (#2718)

Browse source
This commit is contained in:
Max 2026-05-29 17:48:40 +01:00 • committed by GitHub
commit 616476f692
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 647 additions and 42 deletions
Download patch file Download diff file Expand all files Collapse all files

26
src/mcp/server/auth/middleware/bearer_auth.py
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."""

64
src/mcp/server/sse.py
View file

@ -50,6 +50,7 @@ from starlette.responses import Response
from starlette.types import Receive, Scope, Send
from mcp import types
from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser, AuthorizationContext, authorization_context
from mcp.server.transport_security import (
TransportSecurityMiddleware,
TransportSecuritySettings,
@ -73,6 +74,9 @@ class SseServerTransport:
_endpoint: str
_read_stream_writers: dict[UUID, ContextSendStream[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:
@ -112,19 +116,20 @@ 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}")
@asynccontextmanager
async def connect_sse(self, scope: Scope, receive: Receive, send: Send):
if scope["type"] != "http": # pragma: no cover
if scope["type"] != "http":
logger.error("connect_sse received non-HTTP request")
raise ValueError("connect_sse can only handle HTTP requests")
# Validate request headers for DNS rebinding protection
request = Request(scope, receive)
error_response = await self._security.validate_request(request, is_post=False)
if error_response: # pragma: no cover
if error_response:
await error_response(scope, receive, send)
raise ValueError("Request validation failed")
@ -134,6 +139,9 @@ class SseServerTransport:
write_stream, write_stream_reader = create_context_streams[SessionMessage](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}")
@ -169,27 +177,30 @@ 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 sse_stream_reader.aclose()
await read_stream_writer.aclose()
await write_stream_reader.aclose()
self._read_stream_writers.pop(session_id, None)
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:
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:
logger.debug("Handling POST message")
@ -197,7 +208,7 @@ class SseServerTransport:
# Validate request headers for DNS rebinding protection
error_response = await self._security.validate_request(request, is_post=True)
if error_response: # pragma: no cover
if error_response:
return await error_response(scope, receive, send)
session_id_param = request.query_params.get("session_id")
@ -220,13 +231,22 @@ 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}")
try:
message = types.jsonrpc_message_adapter.validate_json(body, by_name=False)
logger.debug(f"Validated client message: {message}")
except ValidationError as err: # pragma: no cover
except ValidationError as err:
logger.exception("Failed to parse message")
response = Response("Could not parse message", status_code=400)
await response(scope, receive, send)

40
src/mcp/server/streamable_http_manager.py
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 TYPE_CHECKING, 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.streamable_http import (
MCP_SESSION_ID_HEADER,
EventStore,
@ -89,6 +89,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, scope: Scope, receive: Receive, send: Send) -> None:
"""Process ASGI request with proper session handling and transport setup.
@ -192,9 +196,29 @@ 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:
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=None, error=ErrorData(code=INVALID_REQUEST, message="Session not found")
)
response = Response(
body.model_dump_json(by_alias=True, exclude_unset=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:
@ -216,6 +240,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}")
@ -246,6 +272,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")
@ -260,6 +287,7 @@ class StreamableHTTPSessionManager:
f"{http_transport.mcp_session_id} from 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
@ -273,15 +301,11 @@ class StreamableHTTPSessionManager:
# TODO: Align error code once spec clarifies
# See: https://github.com/modelcontextprotocol/python-sdk/issues/1821
logger.info(f"Rejected request with unknown or expired session ID: {request_mcp_session_id[:64]}")
error_response = JSONRPCError(
jsonrpc="2.0",
id=None,
error=ErrorData(code=INVALID_REQUEST, message="Session not found"),
body = JSONRPCError(
jsonrpc="2.0", id=None, error=ErrorData(code=INVALID_REQUEST, message="Session not found")
)
response = Response(
content=error_response.model_dump_json(by_alias=True, exclude_unset=True),
status_code=HTTPStatus.NOT_FOUND,
media_type="application/json",
body.model_dump_json(by_alias=True, exclude_unset=True), status_code=404, media_type="application/json"
)
await response(scope, receive, send)

14
src/mcp/server/transport_security.py
View file

@ -42,17 +42,17 @@ class TransportSecurityMiddleware:
def _validate_host(self, host: str | None) -> bool:
"""Validate the Host header against allowed values."""
if not host: # pragma: no cover
if not host:
logger.warning("Missing Host header in request")
return False
# Check exact match first
if host in self.settings.allowed_hosts: # pragma: no cover
if host in self.settings.allowed_hosts:
return True
# Check wildcard port patterns
for allowed in self.settings.allowed_hosts:
if allowed.endswith(":*"): # pragma: no branch
if allowed.endswith(":*"):
# Extract base host from pattern
base_host = allowed[:-2]
# Check if the actual host starts with base host and has a port
@ -65,16 +65,16 @@ class TransportSecurityMiddleware:
def _validate_origin(self, origin: str | None) -> bool:
"""Validate the Origin header against allowed values."""
# Origin can be absent for same-origin requests
if not origin: # pragma: no cover
if not origin:
return True
# Check exact match first
if origin in self.settings.allowed_origins: # pragma: no cover
if origin in self.settings.allowed_origins:
return True
# Check wildcard port patterns
for allowed in self.settings.allowed_origins:
if allowed.endswith(":*"): # pragma: no branch
if allowed.endswith(":*"):
# Extract base origin from pattern
base_origin = allowed[:-2]
# Check if the actual origin starts with base origin and has a port
@ -94,7 +94,7 @@ class TransportSecurityMiddleware:
Returns None if validation passes, or an error Response if validation fails.
"""
# Always validate Content-Type for POST requests
if is_post: # pragma: no branch
if is_post:
content_type = request.headers.get("content-type")
if not self._validate_content_type(content_type):
return Response("Invalid Content-Type header", status_code=400)