Grayscale/nanobot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(security): harden generated image downloads

Browse Source
This commit is contained in:
Xubin Ren
2026-07-27 10:06:19 +08:00
parent cf1e801a29
commit 4408cde019
3 changed files with 214 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nanobot/providers/image_generation.py
+70 -36
View File
@@ -10,11 +10,13 @@ from abc import ABC, abstractmethod
from dataclasses import dataclass
from pathlib import Path
from typing import Any
from urllib.parse import urljoin
import httpx
from loguru import logger
from nanobot.providers.registry import find_by_name
from nanobot.security.network import PinnedDNSAsyncTransport, UnsafeURLRequestError
from nanobot.utils.helpers import detect_image_mime
_OPENROUTER_ATTRIBUTION_HEADERS = {
@@ -23,6 +25,8 @@ _OPENROUTER_ATTRIBUTION_HEADERS = {
"X-OpenRouter-Categories": "cli-agent,personal-agent",
}
_DEFAULT_TIMEOUT_S = 120.0
_IMAGE_DOWNLOAD_MAX_BYTES = 32 * 1024 * 1024
_IMAGE_DOWNLOAD_MAX_REDIRECTS = 5
_AIHUBMIX_TIMEOUT_S = 300.0
_AIHUBMIX_ASPECT_RATIO_SIZES = {
"1:1": "1024x1024",
@@ -131,16 +135,66 @@ def _aihubmix_model_path(model: str) -> str:
async def _download_image_data_url(
client: httpx.AsyncClient,
url: str,
*,
transport: httpx.AsyncBaseTransport | None = None,
) -> str:
response = await client.get(url)
try:
response.raise_for_status()
except httpx.HTTPStatusError as exc:
detail = response.text[:500]
raise ImageGenerationError(f"failed to download generated image: {detail}") from exc
raw = response.content
safe_transport = PinnedDNSAsyncTransport(inner=transport)
# Proxies resolve the target independently and would defeat DNS pinning.
async with httpx.AsyncClient(
transport=safe_transport,
follow_redirects=False,
timeout=_DEFAULT_TIMEOUT_S,
trust_env=False,
) as client:
current_url = url
for _ in range(_IMAGE_DOWNLOAD_MAX_REDIRECTS + 1):
async with client.stream("GET", current_url) as response:
if response.is_redirect:
location = response.headers.get("location")
if not location:
raise ImageGenerationError(
"generated image URL redirected without a location"
)
current_url = urljoin(str(response.url), location)
continue
try:
response.raise_for_status()
except httpx.HTTPStatusError as exc:
raise ImageGenerationError(
f"failed to download generated image (HTTP {response.status_code})"
) from exc
declared_size = response.headers.get("content-length")
if declared_size:
try:
if int(declared_size) > _IMAGE_DOWNLOAD_MAX_BYTES:
raise ImageGenerationError(
"generated image exceeded the 32 MiB download limit"
)
except ValueError:
pass
chunks: list[bytes] = []
total = 0
async for chunk in response.aiter_bytes():
total += len(chunk)
if total > _IMAGE_DOWNLOAD_MAX_BYTES:
raise ImageGenerationError(
"generated image exceeded the 32 MiB download limit"
)
chunks.append(chunk)
raw = b"".join(chunks)
break
else:
raise ImageGenerationError("generated image URL exceeded the redirect limit")
except UnsafeURLRequestError as exc:
raise ImageGenerationError(f"blocked unsafe generated image URL: {exc}") from exc
except httpx.RequestError as exc:
raise ImageGenerationError(f"failed to download generated image: {exc}") from exc
mime = detect_image_mime(raw)
if mime is None:
raise ImageGenerationError("generated image URL did not return a supported image")
@@ -452,7 +506,7 @@ class AIHubMixImageGenerationClient(ImageGenerationProvider):
raise ImageGenerationError(f"AIHubMix image generation failed: {detail}") from exc
payload = response.json()
images = await _aihubmix_images_from_payload(client, payload)
images = await _aihubmix_images_from_payload(payload)
self._require_images(images, payload)
@@ -827,7 +881,6 @@ def _gemini_flash_supported_image_sizes(model: str) -> set[str]:
async def _aihubmix_images_from_payload(
client: httpx.AsyncClient,
payload: dict[str, Any],
) -> list[str]:
images: list[str] = []
@@ -846,7 +899,7 @@ async def _aihubmix_images_from_payload(
if value.startswith("data:image/"):
images.append(value)
elif value.startswith(("http://", "https://")):
images.append(await _download_image_data_url(client, value))
images.append(await _download_image_data_url(value))
return
if not isinstance(value, dict):
return
@@ -1047,15 +1100,7 @@ class OpenAIImageGenerationClient(ImageGenerationProvider):
return model
async def _parse_images_response(self, payload: dict[str, Any]) -> list[str]:
client = self._client
owns_client = client is None
if owns_client:
client = httpx.AsyncClient(timeout=self.timeout)
try:
return await _openai_images_from_payload(client, payload)
finally:
if owns_client:
await client.aclose()
return await _openai_images_from_payload(payload)
async def _post_image_edit(
self,
@@ -1266,15 +1311,7 @@ class CustomImageGenerationClient(ImageGenerationProvider):
logger.info("Custom Images API response ({}): {}", response.status_code,
{k: v for k, v in payload.items() if k != "data"})
client = self._client
owns_client = client is None
if owns_client:
client = httpx.AsyncClient(timeout=self.timeout)
try:
images = await _openai_images_from_payload(client, payload)
finally:
if owns_client:
await client.aclose()
images = await _openai_images_from_payload(payload)
self._require_images(images, payload)
@@ -1467,7 +1504,6 @@ def _openai_explicit_size_supported(
async def _openai_images_from_payload(
client: httpx.AsyncClient,
payload: dict[str, Any],
) -> list[str]:
"""Extract images from OpenAI Images API response.
@@ -1484,7 +1520,7 @@ async def _openai_images_from_payload(
continue
url = item.get("url")
if isinstance(url, str) and url:
images.append(await _download_image_data_url(client, url))
images.append(await _download_image_data_url(url))
return images
@@ -1798,7 +1834,7 @@ class ZhipuImageGenerationClient(ImageGenerationProvider):
raise ImageGenerationError(f"Zhipu image generation failed: {detail}") from exc
payload = response.json()
images = await _zhipu_images_from_payload(client, payload)
images = await _zhipu_images_from_payload(payload)
self._require_images(images, payload)
@@ -1822,7 +1858,6 @@ def _zhipu_size(
async def _zhipu_images_from_payload(
client: httpx.AsyncClient,
payload: dict[str, Any],
) -> list[str]:
"""Extract image data URLs from Zhipu API response.
@@ -1836,7 +1871,7 @@ async def _zhipu_images_from_payload(
continue
url = item.get("url")
if isinstance(url, str) and url:
images.append(await _download_image_data_url(client, url))
images.append(await _download_image_data_url(url))
return images
@@ -1999,7 +2034,7 @@ class ModelScopeImageGenerationClient(ImageGenerationProvider):
status = data.get("task_status")
if status == "SUCCEED":
return await self._collect_images(client, data)
return await self._collect_images(data)
if status == "FAILED":
raise ImageGenerationError(
f"ModelScope image generation task failed: {data}"
@@ -2014,7 +2049,6 @@ class ModelScopeImageGenerationClient(ImageGenerationProvider):
@staticmethod
async def _collect_images(
client: httpx.AsyncClient,
data: dict[str, Any],
) -> list[str]:
images: list[str] = []
@@ -2023,7 +2057,7 @@ class ModelScopeImageGenerationClient(ImageGenerationProvider):
if url.startswith("data:image/"):
images.append(url)
else:
images.append(await _download_image_data_url(client, url))
images.append(await _download_image_data_url(url))
return images
tests/providers/test_image_generation.py
+31 -10
View File
@@ -102,6 +102,22 @@ class CodexStreamingCompleteThenErrorResponse(FakeResponse):
)
@pytest.fixture(autouse=True)
def generated_image_downloads(monkeypatch) -> list[str]:
"""Keep provider response parsing tests independent from outbound HTTP."""
urls: list[str] = []
async def download(url: str) -> str:
urls.append(url)
return PNG_DATA_URL
monkeypatch.setattr(
"nanobot.providers.image_generation._download_image_data_url",
download,
)
return urls
@pytest.mark.asyncio
async def test_openrouter_image_generation_payload_and_response(tmp_path: Path) -> None:
ref = tmp_path / "ref.png"
@@ -277,7 +293,9 @@ async def test_aihubmix_image_edit_payload_uses_reference_images(tmp_path: Path)
@pytest.mark.asyncio
async def test_aihubmix_image_generation_downloads_url_response() -> None:
async def test_aihubmix_image_generation_downloads_url_response(
generated_image_downloads: list[str],
) -> None:
fake = FakeClient(FakeResponse({"data": [{"url": "https://cdn.example/image.png"}]}))
fake.get_response = FakeResponse({}, content=PNG_BYTES)
client = AIHubMixImageGenerationClient(
@@ -288,7 +306,7 @@ async def test_aihubmix_image_generation_downloads_url_response() -> None:
response = await client.generate(prompt="draw", model="gpt-image-2-free")
assert response.images[0].startswith("data:image/png;base64,")
assert fake.get_calls[0]["url"] == "https://cdn.example/image.png"
assert generated_image_downloads == ["https://cdn.example/image.png"]
@pytest.mark.asyncio
@@ -818,7 +836,7 @@ async def test_openai_b64_json_response_uses_detected_mime() -> None:
@pytest.mark.asyncio
async def test_openai_url_download_fallback() -> None:
async def test_openai_url_download_fallback(generated_image_downloads: list[str]) -> None:
fake = FakeClient(FakeResponse({"data": [{"url": "https://cdn.example/image.png"}]}))
fake.get_response = FakeResponse({}, content=PNG_BYTES)
client = OpenAIImageGenerationClient(
@@ -829,7 +847,7 @@ async def test_openai_url_download_fallback() -> None:
response = await client.generate(prompt="draw", model="dall-e-3")
assert response.images[0].startswith("data:image/png;base64,")
assert fake.get_calls[0]["url"] == "https://cdn.example/image.png"
assert generated_image_downloads == ["https://cdn.example/image.png"]
@pytest.mark.asyncio
@@ -1192,7 +1210,9 @@ async def test_custom_generate_maps_one_k_to_openai_dimension() -> None:
@pytest.mark.asyncio
async def test_custom_generate_extra_body_can_override_defaults() -> None:
async def test_custom_generate_extra_body_can_override_defaults(
generated_image_downloads: list[str],
) -> None:
fake = FakeClient(FakeResponse({"data": [{"url": "https://images.example/cat.png"}]}))
fake.get_response = FakeResponse({}, content=PNG_BYTES)
client = CustomImageGenerationClient(
@@ -1208,9 +1228,8 @@ async def test_custom_generate_extra_body_can_override_defaults() -> None:
image_size="1K",
)
expected_data_url = f"data:image/png;base64,{base64.b64encode(PNG_BYTES).decode('ascii')}"
assert response.images == [expected_data_url]
assert fake.get_calls[0]["url"] == "https://images.example/cat.png"
assert response.images == [PNG_DATA_URL]
assert generated_image_downloads == ["https://images.example/cat.png"]
body = fake.calls[0]["json"]
assert body["response_format"] == "url"
assert body["size"] == "2K"
@@ -1616,7 +1635,9 @@ async def test_zhipu_image_generation_with_explicit_size() -> None:
@pytest.mark.asyncio
async def test_zhipu_image_generation_downloads_url_response() -> None:
async def test_zhipu_image_generation_downloads_url_response(
generated_image_downloads: list[str],
) -> None:
fake = FakeClient(FakeResponse({"data": [{"url": "https://cdn.example/image.png"}]}))
fake.get_response = FakeResponse({}, content=PNG_BYTES)
client = ZhipuImageGenerationClient(
@@ -1627,7 +1648,7 @@ async def test_zhipu_image_generation_downloads_url_response() -> None:
response = await client.generate(prompt="draw", model="glm-image")
assert response.images[0].startswith("data:image/png;base64,")
assert fake.get_calls[0]["url"] == "https://cdn.example/image.png"
assert generated_image_downloads == ["https://cdn.example/image.png"]
@pytest.mark.asyncio
tests/providers/test_image_generation_security.py
+113
View File
@@ -0,0 +1,113 @@
from __future__ import annotations
import socket
import httpx
import pytest
from nanobot.providers import image_generation
from nanobot.providers.image_generation import ImageGenerationError, _download_image_data_url
PNG_BYTES = (
b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01"
b"\x00\x00\x00\x01\x08\x04\x00\x00\x00\xb5\x1c\x0c\x02"
b"\x00\x00\x00\x0bIDATx\xdacd\xfc\xff\x1f\x00\x03\x03"
b"\x02\x00\xef\xbf\xa7\xdb\x00\x00\x00\x00IEND\xaeB`\x82"
)
def _resolve_public(host: str, port: int | None, *args, **kwargs):
return [
(
socket.AF_INET,
socket.SOCK_STREAM,
socket.IPPROTO_TCP,
"",
("93.184.216.34", port or 0),
)
]
@pytest.mark.asyncio
async def test_generated_image_download_blocks_private_target() -> None:
requested = False
async def handler(request: httpx.Request) -> httpx.Response:
nonlocal requested
requested = True
return httpx.Response(200, content=PNG_BYTES)
with pytest.raises(ImageGenerationError, match="blocked unsafe generated image URL"):
await _download_image_data_url(
"http://127.0.0.1/admin",
transport=httpx.MockTransport(handler),
)
assert requested is False
@pytest.mark.asyncio
async def test_generated_image_download_revalidates_redirects(monkeypatch) -> None:
original_getaddrinfo = socket.getaddrinfo
def resolve_test_hosts(host: str, port: int | None, *args, **kwargs):
if host == "cdn.example":
return _resolve_public(host, port, *args, **kwargs)
return original_getaddrinfo(host, port, *args, **kwargs)
monkeypatch.setattr("nanobot.security.network.socket.getaddrinfo", resolve_test_hosts)
requested: list[str] = []
async def handler(request: httpx.Request) -> httpx.Response:
requested.append(str(request.url))
return httpx.Response(302, headers={"location": "http://169.254.169.254/latest"})
with pytest.raises(ImageGenerationError, match="blocked unsafe generated image URL"):
await _download_image_data_url(
"https://cdn.example/image.png",
transport=httpx.MockTransport(handler),
)
assert requested == ["https://cdn.example/image.png"]
@pytest.mark.asyncio
async def test_generated_image_download_returns_valid_data_url(monkeypatch) -> None:
monkeypatch.setattr(
"nanobot.security.network.socket.getaddrinfo",
_resolve_public,
)
async def handler(request: httpx.Request) -> httpx.Response:
return httpx.Response(200, content=PNG_BYTES)
result = await _download_image_data_url(
"https://cdn.example/image.png",
transport=httpx.MockTransport(handler),
)
assert result.startswith("data:image/png;base64,")
class _OversizedStream(httpx.AsyncByteStream):
async def __aiter__(self):
yield b"12345"
yield b"6789"
@pytest.mark.asyncio
async def test_generated_image_download_enforces_streaming_size_limit(monkeypatch) -> None:
monkeypatch.setattr(
"nanobot.security.network.socket.getaddrinfo",
_resolve_public,
)
monkeypatch.setattr(image_generation, "_IMAGE_DOWNLOAD_MAX_BYTES", 8)
async def handler(request: httpx.Request) -> httpx.Response:
return httpx.Response(200, stream=_OversizedStream())
with pytest.raises(ImageGenerationError, match="download limit"):
await _download_image_data_url(
"https://cdn.example/image.png",
transport=httpx.MockTransport(handler),
)