1"""SSRF-safe httpx transport with DNS resolution and IP pinning."""23import asyncio4import socket56import httpx78from langchain_core._security._exceptions import SSRFBlockedError9from langchain_core._security._policy import (10    DEFAULT_SSRF_POLICY,11    SSRFPolicy,12    _effective_allowed_hosts,13    validate_resolved_ip,14    validate_url_sync,15)1617# Keys that AsyncHTTPTransport accepts (forwarded from factory kwargs).18_TRANSPORT_KWARGS = frozenset(19    {20        "verify",21        "cert",22        "trust_env",23        "http1",24        "http2",25        "limits",26        "retries",27    }28)293031class SSRFSafeTransport(httpx.AsyncBaseTransport):32    """httpx async transport that validates DNS results against an SSRF policy.3334    For every outgoing request the transport:35    1. Checks the URL scheme against `policy.allowed_schemes`.36    2. Validates the hostname against blocked patterns.37    3. Resolves DNS and validates **all** returned IPs.38    4. Rewrites the request to connect to the first valid IP while39       preserving the original `Host` header and TLS SNI hostname.4041    Redirects are re-validated on each hop because `follow_redirects`42    is set on the *client*, causing `handle_async_request` to be called43    again for each redirect target.44    """4546    def __init__(47        self,48        policy: SSRFPolicy = DEFAULT_SSRF_POLICY,49        **transport_kwargs: object,50    ) -> None:51        self._policy = policy52        self._inner = httpx.AsyncHTTPTransport(**transport_kwargs)  # type: ignore[arg-type]5354    # ------------------------------------------------------------------ #55    # Core request handler56    # ------------------------------------------------------------------ #5758    async def handle_async_request(59        self,60        request: httpx.Request,61    ) -> httpx.Response:62        hostname = request.url.host or ""63        scheme = request.url.scheme.lower()6465        # 1-3. Scheme, hostname, and pattern checks (reuse sync validator).66        validate_url_sync(str(request.url), self._policy)6768        # Allowed-hosts bypass - skip DNS/IP validation entirely.69        allowed = {h.lower() for h in _effective_allowed_hosts(self._policy)}70        if hostname.lower() in allowed:71            return await self._inner.handle_async_request(request)7273        # 4. DNS resolution74        port = request.url.port or (443 if scheme == "https" else 80)75        try:76            addrinfo = await asyncio.to_thread(77                socket.getaddrinfo,78                hostname,79                port,80                type=socket.SOCK_STREAM,81            )82        except socket.gaierror as exc:83            msg = "DNS resolution failed"84            raise SSRFBlockedError(msg) from exc8586        if not addrinfo:87            msg = "DNS resolution returned no results"88            raise SSRFBlockedError(msg)8990        # 5. Validate ALL resolved IPs - any blocked means reject.91        for _family, _type, _proto, _canonname, sockaddr in addrinfo:92            ip_str: str = sockaddr[0]  # type: ignore[assignment]93            validate_resolved_ip(ip_str, self._policy)9495        # 6. Pin to first resolved IP.96        pinned_ip = addrinfo[0][4][0]9798        # 7. Rewrite URL to use pinned IP, preserving Host header and SNI.99        pinned_url = request.url.copy_with(host=pinned_ip)100101        # Build extensions dict, adding sni_hostname for HTTPS so TLS102        # certificate validation uses the original hostname.103        extensions = dict(request.extensions)104        if scheme == "https":105            extensions["sni_hostname"] = hostname.encode("ascii")106107        pinned_request = httpx.Request(108            method=request.method,109            url=pinned_url,110            headers=request.headers,  # Host header already set to original111            content=request.content,112            extensions=extensions,113        )114115        return await self._inner.handle_async_request(pinned_request)116117    # ------------------------------------------------------------------ #118    # Lifecycle119    # ------------------------------------------------------------------ #120121    async def aclose(self) -> None:122        await self._inner.aclose()123124125# ---------------------------------------------------------------------- #126# Factory127# ---------------------------------------------------------------------- #128129130class SSRFSafeSyncTransport(httpx.BaseTransport):131    """httpx sync transport that validates DNS results against an SSRF policy.132133    Sync mirror of `SSRFSafeTransport`. See that class for full documentation.134    """135136    def __init__(137        self,138        policy: SSRFPolicy = DEFAULT_SSRF_POLICY,139        **transport_kwargs: object,140    ) -> None:141        self._policy = policy142        self._inner = httpx.HTTPTransport(**transport_kwargs)  # type: ignore[arg-type]143144    def handle_request(145        self,146        request: httpx.Request,147    ) -> httpx.Response:148        hostname = request.url.host or ""149        scheme = request.url.scheme.lower()150151        validate_url_sync(str(request.url), self._policy)152153        allowed = {h.lower() for h in _effective_allowed_hosts(self._policy)}154        if hostname.lower() in allowed:155            return self._inner.handle_request(request)156157        port = request.url.port or (443 if scheme == "https" else 80)158        try:159            addrinfo = socket.getaddrinfo(160                hostname,161                port,162                type=socket.SOCK_STREAM,163            )164        except socket.gaierror as exc:165            msg = "DNS resolution failed"166            raise SSRFBlockedError(msg) from exc167168        if not addrinfo:169            msg = "DNS resolution returned no results"170            raise SSRFBlockedError(msg)171172        for _family, _type, _proto, _canonname, sockaddr in addrinfo:173            ip_str: str = sockaddr[0]  # type: ignore[assignment]174            validate_resolved_ip(ip_str, self._policy)175176        pinned_ip = addrinfo[0][4][0]177        pinned_url = request.url.copy_with(host=pinned_ip)178179        extensions = dict(request.extensions)180        if scheme == "https":181            extensions["sni_hostname"] = hostname.encode("ascii")182183        pinned_request = httpx.Request(184            method=request.method,185            url=pinned_url,186            headers=request.headers,187            content=request.content,188            extensions=extensions,189        )190191        return self._inner.handle_request(pinned_request)192193    def close(self) -> None:194        self._inner.close()195196197# ---------------------------------------------------------------------- #198# Factories199# ---------------------------------------------------------------------- #200201202def ssrf_safe_client(203    policy: SSRFPolicy = DEFAULT_SSRF_POLICY,204    **kwargs: object,205) -> httpx.Client:206    """Create an `httpx.Client` with SSRF protection."""207    transport_kwargs: dict[str, object] = {}208    client_kwargs: dict[str, object] = {}209    for key, value in kwargs.items():210        if key in _TRANSPORT_KWARGS:211            transport_kwargs[key] = value212        else:213            client_kwargs[key] = value214215    transport = SSRFSafeSyncTransport(policy=policy, **transport_kwargs)216217    client_kwargs.setdefault("follow_redirects", True)218    client_kwargs.setdefault("max_redirects", 10)219220    return httpx.Client(221        transport=transport,222        **client_kwargs,  # type: ignore[arg-type]223    )224225226def ssrf_safe_async_client(227    policy: SSRFPolicy = DEFAULT_SSRF_POLICY,228    **kwargs: object,229) -> httpx.AsyncClient:230    """Create an `httpx.AsyncClient` with SSRF protection.231232    Drop-in replacement for `httpx.AsyncClient(...)` - callers just swap233    the constructor call.  Transport-specific kwargs (`verify`, `cert`,234    `retries`, etc.) are forwarded to the inner `AsyncHTTPTransport`;235    everything else goes to the `AsyncClient`.236    """237    transport_kwargs: dict[str, object] = {}238    client_kwargs: dict[str, object] = {}239    for key, value in kwargs.items():240        if key in _TRANSPORT_KWARGS:241            transport_kwargs[key] = value242        else:243            client_kwargs[key] = value244245    transport = SSRFSafeTransport(policy=policy, **transport_kwargs)246247    # Apply defaults only if not overridden by caller.248    client_kwargs.setdefault("follow_redirects", True)249    client_kwargs.setdefault("max_redirects", 10)250251    return httpx.AsyncClient(252        transport=transport,253        **client_kwargs,  # type: ignore[arg-type]254    )