diff --git a/nanobot/agent/tools/web.py b/nanobot/agent/tools/web.py index 74f17918..b8b7f971 100644 --- a/nanobot/agent/tools/web.py +++ b/nanobot/agent/tools/web.py @@ -118,6 +118,12 @@ def _resolve_url_safe(url: str) -> tuple[bool, str, tuple[str, ...]]: return resolve_url_target(url) +def _pinned_dns_transport(proxy: str | None = None) -> httpx.AsyncBaseTransport: + from nanobot.security.network import PinnedDNSAsyncTransport + + return PinnedDNSAsyncTransport(proxy=proxy) + + async def _get_with_safe_redirects( client: httpx.AsyncClient, url: str, @@ -126,14 +132,11 @@ async def _get_with_safe_redirects( """GET a URL while validating every redirect target before requesting it.""" current_url = url for _ in range(MAX_REDIRECTS + 1): - is_valid, error_msg, resolved_ips = _resolve_url_safe(current_url) + is_valid, error_msg, _ = _resolve_url_safe(current_url) if not is_valid: return None, f"Redirect blocked: {error_msg}" - from nanobot.security.network import pin_resolved_url_dns - - with pin_resolved_url_dns(current_url, resolved_ips): - response = await client.get(current_url, headers=headers, follow_redirects=False) + response = await client.get(current_url, headers=headers, follow_redirects=False) is_redirect = 300 <= response.status_code < 400 if not is_redirect: return response, None @@ -162,20 +165,17 @@ async def _stream_with_safe_redirects( """Open a streamed response while validating every redirect target first.""" current_url = url for _ in range(MAX_REDIRECTS + 1): - is_valid, error_msg, resolved_ips = _resolve_url_safe(current_url) + is_valid, error_msg, _ = _resolve_url_safe(current_url) if not is_valid: return None, None, f"Redirect blocked: {error_msg}" - from nanobot.security.network import pin_resolved_url_dns - stream = client.stream( "GET", current_url, headers=headers, follow_redirects=False, ) - with pin_resolved_url_dns(current_url, resolved_ips): - response = await stream.__aenter__() + response = await stream.__aenter__() is_redirect = 300 <= response.status_code < 400 if not is_redirect: return response, stream, None @@ -966,7 +966,10 @@ class WebFetchTool(Tool): # Detect and fetch images directly to avoid Jina's textual image captioning try: - async with httpx.AsyncClient(proxy=self.proxy, timeout=15.0) as client: + async with httpx.AsyncClient( + transport=_pinned_dns_transport(self.proxy), + timeout=15.0, + ) as client: r, stream, redirect_error = await _stream_with_safe_redirects( client, url, @@ -1037,7 +1040,7 @@ class WebFetchTool(Tool): try: async with httpx.AsyncClient( timeout=30.0, - proxy=self.proxy, + transport=_pinned_dns_transport(self.proxy), ) as client: r, redirect_error = await _get_with_safe_redirects( client, diff --git a/nanobot/security/network.py b/nanobot/security/network.py index a17919a4..3cb7df3c 100644 --- a/nanobot/security/network.py +++ b/nanobot/security/network.py @@ -2,6 +2,7 @@ from __future__ import annotations +import asyncio import ipaddress import re import socket @@ -114,7 +115,12 @@ def validate_url_target(url: str, *, allow_loopback: bool = False) -> tuple[bool @contextmanager def pin_resolved_url_dns(url: str, resolved_ips: tuple[str, ...]): - """Pin DNS lookups for the URL hostname to previously validated IPs.""" + """Pin DNS lookups for the URL hostname to previously validated IPs. + + This temporarily overrides process-global resolver state. Do not use it + directly across awaits unless the caller serializes access; prefer + PinnedDNSAsyncTransport for HTTP requests. + """ try: hostname = urlparse(url).hostname except Exception: @@ -149,17 +155,26 @@ def pin_resolved_url_dns(url: str, resolved_ips: tuple[str, ...]): class PinnedDNSAsyncTransport(httpx.AsyncBaseTransport): """HTTPX transport that pins each request to the IPs validated for its URL.""" - def __init__(self, *, allow_loopback: bool = False) -> None: + _resolver_lock = asyncio.Lock() + + def __init__( + self, + *, + allow_loopback: bool = False, + proxy: httpx.ProxyTypes | None = None, + inner: httpx.AsyncBaseTransport | None = None, + ) -> None: self._allow_loopback = allow_loopback - self._inner = httpx.AsyncHTTPTransport() + self._inner = inner or httpx.AsyncHTTPTransport(proxy=proxy) async def handle_async_request(self, request: httpx.Request) -> httpx.Response: url = str(request.url) ok, error, resolved_ips = resolve_url_target(url, allow_loopback=self._allow_loopback) if not ok: raise httpx.RequestError(error, request=request) - with pin_resolved_url_dns(url, resolved_ips): - return await self._inner.handle_async_request(request) + async with self._resolver_lock: + with pin_resolved_url_dns(url, resolved_ips): + return await self._inner.handle_async_request(request) async def aclose(self) -> None: await self._inner.aclose() diff --git a/tests/tools/test_web_fetch_security.py b/tests/tools/test_web_fetch_security.py index dfda4b92..7523f81b 100644 --- a/tests/tools/test_web_fetch_security.py +++ b/tests/tools/test_web_fetch_security.py @@ -2,6 +2,7 @@ from __future__ import annotations +import asyncio import json import socket from unittest.mock import patch @@ -12,6 +13,7 @@ import pytest from nanobot.agent.tools import web as web_module from nanobot.agent.tools.web import WebFetchTool, _get_with_safe_redirects from nanobot.config.schema import WebFetchConfig +from nanobot.security.network import PinnedDNSAsyncTransport from nanobot.security.workspace_access import ( bind_workspace_scope, build_workspace_scope, @@ -98,27 +100,51 @@ async def test_web_fetch_result_contains_untrusted_flag(): @pytest.mark.asyncio -async def test_safe_redirect_request_pins_validated_dns(monkeypatch): - calls: list[str] = [] +async def test_safe_redirect_requests_use_independent_pinned_dns_concurrently(monkeypatch): + public_ips = { + "a.example": "93.184.216.34", + "b.example": "93.184.216.35", + } + calls: dict[str, int] = {host: 0 for host in public_ips} + seen: dict[str, str] = {} - def _rebinding_resolver(hostname, port, family=0, type_=0): - calls.append(hostname) - ip = "93.184.216.34" if len(calls) == 1 else "169.254.169.254" + def _rebinding_resolver(hostname, port, family=0, type_=0, proto=0, flags=0): + host = str(hostname).rstrip(".").lower() + calls[host] += 1 + ip = public_ips[host] if calls[host] <= 2 else "169.254.169.254" return [(socket.AF_INET, socket.SOCK_STREAM, 0, "", (ip, 0))] - class FakeClient: - async def get(self, url, headers=None, follow_redirects=False): - infos = socket.getaddrinfo("attacker.example", 443, socket.AF_UNSPEC, socket.SOCK_STREAM) - assert infos[0][4][0] == "93.184.216.34" - return httpx.Response(200, request=httpx.Request("GET", url)) + class ResolvingTransport(httpx.AsyncBaseTransport): + async def handle_async_request(self, request: httpx.Request) -> httpx.Response: + await asyncio.sleep(0) + infos = socket.getaddrinfo( + request.url.host, + request.url.port or 443, + socket.AF_UNSPEC, + socket.SOCK_STREAM, + ) + seen[str(request.url)] = infos[0][4][0] + return httpx.Response(200, request=request) + + async def _fetch(url: str) -> tuple[httpx.Response | None, str | None]: + async with httpx.AsyncClient( + transport=PinnedDNSAsyncTransport(inner=ResolvingTransport()) + ) as client: + return await _get_with_safe_redirects(client, url) monkeypatch.setattr("nanobot.security.network.socket.getaddrinfo", _rebinding_resolver) - response, error = await _get_with_safe_redirects(FakeClient(), "https://attacker.example/") + results = await asyncio.gather( + _fetch("https://a.example/"), + _fetch("https://b.example/"), + ) - assert error is None - assert response is not None - assert calls == ["attacker.example"] + assert all(error is None and response is not None for response, error in results) + assert seen == { + "https://a.example/": "93.184.216.34", + "https://b.example/": "93.184.216.35", + } + assert calls == {"a.example": 2, "b.example": 2} @pytest.mark.asyncio @@ -318,6 +344,7 @@ async def test_web_fetch_blocks_private_redirect_before_returning_image(monkeypa class TransportAsyncClient(real_async_client): def __init__(self, *args, **kwargs): kwargs.pop("proxy", None) + kwargs.pop("transport", None) super().__init__(*args, transport=transport, **kwargs) monkeypatch.setattr("nanobot.agent.tools.web.httpx.AsyncClient", TransportAsyncClient)