fix: serialize pinned dns web fetches

This commit is contained in:
hamb1y
2026-07-07 15:40:53 +08:00
committed by Xubin Ren
parent 4353f4680b
commit 97e3b360c2
3 changed files with 76 additions and 31 deletions
+13 -10
View File
@@ -118,6 +118,12 @@ def _resolve_url_safe(url: str) -> tuple[bool, str, tuple[str, ...]]:
return resolve_url_target(url) 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( async def _get_with_safe_redirects(
client: httpx.AsyncClient, client: httpx.AsyncClient,
url: str, url: str,
@@ -126,13 +132,10 @@ async def _get_with_safe_redirects(
"""GET a URL while validating every redirect target before requesting it.""" """GET a URL while validating every redirect target before requesting it."""
current_url = url current_url = url
for _ in range(MAX_REDIRECTS + 1): 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: if not is_valid:
return None, f"Redirect blocked: {error_msg}" 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 is_redirect = 300 <= response.status_code < 400
if not is_redirect: if not is_redirect:
@@ -162,19 +165,16 @@ async def _stream_with_safe_redirects(
"""Open a streamed response while validating every redirect target first.""" """Open a streamed response while validating every redirect target first."""
current_url = url current_url = url
for _ in range(MAX_REDIRECTS + 1): 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: if not is_valid:
return None, None, f"Redirect blocked: {error_msg}" return None, None, f"Redirect blocked: {error_msg}"
from nanobot.security.network import pin_resolved_url_dns
stream = client.stream( stream = client.stream(
"GET", "GET",
current_url, current_url,
headers=headers, headers=headers,
follow_redirects=False, 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 is_redirect = 300 <= response.status_code < 400
if not is_redirect: if not is_redirect:
@@ -966,7 +966,10 @@ class WebFetchTool(Tool):
# Detect and fetch images directly to avoid Jina's textual image captioning # Detect and fetch images directly to avoid Jina's textual image captioning
try: 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( r, stream, redirect_error = await _stream_with_safe_redirects(
client, client,
url, url,
@@ -1037,7 +1040,7 @@ class WebFetchTool(Tool):
try: try:
async with httpx.AsyncClient( async with httpx.AsyncClient(
timeout=30.0, timeout=30.0,
proxy=self.proxy, transport=_pinned_dns_transport(self.proxy),
) as client: ) as client:
r, redirect_error = await _get_with_safe_redirects( r, redirect_error = await _get_with_safe_redirects(
client, client,
+18 -3
View File
@@ -2,6 +2,7 @@
from __future__ import annotations from __future__ import annotations
import asyncio
import ipaddress import ipaddress
import re import re
import socket import socket
@@ -114,7 +115,12 @@ def validate_url_target(url: str, *, allow_loopback: bool = False) -> tuple[bool
@contextmanager @contextmanager
def pin_resolved_url_dns(url: str, resolved_ips: tuple[str, ...]): 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: try:
hostname = urlparse(url).hostname hostname = urlparse(url).hostname
except Exception: except Exception:
@@ -149,15 +155,24 @@ def pin_resolved_url_dns(url: str, resolved_ips: tuple[str, ...]):
class PinnedDNSAsyncTransport(httpx.AsyncBaseTransport): class PinnedDNSAsyncTransport(httpx.AsyncBaseTransport):
"""HTTPX transport that pins each request to the IPs validated for its URL.""" """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._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: async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
url = str(request.url) url = str(request.url)
ok, error, resolved_ips = resolve_url_target(url, allow_loopback=self._allow_loopback) ok, error, resolved_ips = resolve_url_target(url, allow_loopback=self._allow_loopback)
if not ok: if not ok:
raise httpx.RequestError(error, request=request) raise httpx.RequestError(error, request=request)
async with self._resolver_lock:
with pin_resolved_url_dns(url, resolved_ips): with pin_resolved_url_dns(url, resolved_ips):
return await self._inner.handle_async_request(request) return await self._inner.handle_async_request(request)
+41 -14
View File
@@ -2,6 +2,7 @@
from __future__ import annotations from __future__ import annotations
import asyncio
import json import json
import socket import socket
from unittest.mock import patch from unittest.mock import patch
@@ -12,6 +13,7 @@ import pytest
from nanobot.agent.tools import web as web_module from nanobot.agent.tools import web as web_module
from nanobot.agent.tools.web import WebFetchTool, _get_with_safe_redirects from nanobot.agent.tools.web import WebFetchTool, _get_with_safe_redirects
from nanobot.config.schema import WebFetchConfig from nanobot.config.schema import WebFetchConfig
from nanobot.security.network import PinnedDNSAsyncTransport
from nanobot.security.workspace_access import ( from nanobot.security.workspace_access import (
bind_workspace_scope, bind_workspace_scope,
build_workspace_scope, build_workspace_scope,
@@ -98,27 +100,51 @@ async def test_web_fetch_result_contains_untrusted_flag():
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_safe_redirect_request_pins_validated_dns(monkeypatch): async def test_safe_redirect_requests_use_independent_pinned_dns_concurrently(monkeypatch):
calls: list[str] = [] 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): def _rebinding_resolver(hostname, port, family=0, type_=0, proto=0, flags=0):
calls.append(hostname) host = str(hostname).rstrip(".").lower()
ip = "93.184.216.34" if len(calls) == 1 else "169.254.169.254" 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))] return [(socket.AF_INET, socket.SOCK_STREAM, 0, "", (ip, 0))]
class FakeClient: class ResolvingTransport(httpx.AsyncBaseTransport):
async def get(self, url, headers=None, follow_redirects=False): async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
infos = socket.getaddrinfo("attacker.example", 443, socket.AF_UNSPEC, socket.SOCK_STREAM) await asyncio.sleep(0)
assert infos[0][4][0] == "93.184.216.34" infos = socket.getaddrinfo(
return httpx.Response(200, request=httpx.Request("GET", url)) 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) 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 all(error is None and response is not None for response, error in results)
assert response is not None assert seen == {
assert calls == ["attacker.example"] "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 @pytest.mark.asyncio
@@ -318,6 +344,7 @@ async def test_web_fetch_blocks_private_redirect_before_returning_image(monkeypa
class TransportAsyncClient(real_async_client): class TransportAsyncClient(real_async_client):
def __init__(self, *args, **kwargs): def __init__(self, *args, **kwargs):
kwargs.pop("proxy", None) kwargs.pop("proxy", None)
kwargs.pop("transport", None)
super().__init__(*args, transport=transport, **kwargs) super().__init__(*args, transport=transport, **kwargs)
monkeypatch.setattr("nanobot.agent.tools.web.httpx.AsyncClient", TransportAsyncClient) monkeypatch.setattr("nanobot.agent.tools.web.httpx.AsyncClient", TransportAsyncClient)