Fix validated Codex Security findings

assert_ai/core/security.py

@@ -132,6 +132,13 @@ def validate_sys_path_addition(path: Path, *, config_path: Path | None = None) -
     ipaddress.ip_network("fe80::/10"),         # IPv6 link-local
 ]
 
+_NAT64_WELL_KNOWN_PREFIX = ipaddress.ip_network("64:ff9b::/96")
+_IPV4_COMPATIBLE_PREFIX = ipaddress.ip_network("::/96")
+_IPV4_COMPATIBLE_RESERVED = {
+    ipaddress.ip_address("::"),
+    ipaddress.ip_address("::1"),
+}
+
 _BLOCKED_HOSTNAMES = {
     "metadata.google.internal",
     "metadata.google.com",
@@ -189,20 +196,15 @@ def validate_endpoint_url(url: str, *, allow_private: bool = False) -> None:
             f"URL hostname '{hostname}' is blocked (potential metadata endpoint)"
         )
 
-    # Try to parse as IP address
+    # Validate IP literals directly; hostnames are checked after DNS resolution.
     try:
-        ip = ipaddress.ip_address(hostname)
-        for network in _BLOCKED_IP_RANGES:
-            if ip in network:
-                raise ValueError(
-                    f"URL resolves to blocked IP range ({network}): {hostname}"
-                )
-    except ValueError as e:
-        if "blocked" in str(e).lower():
-            raise
-        # Not an IP literal — resolve hostname and check resulting IPs
+        ipaddress.ip_address(hostname)
+    except ValueError:
+        # Not an IP literal — resolve hostname and check resulting IPs.
         if hostname.lower() not in _LOCAL_DEV_HOSTNAMES:
             _validate_resolved_ips(hostname)
+    else:
+        validate_resolved_endpoint_ip(hostname, hostname)
 
 
 def _validate_resolved_ips(hostname: str) -> None:
@@ -219,23 +221,81 @@ def _validate_resolved_ips(hostname: str) -> None:
         log.debug("DNS resolution failed for '%s'; skipping IP validation", hostname)
         return
 
-    for family, _type, _proto, _canonname, sockaddr in addrinfo:
-        ip_str = sockaddr[0]
-        try:
-            ip = ipaddress.ip_address(ip_str)
-        except ValueError:
-            continue
-        for network in _BLOCKED_IP_RANGES:
-            if ip in network:
-                log.warning(
-                    "SSRF protection: hostname '%s' resolves to blocked IP %s (range %s)",
-                    hostname,
-                    ip_str,
-                    network,
-                )
-                raise ValueError(
-                    f"URL hostname '{hostname}' resolves to blocked IP range ({network}): {ip_str}"
-                )
+    for _family, _type, _proto, _canonname, sockaddr in addrinfo:
+        validate_resolved_endpoint_ip(hostname, sockaddr[0])
+
+
+def _canonicalize_endpoint_ip(
+    ip: ipaddress.IPv4Address | ipaddress.IPv6Address,
+) -> ipaddress.IPv4Address | ipaddress.IPv6Address:
+    """Return the address an IPv4-in-IPv6 endpoint ultimately reaches."""
+    if not isinstance(ip, ipaddress.IPv6Address):
+        return ip
+
+    if ip.ipv4_mapped is not None:
+        return ip.ipv4_mapped
+
+    if ip in _NAT64_WELL_KNOWN_PREFIX or (
+        ip in _IPV4_COMPATIBLE_PREFIX and ip not in _IPV4_COMPATIBLE_RESERVED
+    ):
+        return ipaddress.IPv4Address(int(ip) & 0xFFFFFFFF)
+
+    return ip
+
+
+def validate_resolved_endpoint_ip(hostname: str, ip_str: str) -> None:
+    """Validate one DNS answer immediately before an endpoint connection.
+
+    ``validate_endpoint_url`` performs an eager DNS check for fast feedback, but
+    HTTP clients resolve again when they connect.  Call this helper on the
+    resolver results that will actually be used for the socket so a hostname
+    cannot pass validation with a public address and later rebind to a private
+    one.
+    """
+    if os.environ.get("ASSERT_ALLOW_PRIVATE_ENDPOINTS", "").lower() in (
+        "1",
+        "true",
+        "yes",
+    ):
+        return
+    if hostname.lower() in _LOCAL_DEV_HOSTNAMES:
+        return
+
+    try:
+        ip = ipaddress.ip_address(ip_str)
+    except ValueError as exc:
+        raise ValueError(
+            f"Resolver returned a non-IP address for endpoint hostname '{hostname}': {ip_str}"
+        ) from exc
+
+    # Mapped, NAT64, and deprecated IPv4-compatible IPv6 addresses inherit the
+    # security properties of the IPv4 endpoint they ultimately reach.
+    checked_ip = _canonicalize_endpoint_ip(ip)
+
+    for network in _BLOCKED_IP_RANGES:
+        if checked_ip.version == network.version and checked_ip in network:
+            log.warning(
+                "SSRF protection: hostname '%s' resolves to blocked IP %s (range %s)",
+                hostname,
+                ip_str,
+                network,
+            )
+            raise ValueError(
+                f"URL hostname '{hostname}' resolves to blocked IP range ({network}): {ip_str}"
+            )
+
+    # Block unspecified, reserved, documentation, shared, and other special-use
+    # addresses as well as multicast.  SSRF targets must resolve to a publicly
+    # routable unicast address unless the explicit development override is set.
+    if not checked_ip.is_global or checked_ip.is_multicast:
+        log.warning(
+            "SSRF protection: hostname '%s' resolves to non-public IP %s",
+            hostname,
+            ip_str,
+        )
+        raise ValueError(
+            f"URL hostname '{hostname}' resolves to a non-public IP address: {ip_str}"
+        )
 
 
 # ── Credential sanitization ────────────────────────────────────
@@ -278,4 +338,3 @@ def sanitize_payload(payload: Any, *, depth: int = 0, max_depth: int = 10) -> An
             return _REDACTED
         return payload
     return payload
-

assert_ai/core/session.py

@@ -6,6 +6,7 @@
 from __future__ import annotations
 
 import asyncio
+import contextlib
 import inspect
 import json
 import logging
@@ -660,6 +661,7 @@ def __init__(
         self._system_prompt = system_prompt
         self._timeout_s = message_timeout_s
         self._session = None  # aiohttp.ClientSession
+        self._resolver = None
 
     @property
     def runtime_mode(self) -> str:
@@ -675,12 +677,29 @@ async def open(self) -> None:
             )
         self._aiohttp = aiohttp
         timeout = aiohttp.ClientTimeout(total=self._timeout_s or 60)
-        self._session = aiohttp.ClientSession(timeout=timeout)
+        resolver = _ValidatingResolver(aiohttp.DefaultResolver())
+        connector = None
+        try:
+            connector = aiohttp.TCPConnector(resolver=resolver)
+            self._session = aiohttp.ClientSession(timeout=timeout, connector=connector)
+        except Exception:
+            if connector is not None:
+                with contextlib.suppress(Exception):
+                    await connector.close()
+            with contextlib.suppress(Exception):
+                await resolver.close()
+            raise
+        self._resolver = resolver
 
     async def close(self) -> None:
-        if self._session:
-            await self._session.close()
-            self._session = None
+        try:
+            if self._session:
+                await self._session.close()
+                self._session = None
+        finally:
+            if self._resolver:
+                await self._resolver.close()
+                self._resolver = None
 
     async def run_turn(self, messages: list[Message]) -> TurnResult:
         aiohttp = self._aiohttp
@@ -704,7 +723,12 @@ async def run_turn(self, messages: list[Message]) -> TurnResult:
                 self._endpoint,
                 json=payload,
                 headers=self._headers,
+                allow_redirects=False,
             ) as resp:
+                if 300 <= resp.status < 400:
+                    raise RuntimeError(
+                        f"HTTP endpoint {self._endpoint} returned a redirect, which is not allowed"
+                    )
                 resp.raise_for_status()
                 data = await resp.json()
                 response_text = data.get("response", "")
@@ -733,6 +757,29 @@ async def run_turn(self, messages: list[Message]) -> TurnResult:
         )
 
 
+class _ValidatingResolver:
+    """Validate the exact DNS answers aiohttp will use for a connection."""
+
+    def __init__(self, resolver: Any) -> None:
+        self._resolver = resolver
+
+    async def resolve(self, host: str, port: int = 0, family: int = 0) -> list[Any]:
+        from assert_ai.core.security import validate_resolved_endpoint_ip
+
+        results = await self._resolver.resolve(host, port, family)
+        try:
+            for result in results:
+                validate_resolved_endpoint_ip(host, str(result["host"]))
+        except ValueError as exc:
+            # aiohttp converts resolver OSErrors into ClientConnectorError,
+            # which the session already exposes as a scoped RuntimeError.
+            raise OSError(str(exc)) from exc
+        return results
+
+    async def close(self) -> None:
+        await self._resolver.close()
+
+
 class ExternalSession:
     def __init__(
         self,

assert_ai/stages/inference.py

@@ -809,17 +809,16 @@ async def _run_tester_target_loop(
         assert action_message is not None
 
         target_messages.append(Message(role="user", content=action_message))
-        tester_call_id = transcript.append_llm_call(
-            **build_llm_call_trace(tester_response, source="tester")
-        )
+        tester_call_trace = build_llm_call_trace(tester_response, source="tester")
+        tester_call_id = transcript.append_llm_call(**tester_call_trace)
         message_id = f"event:{len(transcript.events)}"
         transcript.add_event(TranscriptEvent(
             view=["target", "combined"],
             actor="tester",
             edit=AddMessageEdit(message=TranscriptMessage(role="user", content=action_message)),
             raw={
                 "call": "tester",
-                "request": to_jsonable(tester_response.request_payload or {}),
+                "request": tester_call_trace["request"],
                 "response": serialize_response(tester_response),
             },
         ))
@@ -1140,7 +1139,7 @@ async def _worker(test_case: tuple[int, dict[str, Any]]) -> dict[str, Any]:
             # LLMInputError) so we emit a quieter debug log instead of
             # the warning the generic input-error handler below would
             # emit for every adversarial case.
-            test_case_id = seed_row.get("seed_id", "?")
+            test_case_id = test_case_row.get("test_case_id", "?")
             log.debug(
                 "Inference worker hit provider content filter for test case %s: %s",
                 test_case_id, exc,