|
| 1 | +"""Hardened HTTP core for discovery fetches. Private module. |
| 2 | +
|
| 3 | +Every network fetch (including every redirect hop and every nested catalog |
| 4 | +follow) is re-admitted under the same rules: http(s) scheme only, plain http |
| 5 | +only to loopback hosts, an SSRF address guard checked before the request, |
| 6 | +bounded redirects, a streamed response size cap, and a media type check. |
| 7 | +""" |
| 8 | + |
| 9 | +import ipaddress |
| 10 | +from dataclasses import dataclass |
| 11 | +from typing import Literal |
| 12 | +from urllib.parse import urljoin, urlsplit |
| 13 | + |
| 14 | +import anyio |
| 15 | +import httpx2 |
| 16 | + |
| 17 | +from mcp.shared._httpx_utils import create_mcp_http_client |
| 18 | +from mcp.shared.experimental._base import is_loopback_host |
| 19 | +from mcp.shared.experimental.ai_catalog import MAX_CATALOG_NESTING_DEPTH |
| 20 | + |
| 21 | +DiscoveryErrorReason = Literal[ |
| 22 | + "status", |
| 23 | + "media_type", |
| 24 | + "response_too_large", |
| 25 | + "too_many_redirects", |
| 26 | + "blocked_address", |
| 27 | + "insecure_transport", |
| 28 | + "invalid_entry", |
| 29 | + "catalog_depth", |
| 30 | +] |
| 31 | + |
| 32 | +_REDIRECT_STATUSES = frozenset({301, 302, 303, 307, 308}) |
| 33 | +_CGNAT_NETWORK = ipaddress.ip_network("100.64.0.0/10") |
| 34 | + |
| 35 | + |
| 36 | +@dataclass(frozen=True, slots=True, kw_only=True) |
| 37 | +class DiscoveryPolicy: |
| 38 | + """Limits applied to every discovery fetch. |
| 39 | +
|
| 40 | + The defaults are the hardened production posture. Set |
| 41 | + `allow_private_addresses=True` for local development, which permits plain |
| 42 | + http and skips the address guard entirely. |
| 43 | +
|
| 44 | + The address guard resolves DNS before the request and re-checks every |
| 45 | + redirect hop, but the HTTP client re-resolves when it connects, so a |
| 46 | + DNS rebinding race remains possible between the check and the connect. |
| 47 | + """ |
| 48 | + |
| 49 | + max_response_bytes: int = 1_048_576 |
| 50 | + max_redirects: int = 3 |
| 51 | + max_catalog_entries: int = 100 |
| 52 | + max_catalog_depth: int = MAX_CATALOG_NESTING_DEPTH |
| 53 | + allow_private_addresses: bool = False |
| 54 | + timeout_seconds: float = 10.0 |
| 55 | + |
| 56 | + |
| 57 | +class DiscoveryError(Exception): |
| 58 | + """A transport or policy failure during discovery. |
| 59 | +
|
| 60 | + `reason` states which rule failed and `url` names the offending target. |
| 61 | + Document shape failures raise `pydantic.ValidationError` instead, and raw |
| 62 | + connection failures stay `httpx2.HTTPError`. |
| 63 | + """ |
| 64 | + |
| 65 | + url: str |
| 66 | + reason: DiscoveryErrorReason |
| 67 | + |
| 68 | + def __init__(self, message: str, *, url: str, reason: DiscoveryErrorReason) -> None: |
| 69 | + super().__init__(message) |
| 70 | + self.url = url |
| 71 | + self.reason = reason |
| 72 | + |
| 73 | + |
| 74 | +def check_status(response: httpx2.Response, url: str) -> None: |
| 75 | + """Reject any non-2xx response. |
| 76 | +
|
| 77 | + Raises: |
| 78 | + DiscoveryError: With `reason="status"`, chaining the underlying |
| 79 | + `httpx2.HTTPStatusError`. |
| 80 | + """ |
| 81 | + try: |
| 82 | + response.raise_for_status() |
| 83 | + except httpx2.HTTPStatusError as exc: |
| 84 | + raise DiscoveryError(f"unexpected status {response.status_code} from {url}", url=url, reason="status") from exc |
| 85 | + |
| 86 | + |
| 87 | +def check_media_type(response: httpx2.Response, url: str, expected: str) -> None: |
| 88 | + """Reject a response whose media type is neither `expected` nor JSON. |
| 89 | +
|
| 90 | + Plain `application/json` is accepted because static hosts and CDNs |
| 91 | + commonly serve it. |
| 92 | +
|
| 93 | + Raises: |
| 94 | + DiscoveryError: With `reason="media_type"`. |
| 95 | + """ |
| 96 | + content_type = response.headers.get("content-type", "") |
| 97 | + media_type = content_type.split(";", 1)[0].strip().lower() |
| 98 | + if media_type not in (expected, "application/json"): |
| 99 | + raise DiscoveryError(f"unexpected media type {content_type!r} from {url}", url=url, reason="media_type") |
| 100 | + |
| 101 | + |
| 102 | +async def _host_addresses(host: str) -> list[ipaddress.IPv4Address | ipaddress.IPv6Address]: |
| 103 | + """The IP addresses `host` names: the literal itself, or its DNS resolution.""" |
| 104 | + try: |
| 105 | + return [ipaddress.ip_address(host)] |
| 106 | + except ValueError: |
| 107 | + infos = await anyio.getaddrinfo(host, None) |
| 108 | + # sockaddr[0] is the address; IPv6 link-local entries carry a %scope suffix. |
| 109 | + return [ipaddress.ip_address(info[4][0].split("%", 1)[0]) for info in infos] |
| 110 | + |
| 111 | + |
| 112 | +def _is_blocked_address(address: ipaddress.IPv4Address | ipaddress.IPv6Address) -> bool: |
| 113 | + """Whether `address` is off-limits for discovery: anything not publicly routable.""" |
| 114 | + if address.version == 4 and address in _CGNAT_NETWORK: |
| 115 | + return True |
| 116 | + # is_private covers loopback, RFC 1918, ULA and the unspecified address. |
| 117 | + return address.is_private or address.is_link_local or address.is_multicast or address.is_reserved |
| 118 | + |
| 119 | + |
| 120 | +async def _admit_url(url: str, policy: DiscoveryPolicy) -> None: |
| 121 | + """Apply the scheme rule and the SSRF address guard to one fetch target. |
| 122 | +
|
| 123 | + Raises: |
| 124 | + DiscoveryError: `reason="insecure_transport"` for a non-http(s) URL or |
| 125 | + plain http to a non-loopback host, `reason="blocked_address"` when |
| 126 | + the host is (or resolves to) a non-public address. |
| 127 | + """ |
| 128 | + parts = urlsplit(url) |
| 129 | + host = parts.hostname |
| 130 | + if parts.scheme not in ("http", "https") or not host: |
| 131 | + raise DiscoveryError( |
| 132 | + f"discovery requires an absolute http(s) URL, got {url!r}", url=url, reason="insecure_transport" |
| 133 | + ) |
| 134 | + if policy.allow_private_addresses: |
| 135 | + return |
| 136 | + if parts.scheme == "http" and not is_loopback_host(host): |
| 137 | + raise DiscoveryError( |
| 138 | + f"plain http is only allowed for loopback hosts, got {url!r}", url=url, reason="insecure_transport" |
| 139 | + ) |
| 140 | + for address in await _host_addresses(host): |
| 141 | + if _is_blocked_address(address): |
| 142 | + raise DiscoveryError(f"{url!r} points at blocked address {address}", url=url, reason="blocked_address") |
| 143 | + |
| 144 | + |
| 145 | +async def _read_limited(response: httpx2.Response, url: str, policy: DiscoveryPolicy) -> bytes: |
| 146 | + """Stream the body, failing as soon as it exceeds the size cap.""" |
| 147 | + chunks: list[bytes] = [] |
| 148 | + total = 0 |
| 149 | + async for chunk in response.aiter_bytes(): |
| 150 | + total += len(chunk) |
| 151 | + if total > policy.max_response_bytes: |
| 152 | + raise DiscoveryError( |
| 153 | + f"response from {url} exceeds {policy.max_response_bytes} bytes", url=url, reason="response_too_large" |
| 154 | + ) |
| 155 | + chunks.append(chunk) |
| 156 | + return b"".join(chunks) |
| 157 | + |
| 158 | + |
| 159 | +async def _fetch(client: httpx2.AsyncClient, url: str, media_type: str, policy: DiscoveryPolicy) -> bytes: |
| 160 | + """One admitted fetch, walking redirects manually so each hop is re-checked.""" |
| 161 | + with anyio.fail_after(policy.timeout_seconds): |
| 162 | + current = url |
| 163 | + for _ in range(policy.max_redirects + 1): |
| 164 | + await _admit_url(current, policy) |
| 165 | + headers = {"Accept": f"{media_type}, application/json;q=0.5"} |
| 166 | + request = client.build_request("GET", current, headers=headers) |
| 167 | + response = await client.send(request, stream=True, follow_redirects=False) |
| 168 | + try: |
| 169 | + if response.status_code in _REDIRECT_STATUSES: |
| 170 | + location = response.headers.get("location") |
| 171 | + if location is None: |
| 172 | + raise DiscoveryError( |
| 173 | + f"redirect from {current} carries no Location", url=current, reason="status" |
| 174 | + ) |
| 175 | + current = urljoin(current, location) |
| 176 | + continue |
| 177 | + check_status(response, current) |
| 178 | + check_media_type(response, current, media_type) |
| 179 | + return await _read_limited(response, current, policy) |
| 180 | + finally: |
| 181 | + await response.aclose() |
| 182 | + raise DiscoveryError( |
| 183 | + f"more than {policy.max_redirects} redirects while fetching {url}", url=current, reason="too_many_redirects" |
| 184 | + ) |
| 185 | + |
| 186 | + |
| 187 | +async def fetch_discovery_document( |
| 188 | + url: str, |
| 189 | + media_type: str, |
| 190 | + *, |
| 191 | + http_client: httpx2.AsyncClient | None = None, |
| 192 | + policy: DiscoveryPolicy | None = None, |
| 193 | +) -> bytes: |
| 194 | + """Fetch one discovery document under `policy`, returning its raw bytes. |
| 195 | +
|
| 196 | + With `http_client=None` a fresh client is created per call, so no cookies |
| 197 | + or ambient credentials ever accompany a discovery request. A caller-owned |
| 198 | + client still gets every URL, redirect, size and media type check. |
| 199 | +
|
| 200 | + Raises: |
| 201 | + DiscoveryError: When any policy rule fails. |
| 202 | + httpx2.HTTPError: For raw connection failures. |
| 203 | + OSError: When DNS resolution itself fails. |
| 204 | + TimeoutError: When the fetch exceeds `policy.timeout_seconds`. |
| 205 | + """ |
| 206 | + resolved_policy = policy if policy is not None else DiscoveryPolicy() |
| 207 | + if http_client is None: |
| 208 | + async with create_mcp_http_client() as own_client: |
| 209 | + return await _fetch(own_client, url, media_type, resolved_policy) |
| 210 | + return await _fetch(http_client, url, media_type, resolved_policy) |
0 commit comments