Fix document query SSRF regression
Route remote document downloads through the existing public-only HTTP fetcher so direct and redirected non-public destinations are rejected while public redirects, request headers, proxy isolation, retries, timeouts, and size limits remain intact. Add focused CVE-2026-4308 regression coverage for private targets and genuine public-fetch compatibility.
Alessandro committed
Aug 12, 2026 at 04:18 UTC
b40874e7c03775c53989e206769e33ff23a4384e
2 files changed
+151
-63
plugins/_document_query/helpers/fetch.py
+31
-63
@@ -12,9 +12,8 @@ from pathlib import Path
12
from typing import Awaitable, Callable
13
from urllib.parse import urlparse
14
15
-import aiohttp
16
-
15
from helpers import files
16
+from helpers.network import fetch_public_http_resource
17
18
19
InterventionCallback = Callable[[], Awaitable[None]]
@@ -140,53 +139,36 @@ async def _fetch_http(
139
last_error = ""
140
for attempt in range(retries):
141
try:
143
- async with aiohttp.ClientSession(
144
- timeout=aiohttp.ClientTimeout(total=timeout)
145
- ) as session:
146
- async with session.get(uri, allow_redirects=True) as response:
147
- if response.status > 399:
148
- raise ValueError(f"HTTP {response.status}")
149
-
150
- content_length = response.headers.get("content-length")
151
- if content_length and int(content_length) > max_remote_bytes:
152
- size_mb = int(content_length) / 1024 / 1024
153
- raise ValueError(
154
- f"Document exceeds max {max_remote_bytes / 1024 / 1024:.0f}MB: "
155
- f"{size_mb:.2f} MB ({uri})"
156
- )
157
-
158
- chunks: list[bytes] = []
159
- downloaded = 0
160
- async for chunk in response.content.iter_chunked(64 * 1024):
161
- downloaded += len(chunk)
162
- if downloaded > max_remote_bytes:
163
- size_mb = downloaded / 1024 / 1024
164
- raise ValueError(
165
- f"Document exceeds max {max_remote_bytes / 1024 / 1024:.0f}MB: "
166
- f"{size_mb:.2f} MB ({uri})"
167
- )
168
- chunks.append(chunk)
169
- if intervention_callback:
170
- await intervention_callback()
171
-
172
- content_type = response.headers.get("content-type", "")
173
- mimetype, charset = _parse_content_type(content_type)
174
- if not mimetype or mimetype == "application/octet-stream":
175
- mimetype = guessed_mimetype or "application/octet-stream"
176
- if mimetype == "application/octet-stream":
177
- raise ValueError(
178
- f"Unsupported document mimetype '{mimetype}' ({uri})"
179
- )
180
-
181
- return FetchedDocument(
182
- uri=str(response.url),
183
- source_uri=uri,
184
- scheme=response.url.scheme or scheme,
185
- mimetype=mimetype,
186
- encoding=encoding,
187
- charset=charset,
188
- content=b"".join(chunks),
189
- )
142
+ if intervention_callback:
143
+ await intervention_callback()
144
+ resource = await asyncio.to_thread(
145
+ fetch_public_http_resource,
146
+ uri,
147
+ max_bytes=max_remote_bytes,
148
+ timeout=(timeout, timeout),
149
+ )
150
+ if intervention_callback:
151
+ await intervention_callback()
152
+
153
+ mimetype = (
154
+ resource.content_type
155
+ or guessed_mimetype
156
+ or "application/octet-stream"
157
+ )
158
+ if mimetype == "application/octet-stream":
159
+ raise ValueError(
160
+ f"Unsupported document mimetype '{mimetype}' ({uri})"
161
+ )
162
+
163
+ return FetchedDocument(
164
+ uri=resource.url,
165
+ source_uri=uri,
166
+ scheme=urlparse(resource.url).scheme or scheme,
167
+ mimetype=mimetype,
168
+ encoding=encoding,
169
+ charset=resource.encoding,
170
+ content=resource.content,
171
+ )
172
except Exception as e:
173
last_error = str(e)
174
if attempt < retries - 1:
@@ -196,20 +178,6 @@ async def _fetch_http(
178
179
raise ValueError(f"Document fetch error: {uri} ({last_error})")
180
199
-
200
-def _parse_content_type(value: str) -> tuple[str | None, str | None]:
201
- if not value:
202
- return None, None
203
- parts = [part.strip() for part in value.split(";") if part.strip()]
204
- mimetype = parts[0].lower() if parts else None
205
- charset = None
206
- for part in parts[1:]:
207
- if part.lower().startswith("charset="):
208
- charset = part.split("=", 1)[1].strip("\"'")
209
- break
210
- return mimetype, charset
211
-
212
-
181
register_protocol_handler("file", _fetch_file)
182
register_protocol_handler("http", _fetch_http)
183
register_protocol_handler("https", _fetch_http)
tests/test_document_query_plugin.py
+120
@@ -1,6 +1,7 @@
1
from __future__ import annotations
2
3
import asyncio
4
+import ipaddress
5
import sys
6
from pathlib import Path
7
@@ -11,6 +12,7 @@ ROOT = Path(__file__).resolve().parents[1]
12
if str(ROOT) not in sys.path:
13
sys.path.insert(0, str(ROOT))
14
15
+from helpers import network as network_helper
16
from plugins._document_query import hooks as document_query_hooks
17
from plugins._document_query.helpers.fetch import FetchedDocument, fetch_public_resource
18
import plugins._document_query.helpers.document_query as document_query_module
@@ -113,6 +115,124 @@ def test_fetch_file_detects_mimetype_and_reads_once(tmp_path):
115
assert fetched.text() == "hello\nworld\n"
116
117
118
+def test_fetch_http_blocks_non_public_destinations():
119
+ with pytest.raises(ValueError, match="Blocked non-public address"):
120
+ run_async(
121
+ fetch_public_resource(
122
+ "http://127.0.0.1/internal.txt",
123
+ {"fetch_retries": 1},
124
+ )
125
+ )
126
+
127
+
128
+def test_fetch_http_blocks_redirects_to_non_public_destinations(monkeypatch):
129
+ source = "https://public.example/start"
130
+ calls = []
131
+
132
+ class RedirectResponse:
133
+ status_code = 302
134
+ headers = {"Location": "http://127.0.0.1/internal.txt"}
135
+ encoding = None
136
+
137
+ def __enter__(self):
138
+ return self
139
+
140
+ def __exit__(self, *_args):
141
+ return False
142
+
143
+ class FakeSession:
144
+ trust_env = True
145
+
146
+ def get(self, url, **kwargs):
147
+ calls.append((url, kwargs, self.trust_env))
148
+ return RedirectResponse()
149
+
150
+ def resolve_host_ips(hostname):
151
+ address = "127.0.0.1" if hostname == "127.0.0.1" else "93.184.216.34"
152
+ return (ipaddress.ip_address(address),)
153
+
154
+ monkeypatch.setattr(network_helper, "resolve_host_ips", resolve_host_ips)
155
+ monkeypatch.setattr(network_helper.requests, "Session", FakeSession)
156
+
157
+ with pytest.raises(ValueError, match="Blocked non-public address"):
158
+ run_async(fetch_public_resource(source, {"fetch_retries": 1}))
159
+
160
+ assert [url for url, _kwargs, _trust_env in calls] == [source]
161
+
162
+
163
+def test_fetch_http_preserves_public_redirects_and_request_compatibility(monkeypatch):
164
+ source = "https://public.example/start"
165
+ destination = "https://public.example/report.txt"
166
+ calls = []
167
+
168
+ class FakeResponse:
169
+ def __init__(self, status_code, headers, body=b""):
170
+ self.status_code = status_code
171
+ self.headers = headers
172
+ self.encoding = "utf-8"
173
+ self.body = body
174
+
175
+ def __enter__(self):
176
+ return self
177
+
178
+ def __exit__(self, *_args):
179
+ return False
180
+
181
+ def iter_content(self, chunk_size):
182
+ assert chunk_size == 64 * 1024
183
+ yield self.body
184
+
185
+ class FakeSession:
186
+ trust_env = True
187
+
188
+ def get(self, url, **kwargs):
189
+ calls.append((url, kwargs, self.trust_env))
190
+ if url == source:
191
+ return FakeResponse(302, {"Location": "/report.txt"})
192
+ return FakeResponse(
193
+ 200,
194
+ {
195
+ "Content-Length": "9",
196
+ "Content-Type": "text/plain; charset=utf-8",
197
+ },
198
+ b"public ok",
199
+ )
200
+
201
+ monkeypatch.setattr(
202
+ network_helper,
203
+ "resolve_host_ips",
204
+ lambda _hostname: (ipaddress.ip_address("93.184.216.34"),),
205
+ )
206
+ monkeypatch.setattr(network_helper.requests, "Session", FakeSession)
207
+ monkeypatch.setenv("USER_AGENT", "AgentZeroTest")
208
+
209
+ fetched = run_async(
210
+ fetch_public_resource(
211
+ source,
212
+ {
213
+ "fetch_retries": 1,
214
+ "fetch_timeout": 2,
215
+ "max_remote_bytes": 1024,
216
+ },
217
+ )
218
+ )
219
+
220
+ assert fetched.uri == destination
221
+ assert fetched.mimetype == "text/plain"
222
+ assert fetched.text() == "public ok"
223
+ assert [url for url, _kwargs, _trust_env in calls] == [source, destination]
224
+ assert all(not trust_env for _url, _kwargs, trust_env in calls)
225
+ assert all(
226
+ kwargs["allow_redirects"] is False
227
+ for _url, kwargs, _trust_env in calls
228
+ )
229
+ assert all(kwargs["timeout"] == (2.0, 2.0) for _url, kwargs, _trust_env in calls)
230
+ assert all(
231
+ kwargs["headers"] == {"User-Agent": "AgentZeroTest"}
232
+ for _url, kwargs, _trust_env in calls
233
+ )
234
+
235
+
236
def test_parser_registry_prefers_liteparse_for_pdf():
237
parsers = get_parsers_for_mimetype("application/pdf", {"liteparse_enabled": True})
238