fix(security): harden generated image downloads

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:
+        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:
-        detail = response.text[:500]
-        raise ImageGenerationError(f"failed to download generated image: {detail}") from exc
-    raw = response.content
+                        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

@@ -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

@@ -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),
+        )