From 4408cde019eeba1ebf75ab1f4ebffaa425e278ae Mon Sep 17 00:00:00 2001 From: Xubin Ren <52506698+Re-bin@users.noreply.github.com> Date: Sun, 26 Jul 2026 16:50:08 +0800 Subject: [PATCH] fix(security): harden generated image downloads --- nanobot/providers/image_generation.py | 106 ++++++++++------ tests/providers/test_image_generation.py | 41 +++++-- .../test_image_generation_security.py | 113 ++++++++++++++++++ 3 files changed, 214 insertions(+), 46 deletions(-) create mode 100644 tests/providers/test_image_generation_security.py diff --git a/nanobot/providers/image_generation.py b/nanobot/providers/image_generation.py index b3843b11..7f824651 100644 --- a/nanobot/providers/image_generation.py +++ b/nanobot/providers/image_generation.py @@ -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 diff --git a/tests/providers/test_image_generation.py b/tests/providers/test_image_generation.py index 0ec35c0e..8504d4d6 100644 --- a/tests/providers/test_image_generation.py +++ b/tests/providers/test_image_generation.py @@ -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 diff --git a/tests/providers/test_image_generation_security.py b/tests/providers/test_image_generation_security.py new file mode 100644 index 00000000..466b74cd --- /dev/null +++ b/tests/providers/test_image_generation_security.py @@ -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), + )