mirror of
https://github.com/zhayujie/chatgpt-on-wechat.git
synced 2026-07-17 19:21:26 +08:00
ea47f3097e
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>
189 lines
7.7 KiB
Python
189 lines
7.7 KiB
Python
# 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()
|