[v1.x] Bind transport sessions to the authenticated principal (#2719)
This commit is contained in:
parent
1abcca2408
commit
ce267b6fc5
3 changed files with 404 additions and 38 deletions
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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}")
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Reference in a new issue