"""Request checks shared by the HTTP server transports: Host/Origin header validation and body size limits.""" import logging from collections import deque from typing import Final from pydantic import BaseModel, Field from starlette.datastructures import Headers from starlette.requests import HTTPConnection from starlette.responses import Response from starlette.types import ASGIApp, Message, Receive, Scope, Send logger = logging.getLogger(__name__) DEFAULT_MAX_REQUEST_BODY_SIZE: Final = 4 * 1024 * 1024 """Default maximum HTTP request body size in bytes (4 MiB).""" class TransportSecuritySettings(BaseModel): """Settings for MCP transport security features. These settings help protect against DNS rebinding attacks by validating incoming request headers. """ enable_dns_rebinding_protection: bool = Field( default=True, description="Enable DNS rebinding protection (recommended for production)", ) allowed_hosts: list[str] = Field( default=[], description="List of allowed Host header values. Only applies when " + "enable_dns_rebinding_protection is True.", ) allowed_origins: list[str] = Field( default=[], description="List of allowed Origin header values. Only applies when " + "enable_dns_rebinding_protection is True.", ) class TransportSecurityMiddleware: """Middleware to enforce DNS rebinding protection for MCP transport endpoints.""" def __init__(self, settings: TransportSecuritySettings | None = None): # If not specified, disable DNS rebinding protection by default # for backwards compatibility self.settings = settings or TransportSecuritySettings(enable_dns_rebinding_protection=False) def _validate_host(self, host: str | None) -> bool: # pragma: no cover """Validate the Host header against allowed values.""" if not host: logger.warning("Missing Host header in request") return False # Check exact match first if host in self.settings.allowed_hosts: return True # Check wildcard port patterns for allowed in self.settings.allowed_hosts: 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 if host.startswith(base_host + ":"): return True logger.warning(f"Invalid Host header: {host}") return False def _validate_origin(self, origin: str | None) -> bool: # pragma: no cover """Validate the Origin header against allowed values.""" # Origin can be absent for same-origin requests if not origin: return True # Check exact match first if origin in self.settings.allowed_origins: return True # Check wildcard port patterns for allowed in self.settings.allowed_origins: 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 if origin.startswith(base_origin + ":"): return True logger.warning(f"Invalid Origin header: {origin}") return False def _validate_content_type(self, content_type: str | None) -> bool: # pragma: no cover """Validate the Content-Type header for POST requests.""" if not content_type: logger.warning("Missing Content-Type header in POST request") return False # Content-Type must start with application/json if not content_type.lower().startswith("application/json"): logger.warning(f"Invalid Content-Type header: {content_type}") return False return True async def validate_request(self, request: HTTPConnection, is_post: bool = False) -> Response | None: """Validate request headers for DNS rebinding protection. 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 content_type = request.headers.get("content-type") if not self._validate_content_type(content_type): # pragma: no cover return Response("Invalid Content-Type header", status_code=400) # Skip remaining validation if DNS rebinding protection is disabled if not self.settings.enable_dns_rebinding_protection: return None # Validate Host header # pragma: no cover host = request.headers.get("host") # pragma: no cover if not self._validate_host(host): # pragma: no cover return Response("Invalid Host header", status_code=421) # pragma: no cover # Validate Origin header # pragma: no cover origin = request.headers.get("origin") # pragma: no cover if not self._validate_origin(origin): # pragma: no cover return Response("Invalid Origin header", status_code=403) # pragma: no cover return None # pragma: no cover class RequestBodyLimitMiddleware: """Reject oversized HTTP request bodies before invoking an ASGI application.""" def __init__(self, app: ASGIApp, max_body_size: int) -> None: self.app = app self.max_body_size = max_body_size async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: if scope["type"] != "http": await self.app(scope, receive, send) return headers = Headers(scope=scope) content_length = headers.get("content-length") if content_length is not None: try: declared_size = int(content_length) except ValueError: pass else: if declared_size > self.max_body_size: response = Response("Request body too large", status_code=413) return await response(scope, receive, send) received_body = bytearray() received_request = False body_complete = False trailing_message: Message | None = None while True: message = await receive() if message["type"] != "http.request": trailing_message = message break received_request = True body = message.get("body", b"") if len(received_body) + len(body) > self.max_body_size: response = Response("Request body too large", status_code=413) return await response(scope, receive, send) received_body.extend(body) if not message.get("more_body", False): body_complete = True break cached_messages: deque[Message] = deque() if received_request: cached_messages.append( {"type": "http.request", "body": bytes(received_body), "more_body": not body_complete} ) if trailing_message is not None: cached_messages.append(trailing_message) async def replay() -> Message: if cached_messages: return cached_messages.popleft() return await receive() await self.app(scope, replay, send)