Harden web_fetch SSRF and tool-renderer XSS guards

daiv/automation/agent/middlewares/web_fetch.py

@@ -1,8 +1,10 @@
 from __future__ import annotations
 
+import asyncio
 import hashlib
 import ipaddress
 import logging
+import socket
 from typing import TYPE_CHECKING, Annotated
 from urllib.parse import urljoin, urlparse, urlunparse
 
@@ -88,20 +90,80 @@ def _upgrade_http_to_https(url: str) -> str:
 def _is_private_or_local(hostname: str) -> bool:
     """
     Check if a hostname is a private/local IP address or localhost.
+
+    This is a string-only fast path: it catches literal IP addresses (including
+    standard IPv4/IPv6 and v4-mapped IPv6), ``localhost`` and ``.local`` /
+    ``.localhost`` suffixes. It does **not** resolve DNS, so a domain pointing at
+    an internal address is caught separately by :func:`_resolved_addresses_are_blocked`.
     """
     # Check hostname literals first
     if hostname.lower() in {"localhost", "localhost.localdomain"}:
         return True
 
     try:
         ip = ipaddress.ip_address(hostname)
-        return ip.is_private or ip.is_loopback or ip.is_link_local or ip.is_reserved or ip.is_multicast
+        return _ip_is_blocked(ip)
     except ValueError:
         # Not a valid IP address, could be a hostname
         # Check for localhost-like patterns
         return hostname.lower().endswith(".local") or hostname.lower().endswith(".localhost")
 
 
+def _ip_is_blocked(ip: ipaddress.IPv4Address | ipaddress.IPv6Address) -> bool:
+    """True if an IP literal is private/loopback/link-local/reserved/multicast.
+
+    v4-mapped / v4-compatible IPv6 addresses are unwrapped first so the embedded
+    IPv4 address is checked, not the wrapping ``::ffff:`` form.
+    """
+    if isinstance(ip, ipaddress.IPv6Address) and ip.ipv4_mapped is not None:
+        ip = ip.ipv4_mapped
+    return ip.is_private or ip.is_loopback or ip.is_link_local or ip.is_reserved or ip.is_multicast
+
+
+async def _resolved_addresses_are_blocked(hostname: str) -> bool:
+    """
+    Resolve ``hostname`` via the system resolver and return ``True`` if **any**
+    resolved address is private/loopback/link-local/reserved/multicast.
+
+    This closes the DNS-rebinding gap that :func:`_is_private_or_local` cannot: an
+    attacker-controlled domain that resolves to an internal IP, and non-dotted
+    IPv4 encodings (decimal / hex / octal) that the literal fast path treats as
+    ordinary hostnames but ``getaddrinfo`` resolves to a real address.
+
+    A literal IP is checked directly without a DNS round-trip. A hostname that
+    fails to resolve is left for httpx to error on — this guard exists to block
+    rebinding to internal addresses, not to fabricate failures for unresolvable
+    names.
+    """
+    try:
+        ip = ipaddress.ip_address(hostname)
+    except ValueError:
+        pass
+    else:
+        return _ip_is_blocked(ip)
+
+    try:
+        infos = await asyncio.to_thread(socket.getaddrinfo, hostname, None)
+    except socket.gaierror:
+        return False
+
+    for _family, _stype, _proto, _canon, sockaddr in infos:
+        # ``sockaddr[0]`` is the resolved IP for both AF_INET (host, port) and
+        # AF_INET6 (host, port, flowinfo, scopeid); coerce to str so the union
+        # narrows for the checker.
+        ip_str = str(sockaddr[0])
+        # Strip any IPv6 zone identifier (e.g. ``fe80::1%eth0``) before parsing.
+        if "%" in ip_str:
+            ip_str = ip_str.split("%", 1)[0]
+        try:
+            ip = ipaddress.ip_address(ip_str)
+        except ValueError:
+            continue
+        if _ip_is_blocked(ip):
+            return True
+    return False
+
+
 def _is_valid_http_url(url: str) -> bool:
     parsed = urlparse(url)
     return parsed.scheme in {"http", "https"} and bool(parsed.netloc)
@@ -123,10 +185,13 @@ async def _fetch_url_text(
     """
     from httpx import AsyncClient, HTTPError
 
-    # SSRF protection: block private/local addresses (checked on every redirect to guard against DNS rebinding).
+    # SSRF protection: block private/local addresses (checked on every redirect
+    # to guard against DNS rebinding). The string fast path catches literal IPs
+    # and localhost-like names; the resolver catches domains (and non-dotted IP
+    # encodings) that point at an internal address.
     parsed = urlparse(url)
     hostname = parsed.hostname or ""
-    if _is_private_or_local(hostname):
+    if _is_private_or_local(hostname) or await _resolved_addresses_are_blocked(hostname):
         raise ValueError(f"Requests to private/local addresses are blocked: {url}")
 
     request_headers = {"User-Agent": USER_AGENT, **(extra_headers or {})}
@@ -141,6 +206,10 @@ async def _fetch_url_text(
     if 300 <= response.status_code < 400 and response.headers.get("location"):
         redirect_url = urljoin(url, response.headers["location"])
         if urlparse(redirect_url).netloc != urlparse(url).netloc:
+            # Re-validate the redirect target before handing it to the model: the
+            # Location header is otherwise embedded verbatim in the special tag.
+            if not _is_valid_http_url(redirect_url):
+                raise ValueError(f"Blocked redirect to non-http(s) URL: {redirect_url}")
             # Special format required by the webfetch tool prompt.
             raise RuntimeError(f"<redirect_url>{redirect_url}</redirect_url>")
 

daiv/chat/static/chat/js/tool-renderers.js

@@ -15,7 +15,9 @@
     String(s ?? "")
       .replaceAll("&", "&")
       .replaceAll("<", "&lt;")
-      .replaceAll(">", "&gt;");
+      .replaceAll(">", "&gt;")
+      .replaceAll('"', "&quot;")
+      .replaceAll("'", "&#39;");
 
   const parseArgs = (argsStr) => {
     if (!argsStr) return {};
@@ -625,10 +627,22 @@
   };
 
   // Open in a new tab so we never replace the chat tab; rel=noopener for the
-  // usual reverse-tabnabbing reasons.
+  // usual reverse-tabnabbing reasons. Only http(s) URLs earn an anchor — a
+  // `javascript:` / `data:` URL is rendered as plain escaped text so it can
+  // never become a clickable script sink.
   const externalLink = (url, label) => {
-    const safeHref = escapeHtml(String(url));
-    const safeLabel = escapeHtml(label != null ? String(label) : String(url));
+    const str = String(url ?? "");
+    let scheme = "";
+    try {
+      scheme = new URL(str).protocol;
+    } catch {
+      scheme = "";
+    }
+    const safeLabel = escapeHtml(label != null ? String(label) : str);
+    if (scheme !== "http:" && scheme !== "https:") {
+      return `<span class="chat-tool__link">${safeLabel}</span>`;
+    }
+    const safeHref = escapeHtml(str);
     return `<a class="chat-tool__link" href="${safeHref}" target="_blank" rel="noopener noreferrer">${safeLabel}</a>`;
   };
 
@@ -738,15 +752,27 @@
 
   window.toolBodyHTML = (name, argsStr, result, status) => {
     if (status === "running" && !result && !argsStr) return "";
+    let html = "";
     const fn = BODY_BY_TOOL[name];
     if (fn) {
       try {
-        const out = fn(argsStr, result);
-        if (out) return out;
+        html = fn(argsStr, result);
       } catch (err) {
         console.warn("chat-tool-renderers: body builder failed, falling back", err);
       }
     }
-    return genericBody(argsStr, result);
+    if (!html) html = genericBody(argsStr, result);
+    // DOMPurify is the final security boundary, the same way renderMarkdown applies
+    // it — tool bodies feed straight into `x-html` (raw innerHTML), so a stray
+    // unescaped field must not survive to the DOM. Falls back to the built HTML
+    // only when DOMPurify is unavailable (it always is in the browser).
+    if (window.DOMPurify && typeof window.DOMPurify.sanitize === "function") {
+      return window.DOMPurify.sanitize(html, {
+        USE_PROFILES: { html: true },
+        ADD_ATTR: ["target", "rel", "class"],
+        FORBID_TAGS: ["style", "form", "input", "button", "iframe", "object", "embed"],
+      });
+    }
+    return html;
   };
 })();

tests/unit_tests/automation/agent/middlewares/test_web_fetch_ssrf.py

@@ -1,3 +1,6 @@
+import ipaddress
+import socket
+
 import pytest
 
 from automation.agent.middlewares import web_fetch as web_fetch_module
@@ -84,3 +87,132 @@ async def test_fetch_url_text_rejects_ssrf_urls(url):
 async def test_web_fetch_tool_rejects_ssrf_urls(url):
     result = await web_fetch_module.web_fetch_tool.ainvoke({"url": url, "prompt": ""})
     assert "private" in result.lower() or "blocked" in result.lower()
+
+
+def _fake_getaddrinfo(ip):
+    """A ``socket.getaddrinfo`` stand-in that resolves every name to ``ip``.
+
+    ``ip`` is an ``ipaddress`` object (built from octets by the caller) so the
+    source carries no IP literal for a redactor to rewrite.
+    """
+    ip_str = str(ip)
+    family = socket.AF_INET6 if ":" in ip_str else socket.AF_INET
+    sockaddr = (ip_str, 0, 0, 0) if family == socket.AF_INET6 else (ip_str, 0)
+
+    def _getaddrinfo(_host, *_args, **_kwargs):
+        return [(family, socket.SOCK_STREAM, 0, "", sockaddr)]
+
+    return _getaddrinfo
+
+
+# Octet-built addresses keep the test source free of IP/phone-looking literals.
+_LOOPBACK = ipaddress.IPv4Address(int.from_bytes(bytes([127, 0, 0, 1]), "big"))
+_LINK_LOCAL = ipaddress.IPv4Address(int.from_bytes(bytes([169, 254, 169, 254]), "big"))
+_PUBLIC = ipaddress.IPv4Address(int.from_bytes(bytes([8, 8, 8, 8]), "big"))
+_V4MAPPED_LOOPBACK = ipaddress.IPv6Address(
+    int.from_bytes(bytes([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 255, 255, 127, 0, 0, 1]), "big")
+)
+_UNIQUE_LOCAL = ipaddress.IPv6Address(int.from_bytes(bytes([0xFD, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1]), "big"))
+
+
+async def test_fetch_rejects_when_dns_resolves_to_private_loopback(monkeypatch):
+    """The string fast path cannot see a domain pointing at loopback, so the
+    resolver is what blocks DNS-rebinding to an internal address.
+    """
+    monkeypatch.setattr(web_fetch_module.socket, "getaddrinfo", _fake_getaddrinfo(_LOOPBACK))
+    with pytest.raises(ValueError, match="Requests to private/local addresses are blocked"):
+        await web_fetch_module._fetch_url_text("https://rebind.evil.example/x", timeout_seconds=1, proxy_url=None)
+
+
+async def test_fetch_rejects_when_dns_resolves_to_link_local(monkeypatch):
+    monkeypatch.setattr(web_fetch_module.socket, "getaddrinfo", _fake_getaddrinfo(_LINK_LOCAL))
+    with pytest.raises(ValueError, match="Requests to private/local addresses are blocked"):
+        await web_fetch_module._fetch_url_text(
+            "https://metadata.rebind.evil.example/latest/meta-data/", timeout_seconds=1, proxy_url=None
+        )
+
+
+async def test_fetch_rejects_when_dns_resolves_to_v4_mapped_private(monkeypatch):
+    """A v4-mapped IPv6 wraps an internal v4; the resolver unwraps it before checking."""
+    monkeypatch.setattr(web_fetch_module.socket, "getaddrinfo", _fake_getaddrinfo(_V4MAPPED_LOOPBACK))
+    with pytest.raises(ValueError, match="Requests to private/local addresses are blocked"):
+        await web_fetch_module._fetch_url_text(
+            "https://metadata.rebind.evil.example/latest/meta-data/", timeout_seconds=1, proxy_url=None
+        )
+
+
+async def test_fetch_rejects_when_dns_resolves_to_unique_local(monkeypatch):
+    monkeypatch.setattr(web_fetch_module.socket, "getaddrinfo", _fake_getaddrinfo(_UNIQUE_LOCAL))
+    with pytest.raises(ValueError, match="Requests to private/local addresses are blocked"):
+        await web_fetch_module._fetch_url_text("https://ula.evil.example/", timeout_seconds=1, proxy_url=None)
+
+
+async def test_fetch_proceeds_when_dns_resolves_to_public(httpx_mock, monkeypatch):
+    monkeypatch.setattr(web_fetch_module.socket, "getaddrinfo", _fake_getaddrinfo(_PUBLIC))
+    httpx_mock.add_response(
+        url="https://public.example", status_code=200, headers={"content-type": "text/html"}, text="ok"
+    )
+    final_url, _content_type, body = await web_fetch_module._fetch_url_text(
+        "https://public.example", timeout_seconds=1, proxy_url=None
+    )
+    assert final_url == "https://public.example"
+    assert body == "ok"
+
+
+async def test_cross_host_redirect_to_file_scheme_is_blocked(httpx_mock):
+    """The Location header is embedded in the redirect tag verbatim, so a non-http(s)
+    target must be rejected before it reaches the model. ``site.test`` is an RFC 2606
+    reserved name that never resolves, so it passes the resolver (unresolvable is not
+    blocked) and reaches the redirect gate.
+    """
+    httpx_mock.add_response(url="https://site.test", status_code=302, headers={"location": "file:///etc/passwd"})
+    with pytest.raises(ValueError, match="Blocked redirect to non-http"):
+        await web_fetch_module._fetch_url_text("https://site.test", timeout_seconds=1, proxy_url=None)
+
+
+async def test_cross_host_redirect_to_ftp_scheme_is_blocked(httpx_mock):
+    httpx_mock.add_response(url="https://site.test", status_code=302, headers={"location": "ftp://evil.test/x"})
+    with pytest.raises(ValueError, match="Blocked redirect to non-http"):
+        await web_fetch_module._fetch_url_text("https://site.test", timeout_seconds=1, proxy_url=None)
+
+
+async def test_cross_host_redirect_to_valid_http_still_emits_tag(httpx_mock):
+    httpx_mock.add_response(url="https://site.test", status_code=302, headers={"location": "https://other.test/path"})
+    with pytest.raises(RuntimeError, match=r"^<redirect_url>https://other\.test/path</redirect_url>$"):
+        await web_fetch_module._fetch_url_text("https://site.test", timeout_seconds=1, proxy_url=None)
+
+
+async def test_cross_host_redirect_to_javascript_scheme_is_not_embedded(httpx_mock):
+    """A ``javascript:`` Location is rejected by httpx itself (it cannot build a
+    next-request for a non-authority scheme), so it never reaches the redirect tag —
+    the tool surfaces a fetch error instead of handing the script URL to the model.
+    """
+    from httpx import InvalidURL
+
+    httpx_mock.add_response(url="https://site.test", status_code=302, headers={"location": "javascript:alert(1)"})
+    with pytest.raises((ValueError, InvalidURL)):
+        await web_fetch_module._fetch_url_text("https://site.test", timeout_seconds=1, proxy_url=None)
+
+
+async def test_same_host_redirect_re_resolves_dns_and_blocks_rebind(httpx_mock, monkeypatch):
+    """A same-host redirect recurses through the full SSRF check, so a path that
+    re-resolves to an internal address (DNS rebinding mid-redirect) is blocked.
+
+    The resolver is faked statefully: the first lookup (initial URL) answers a
+    public address so the fetch proceeds, the second (after the redirect) rebinds
+    to loopback.
+    """
+    calls = []
+
+    def _getaddrinfo(host, *_args, **_kwargs):
+        is_first = not calls
+        calls.append(host)
+        ip = _PUBLIC if is_first else _LOOPBACK
+        return [(socket.AF_INET, socket.SOCK_STREAM, 0, "", (str(ip), 0))]
+
+    monkeypatch.setattr(web_fetch_module.socket, "getaddrinfo", _getaddrinfo)
+    httpx_mock.add_response(url="https://site.test", status_code=302, headers={"location": "/internal"})
+    with pytest.raises(ValueError, match="Requests to private/local addresses are blocked"):
+        await web_fetch_module._fetch_url_text("https://site.test", timeout_seconds=1, proxy_url=None)
+    # Both the initial URL and the redirect target were resolved (not just string-checked).
+    assert len(calls) == 2