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