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