mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-08-07 01:48:53 +00:00
fix: serialize pinned dns web fetches
This commit is contained in:
parent
4353f4680b
commit
97e3b360c2
@ -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,14 +132,11 @@ 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
|
response = await client.get(current_url, headers=headers, follow_redirects=False)
|
||||||
|
|
||||||
with pin_resolved_url_dns(current_url, resolved_ips):
|
|
||||||
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:
|
||||||
return response, None
|
return response, None
|
||||||
@ -162,20 +165,17 @@ 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:
|
||||||
return response, stream, None
|
return response, stream, None
|
||||||
@ -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,
|
||||||
|
|||||||
@ -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,17 +155,26 @@ 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)
|
||||||
with pin_resolved_url_dns(url, resolved_ips):
|
async with self._resolver_lock:
|
||||||
return await self._inner.handle_async_request(request)
|
with pin_resolved_url_dns(url, resolved_ips):
|
||||||
|
return await self._inner.handle_async_request(request)
|
||||||
|
|
||||||
async def aclose(self) -> None:
|
async def aclose(self) -> None:
|
||||||
await self._inner.aclose()
|
await self._inner.aclose()
|
||||||
|
|||||||
@ -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)
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user