kaiyuan/chatgpt-on-wechat
1
0
Fork 0
mirror of https://github.com/zhayujie/chatgpt-on-wechat.git synced 2026-07-17 11:07:11 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(web_fetch): add SSRF guard for model-supplied URLs

Browse Source 
web_fetch fetched any http/https URL a model emitted, checking only the
scheme. It performed requests.get(..., allow_redirects=True) with no
hostname resolution check and no private/loopback/link-local/cloud-metadata
filtering, and never re-validated redirect targets. A model (including one
under prompt injection) could make CowAgent fetch 127.0.0.1, RFC1918,
169.254.169.254 or other internal endpoints and return their bodies into
the conversation; a public URL could also 302-bounce into a private target.

The repo already shipped an SSRF validator for the vision tool
(Vision._validate_url_safe). Extract that logic into a shared helper
(agent/tools/utils/url_safety.py) and reuse it:

- execute() now validates the URL before dispatching to either fetch path.
- A new _safe_get() helper disables auto-redirect and follows redirects
  manually, re-validating every hop so a public URL cannot bounce into an
  internal address. Both the webpage and document fetch paths use it.
- Vision._validate_url_safe now delegates to the shared helper (public API
  unchanged), so both URL-consuming tools share one guard.

Stdlib only (ipaddress, socket, urllib.parse); no new dependency. Adds
tests/test_security_ssrf_web_fetch.py covering loopback, cloud-metadata,
RFC1918 and a public->loopback redirect.

Sink: agent/tools/web_fetch/web_fetch.py (_fetch_webpage / _fetch_document).
Signed-off-by: christop <825583681@qq.com>
This commit is contained in:
christop
2026-06-17 00:59:53 +08:00
parent e3dce45b2a
commit ea47f3097e
5 changed files with 325 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
agent/tools/utils/__init__.py
View File

@@ -20,6 +20,11 @@ from .diff import (
FuzzyMatchResult
)
from .url_safety import (
validate_url_safe,
assert_public_ip
)
__all__ = [
'truncate_head',
'truncate_tail',
@@ -36,5 +41,7 @@ __all__ = [
'normalize_for_fuzzy_match',
'fuzzy_find_text',
'generate_diff_string',
'FuzzyMatchResult'
'FuzzyMatchResult',
'validate_url_safe',
'assert_public_ip'
]

66
agent/tools/utils/url_safety.py Normal file
View File

@@ -0,0 +1,66 @@
"""
Shared SSRF guard utilities for tools that fetch model-supplied URLs.
A URL is only considered safe when it uses an http/https scheme, has a
hostname, that hostname resolves, and every resolved address is a public
(internet-routable) address. Loopback, private (RFC1918 / ULA), link-local
(incl. the 169.254.169.254 cloud-metadata endpoint) and otherwise reserved
addresses are rejected, for both IPv4 and IPv6.
"""
import ipaddress
import socket
from urllib.parse import urlparse
def _is_blocked_ip(ip: "ipaddress._BaseAddress") -> bool:
"""Return True if the address is not safe to connect to (non-public)."""
return (
ip.is_private
or ip.is_loopback
or ip.is_link_local
or ip.is_reserved
or ip.is_multicast
or ip.is_unspecified
)
def assert_public_ip(ip_str: str) -> None:
"""Raise ValueError if the given literal IP is a non-public address.
Used to re-validate the concrete address a redirect resolved to.
"""
ip = ipaddress.ip_address(ip_str)
if _is_blocked_ip(ip):
raise ValueError(
f"URL resolves to a non-public address ({ip_str}), "
f"request blocked for security"
)
def validate_url_safe(url: str) -> None:
"""Reject URLs that target private/loopback/link-local addresses (SSRF guard).
Resolves the hostname to its IP address(es) and blocks any that fall
into non-public ranges. Also rejects URLs with no host, non-HTTP(S)
schemes, or hosts that fail DNS resolution.
Raises:
ValueError: if the URL targets a disallowed address.
"""
parsed = urlparse(url)
if parsed.scheme not in ("http", "https"):
raise ValueError(f"Unsupported URL scheme: {parsed.scheme}")
hostname = parsed.hostname
if not hostname:
raise ValueError("URL has no hostname")
try:
# Resolve all addresses for the hostname.
addr_infos = socket.getaddrinfo(hostname, None, socket.AF_UNSPEC, socket.SOCK_STREAM)
except socket.gaierror:
raise ValueError(f"Cannot resolve hostname: {hostname}")
for family, _, _, _, sockaddr in addr_infos:
assert_public_ip(sockaddr[0])

30
agent/tools/vision/vision.py
View File

@@ -17,18 +17,16 @@ Provider resolution:
"""
import base64
import ipaddress
import os
import socket
import subprocess
import tempfile
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional
from urllib.parse import urlparse
import requests
from agent.tools.base_tool import BaseTool, ToolResult
from agent.tools.utils.url_safety import validate_url_safe
from common import const
from common.log import logger
from config import conf
@@ -665,31 +663,13 @@ class Vision(BaseTool):
into non-public ranges. Also rejects URLs with no host, non-HTTP(S)
schemes, or hosts that fail DNS resolution.
Delegates to the shared ``agent.tools.utils.url_safety`` helper so the
same guard protects every tool that fetches model-supplied URLs.
Raises:
ValueError: if the URL targets a disallowed address.
"""
parsed = urlparse(url)
if parsed.scheme not in ("http", "https"):
raise ValueError(f"Unsupported URL scheme: {parsed.scheme}")
hostname = parsed.hostname
if not hostname:
raise ValueError("URL has no hostname")
try:
# Resolve all addresses for the hostname.
addr_infos = socket.getaddrinfo(hostname, None, socket.AF_UNSPEC, socket.SOCK_STREAM)
except socket.gaierror:
raise ValueError(f"Cannot resolve hostname: {hostname}")
for family, _, _, _, sockaddr in addr_infos:
ip_str = sockaddr[0]
ip = ipaddress.ip_address(ip_str)
if ip.is_private or ip.is_loopback or ip.is_link_local or ip.is_reserved:
raise ValueError(
f"URL resolves to a non-public address ({ip_str}), "
f"request blocked for security"
)
validate_url_safe(url)
def _build_image_content(self, image: str) -> dict:
"""

71
agent/tools/web_fetch/web_fetch.py
View File

@@ -16,11 +16,15 @@ import requests
from agent.tools.base_tool import BaseTool, ToolResult
from agent.tools.utils.truncate import truncate_head, format_size
from agent.tools.utils.url_safety import validate_url_safe
from common.log import logger
DEFAULT_TIMEOUT = 30
MAX_FILE_SIZE = 50 * 1024 * 1024 # 50MB
# Cap on how many redirects we follow; each hop's target is re-validated
# against the SSRF guard so a public URL cannot bounce us into an internal one.
MAX_REDIRECTS = 10
DEFAULT_HEADERS = {
"User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36",
@@ -107,23 +111,65 @@ class WebFetch(BaseTool):
if parsed.scheme not in ("http", "https"):
return ToolResult.fail("Error: Invalid URL (must start with http:// or https://)")
# SSRF guard: reject URLs that resolve to private/loopback/link-local/
# cloud-metadata addresses before any request is issued.
try:
validate_url_safe(url)
except ValueError as e:
return ToolResult.fail(f"Error: {e}")
if _is_document_url(url):
return self._fetch_document(url)
return self._fetch_webpage(url)
# ---- Safe request helper ----
@staticmethod
def _safe_get(url: str, **kwargs) -> requests.Response:
"""Issue a GET request while re-validating every redirect hop (SSRF guard).
Auto-redirect is disabled and each hop is followed manually so the
target of every redirect is re-resolved and checked against the SSRF
guard. This prevents a public URL from 3xx-bouncing into a private,
loopback, link-local or cloud-metadata address. ``kwargs`` are passed
through to ``requests.get`` (e.g. ``stream``).
Raises:
ValueError: if any hop resolves to a non-public address.
"""
kwargs.pop("allow_redirects", None)
current = url
for _ in range(MAX_REDIRECTS + 1):
response = requests.get(
current,
headers=DEFAULT_HEADERS,
timeout=DEFAULT_TIMEOUT,
allow_redirects=False,
**kwargs,
)
if not response.is_redirect and not response.is_permanent_redirect:
return response
location = response.headers.get("Location")
if not location:
return response
# Resolve the redirect target relative to the current URL, then
# re-validate it before following.
current = requests.compat.urljoin(current, location)
validate_url_safe(current)
response.close()
raise ValueError(f"Too many redirects (>{MAX_REDIRECTS})")
# ---- Web page fetching ----
def _fetch_webpage(self, url: str) -> ToolResult:
"""Fetch and extract readable text from an HTML web page."""
parsed = urlparse(url)
try:
response = requests.get(
url,
headers=DEFAULT_HEADERS,
timeout=DEFAULT_TIMEOUT,
allow_redirects=True,
)
response = self._safe_get(url)
response.raise_for_status()
except requests.Timeout:
return ToolResult.fail(f"Error: Request timed out after {DEFAULT_TIMEOUT}s")
@@ -131,6 +177,8 @@ class WebFetch(BaseTool):
return ToolResult.fail(f"Error: Failed to connect to {parsed.netloc}")
except requests.HTTPError as e:
return ToolResult.fail(f"Error: HTTP {e.response.status_code} for URL: {url}")
except ValueError as e:
return ToolResult.fail(f"Error: {e}")
except Exception as e:
return ToolResult.fail(f"Error: Failed to fetch URL: {e}")
@@ -158,13 +206,7 @@ class WebFetch(BaseTool):
logger.info(f"[WebFetch] Downloading document: {url} -> {local_path}")
try:
response = requests.get(
url,
headers=DEFAULT_HEADERS,
timeout=DEFAULT_TIMEOUT,
stream=True,
allow_redirects=True,
)
response = self._safe_get(url, stream=True)
response.raise_for_status()
content_length = int(response.headers.get("Content-Length", 0))
@@ -191,6 +233,9 @@ class WebFetch(BaseTool):
return ToolResult.fail(f"Error: Failed to connect to {parsed.netloc}")
except requests.HTTPError as e:
return ToolResult.fail(f"Error: HTTP {e.response.status_code} for URL: {url}")
except ValueError as e:
self._cleanup_file(local_path)
return ToolResult.fail(f"Error: {e}")
except Exception as e:
self._cleanup_file(local_path)
return ToolResult.fail(f"Error: Failed to download file: {e}")

188
tests/test_security_ssrf_web_fetch.py Normal file
View File

@@ -0,0 +1,188 @@
# encoding:utf-8
"""
Regression tests for web_fetch SSRF protection.
The web_fetch tool fetches model-supplied URLs. Without a guard, a model
(including one under prompt injection) can point it at loopback, RFC1918,
link-local or cloud-metadata (169.254.169.254) endpoints, or use a public
URL that 3xx-redirects into such a target. These tests ensure web_fetch
refuses the request instead of connecting to the internal address.
No real network is used: DNS resolution and ``requests.get`` are stubbed.
"""
import os
import sys
import types
import unittest
from unittest.mock import patch, MagicMock
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
# Stub 'requests' if not installed so the module can be imported for testing.
if "requests" not in sys.modules:
_requests_stub = types.ModuleType("requests")
_requests_stub.get = lambda *a, **k: None
class _Exc(Exception):
pass
_requests_stub.Timeout = type("Timeout", (_Exc,), {})
_requests_stub.ConnectionError = type("ConnectionError", (_Exc,), {})
_requests_stub.HTTPError = type("HTTPError", (_Exc,), {})
_requests_stub.Response = object
_compat = types.SimpleNamespace(urljoin=__import__("urllib.parse", fromlist=["urljoin"]).urljoin)
_requests_stub.compat = _compat
sys.modules["requests"] = _requests_stub
def _gai(ip_str):
"""Build a socket.getaddrinfo return value for a single IPv4 address."""
return [(2, 1, 6, "", (ip_str, 0))]
class _FakeRedirect:
"""Minimal stand-in for a requests redirect Response."""
def __init__(self, location):
self.is_redirect = True
self.is_permanent_redirect = False
self.headers = {"Location": location}
self.closed = False
def close(self):
self.closed = True
def _fake_ok_response(body=b"<html><head><title>internal</title></head><body>secret</body></html>"):
"""A well-formed non-redirect response.
Returned by the mocked ``requests.get`` so that on UNPATCHED code the
fetch path runs to completion and the test fails specifically on the
``assert_not_called`` guard (proving a request reached the internal
target), rather than on an incidental error.
"""
resp = MagicMock()
resp.is_redirect = False
resp.is_permanent_redirect = False
resp.status_code = 200
resp.headers = {"Content-Type": "text/html; charset=utf-8"}
resp.content = body
resp.text = body.decode("utf-8")
resp.apparent_encoding = "utf-8"
resp.raise_for_status = lambda: None
return resp
class TestWebFetchSSRF(unittest.TestCase):
"""web_fetch must refuse internal targets and never connect to them."""
def setUp(self):
from agent.tools.web_fetch.web_fetch import WebFetch
self.tool = WebFetch()
# --- Literal internal IPs: rejected before any socket call ---
def test_loopback_literal_blocked(self):
"""http://127.0.0.1:<port>/x must be refused, no request issued."""
with patch("requests.get", return_value=_fake_ok_response()) as mock_get:
result = self.tool.execute({"url": "http://127.0.0.1:8080/canary"})
self.assertEqual(result.status, "error")
self.assertIn("non-public", result.result)
mock_get.assert_not_called()
def test_cloud_metadata_literal_blocked(self):
"""http://169.254.169.254/latest/meta-data/ must be refused."""
with patch("requests.get", return_value=_fake_ok_response()) as mock_get:
result = self.tool.execute(
{"url": "http://169.254.169.254/latest/meta-data/"}
)
self.assertEqual(result.status, "error")
self.assertIn("non-public", result.result)
mock_get.assert_not_called()
def test_ipv6_loopback_literal_blocked(self):
"""http://[::1]/x must be refused."""
with patch("requests.get", return_value=_fake_ok_response()) as mock_get:
result = self.tool.execute({"url": "http://[::1]/canary"})
self.assertEqual(result.status, "error")
self.assertIn("non-public", result.result)
mock_get.assert_not_called()
# --- RFC1918 host resolved via DNS: rejected after resolution ---
def test_rfc1918_hostname_blocked(self):
"""A hostname that resolves to 10.x.x.x must be refused, no request."""
with patch("socket.getaddrinfo", return_value=_gai("10.1.2.3")), \
patch("requests.get", return_value=_fake_ok_response()) as mock_get:
result = self.tool.execute({"url": "http://internal.corp/secret"})
self.assertEqual(result.status, "error")
self.assertIn("non-public", result.result)
mock_get.assert_not_called()
def test_192_168_hostname_blocked(self):
"""A hostname that resolves to 192.168.x.x must be refused."""
with patch("socket.getaddrinfo", return_value=_gai("192.168.0.5")), \
patch("requests.get", return_value=_fake_ok_response()) as mock_get:
result = self.tool.execute({"url": "http://router.local/admin"})
self.assertEqual(result.status, "error")
self.assertIn("non-public", result.result)
mock_get.assert_not_called()
# --- Redirect bounce: public entry URL 302 -> loopback ---
def test_public_to_loopback_redirect_blocked(self):
"""A public URL that redirects to a loopback target must be refused.
The first hop resolves to a public IP and returns a 302 pointing at
127.0.0.1; the guard must re-validate the redirect target and refuse
instead of fetching the internal address.
"""
redirect = _FakeRedirect("http://127.0.0.1:8080/canary")
def fake_getaddrinfo(host, *a, **k):
# Public entry host resolves to a public IP; the loopback literal
# echoes back (as the real getaddrinfo does for an IP literal).
if host == "evil.example.com":
return _gai("93.184.216.34")
return _gai(host)
with patch("socket.getaddrinfo", side_effect=fake_getaddrinfo), \
patch("requests.get", return_value=redirect) as mock_get:
result = self.tool.execute({"url": "http://evil.example.com/start"})
self.assertEqual(result.status, "error")
self.assertIn("non-public", result.result)
# The first (public) hop is issued exactly once; the loopback hop is
# rejected by the guard BEFORE a second requests.get to the internal
# target is made.
self.assertEqual(mock_get.call_count, 1)
first_call_url = mock_get.call_args[0][0]
self.assertEqual(first_call_url, "http://evil.example.com/start")
# The follow-up request to the internal target was never issued.
for call in mock_get.call_args_list:
self.assertNotIn("127.0.0.1", call[0][0])
# --- Sanity: a public URL is allowed to proceed to the fetch path ---
def test_public_url_allowed_through_guard(self):
"""A public URL passes the guard and a (mocked) request is issued."""
ok = MagicMock()
ok.is_redirect = False
ok.is_permanent_redirect = False
ok.headers = {"Content-Type": "text/html; charset=utf-8"}
ok.content = b"<html><head><title>Hi</title></head><body>ok</body></html>"
ok.text = "<html><head><title>Hi</title></head><body>ok</body></html>"
ok.apparent_encoding = "utf-8"
ok.raise_for_status = lambda: None
with patch("socket.getaddrinfo", return_value=_gai("93.184.216.34")), \
patch("requests.get", return_value=ok) as mock_get:
result = self.tool.execute({"url": "http://example.com/page"})
self.assertEqual(result.status, "success")
mock_get.assert_called_once()
self.assertEqual(mock_get.call_args[0][0], "http://example.com/page")
if __name__ == "__main__":
unittest.main()