From d73794bc688971428c5612bec90e19234ba81f2a Mon Sep 17 00:00:00 2001 From: chengyongru Date: Mon, 27 Jul 2026 00:33:58 +0800 Subject: [PATCH] fix(image): honor provider proxy for URL downloads --- docs/image-generation.md | 3 + nanobot/providers/image_generation.py | 81 ++++++++----- tests/providers/test_image_generation.py | 58 ++++++--- .../test_image_generation_security.py | 111 +++++++++++++++++- 4 files changed, 211 insertions(+), 42 deletions(-) diff --git a/docs/image-generation.md b/docs/image-generation.md index bcbb025d..bf727411 100644 --- a/docs/image-generation.md +++ b/docs/image-generation.md @@ -70,6 +70,9 @@ Provider settings reuse normal provider config fields: | `providers..apiBase` | Optional custom base URL | | `providers..extraHeaders` | Headers merged into provider requests | | `providers..extraBody` | Extra JSON fields merged into provider request bodies | +| `providers..proxy` | Explicit trusted HTTP proxy for provider requests and returned image URL downloads | + +For providers that return image URLs, direct downloads use DNS pinning. When an explicit provider `proxy` is configured, nanobot validates the initial URL and every redirect locally, then relies on that trusted proxy for final DNS resolution and network egress. Process-wide proxy environment variables are not used for these downloads. Both camelCase and snake_case config keys are accepted, but docs use camelCase to match `config.json`. diff --git a/nanobot/providers/image_generation.py b/nanobot/providers/image_generation.py index 7f824651..d3ad017e 100644 --- a/nanobot/providers/image_generation.py +++ b/nanobot/providers/image_generation.py @@ -16,7 +16,11 @@ import httpx from loguru import logger from nanobot.providers.registry import find_by_name -from nanobot.security.network import PinnedDNSAsyncTransport, UnsafeURLRequestError +from nanobot.security.network import ( + PinnedDNSAsyncTransport, + UnsafeURLRequestError, + resolve_url_target, +) from nanobot.utils.helpers import detect_image_mime _OPENROUTER_ATTRIBUTION_HEADERS = { @@ -137,19 +141,31 @@ def _aihubmix_model_path(model: str) -> str: async def _download_image_data_url( url: str, *, + proxy: str | None = None, transport: httpx.AsyncBaseTransport | None = None, ) -> str: 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: + client_kwargs: dict[str, Any] = { + "follow_redirects": False, + "timeout": _DEFAULT_TIMEOUT_S, + "trust_env": False, + } + if proxy: + # An explicit provider proxy is a user-selected trusted egress boundary. + # Validate each URL locally, while the proxy owns final DNS resolution. + client_kwargs["proxy"] = proxy + else: + client_kwargs["transport"] = PinnedDNSAsyncTransport(inner=transport) + + async with httpx.AsyncClient(**client_kwargs) as client: current_url = url for _ in range(_IMAGE_DOWNLOAD_MAX_REDIRECTS + 1): + if proxy: + ok, error, _ = resolve_url_target(current_url) + if not ok: + raise ImageGenerationError( + f"blocked unsafe generated image URL: {error}" + ) async with client.stream("GET", current_url) as response: if response.is_redirect: location = response.headers.get("location") @@ -302,6 +318,13 @@ class ImageGenerationProvider(ABC): raise ImageGenerationError(f"{label} returned no images: {provider_error}") raise ImageGenerationError(f"{label} returned no images for this request") + def _http_client_kwargs(self) -> dict[str, Any]: + kwargs: dict[str, Any] = {"timeout": self.timeout} + if self.proxy: + kwargs["proxy"] = self.proxy + kwargs["trust_env"] = False + return kwargs + async def _http_post( self, url: str, @@ -314,11 +337,7 @@ class ImageGenerationProvider(ABC): return await client.post(url, headers=headers, json=body) if self._client is not None: return await self._client.post(url, headers=headers, json=body) - client_kwargs: dict[str, Any] = {"timeout": self.timeout} - if self.proxy: - client_kwargs["proxy"] = self.proxy - client_kwargs["trust_env"] = False - async with httpx.AsyncClient(**client_kwargs) as c: + async with httpx.AsyncClient(**self._http_client_kwargs()) as c: return await c.post(url, headers=headers, json=body) @@ -446,7 +465,7 @@ class AIHubMixImageGenerationClient(ImageGenerationProvider): } size = _aihubmix_size(aspect_ratio, image_size) - client = self._client or httpx.AsyncClient(timeout=self.timeout) + client = self._client or httpx.AsyncClient(**self._http_client_kwargs()) try: return await self._generate_with_client( client, @@ -506,7 +525,7 @@ class AIHubMixImageGenerationClient(ImageGenerationProvider): raise ImageGenerationError(f"AIHubMix image generation failed: {detail}") from exc payload = response.json() - images = await _aihubmix_images_from_payload(payload) + images = await _aihubmix_images_from_payload(payload, proxy=self.proxy) self._require_images(images, payload) @@ -882,6 +901,8 @@ def _gemini_flash_supported_image_sizes(model: str) -> set[str]: async def _aihubmix_images_from_payload( payload: dict[str, Any], + *, + proxy: str | None = None, ) -> list[str]: images: list[str] = [] candidates: list[Any] = [] @@ -899,7 +920,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(value)) + images.append(await _download_image_data_url(value, proxy=proxy)) return if not isinstance(value, dict): return @@ -1100,7 +1121,7 @@ class OpenAIImageGenerationClient(ImageGenerationProvider): return model async def _parse_images_response(self, payload: dict[str, Any]) -> list[str]: - return await _openai_images_from_payload(payload) + return await _openai_images_from_payload(payload, proxy=self.proxy) async def _post_image_edit( self, @@ -1130,7 +1151,7 @@ class OpenAIImageGenerationClient(ImageGenerationProvider): data=body, files=files, ) - async with httpx.AsyncClient(timeout=self.timeout) as c: + async with httpx.AsyncClient(**self._http_client_kwargs()) as c: return await c.post( f"{self.api_base}/images/edits", headers=headers, @@ -1311,7 +1332,7 @@ class CustomImageGenerationClient(ImageGenerationProvider): logger.info("Custom Images API response ({}): {}", response.status_code, {k: v for k, v in payload.items() if k != "data"}) - images = await _openai_images_from_payload(payload) + images = await _openai_images_from_payload(payload, proxy=self.proxy) self._require_images(images, payload) @@ -1505,6 +1526,8 @@ def _openai_explicit_size_supported( async def _openai_images_from_payload( payload: dict[str, Any], + *, + proxy: str | None = None, ) -> list[str]: """Extract images from OpenAI Images API response. @@ -1520,7 +1543,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(url)) + images.append(await _download_image_data_url(url, proxy=proxy)) return images @@ -1800,7 +1823,7 @@ class ZhipuImageGenerationClient(ImageGenerationProvider): url = f"{self.api_base}/images/generations" - client = self._client or httpx.AsyncClient(timeout=self.timeout) + client = self._client or httpx.AsyncClient(**self._http_client_kwargs()) try: return await self._generate_with_client( client, @@ -1834,7 +1857,7 @@ class ZhipuImageGenerationClient(ImageGenerationProvider): raise ImageGenerationError(f"Zhipu image generation failed: {detail}") from exc payload = response.json() - images = await _zhipu_images_from_payload(payload) + images = await _zhipu_images_from_payload(payload, proxy=self.proxy) self._require_images(images, payload) @@ -1859,6 +1882,8 @@ def _zhipu_size( async def _zhipu_images_from_payload( payload: dict[str, Any], + *, + proxy: str | None = None, ) -> list[str]: """Extract image data URLs from Zhipu API response. @@ -1871,7 +1896,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(url)) + images.append(await _download_image_data_url(url, proxy=proxy)) return images @@ -1957,7 +1982,7 @@ class ModelScopeImageGenerationClient(ImageGenerationProvider): body.update(self.extra_body) url = f"{self.api_base}/images/generations" - client = self._client or httpx.AsyncClient(timeout=self.timeout) + client = self._client or httpx.AsyncClient(**self._http_client_kwargs()) try: return await self._generate_with_client( client, @@ -2047,8 +2072,8 @@ class ModelScopeImageGenerationClient(ImageGenerationProvider): f"{_MODELSCOPE_POLL_MAX_ATTEMPTS} polls" ) - @staticmethod async def _collect_images( + self, data: dict[str, Any], ) -> list[str]: images: list[str] = [] @@ -2057,7 +2082,9 @@ class ModelScopeImageGenerationClient(ImageGenerationProvider): if url.startswith("data:image/"): images.append(url) else: - images.append(await _download_image_data_url(url)) + images.append( + await _download_image_data_url(url, proxy=self.proxy) + ) return images diff --git a/tests/providers/test_image_generation.py b/tests/providers/test_image_generation.py index 8504d4d6..8181ab1d 100644 --- a/tests/providers/test_image_generation.py +++ b/tests/providers/test_image_generation.py @@ -103,19 +103,19 @@ class CodexStreamingCompleteThenErrorResponse(FakeResponse): @pytest.fixture(autouse=True) -def generated_image_downloads(monkeypatch) -> list[str]: +def generated_image_downloads(monkeypatch) -> list[tuple[str, str | None]]: """Keep provider response parsing tests independent from outbound HTTP.""" - urls: list[str] = [] + downloads: list[tuple[str, str | None]] = [] - async def download(url: str) -> str: - urls.append(url) + async def download(url: str, *, proxy: str | None = None) -> str: + downloads.append((url, proxy)) return PNG_DATA_URL monkeypatch.setattr( "nanobot.providers.image_generation._download_image_data_url", download, ) - return urls + return downloads @pytest.mark.asyncio @@ -294,19 +294,21 @@ 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( - generated_image_downloads: list[str], + generated_image_downloads: list[tuple[str, str | None]], ) -> None: fake = FakeClient(FakeResponse({"data": [{"url": "https://cdn.example/image.png"}]})) fake.get_response = FakeResponse({}, content=PNG_BYTES) + proxy = "http://127.0.0.1:23458" client = AIHubMixImageGenerationClient( api_key="sk-ahm-test", + proxy=proxy, client=fake, # type: ignore[arg-type] ) response = await client.generate(prompt="draw", model="gpt-image-2-free") assert response.images[0].startswith("data:image/png;base64,") - assert generated_image_downloads == ["https://cdn.example/image.png"] + assert generated_image_downloads == [("https://cdn.example/image.png", proxy)] @pytest.mark.asyncio @@ -836,18 +838,22 @@ async def test_openai_b64_json_response_uses_detected_mime() -> None: @pytest.mark.asyncio -async def test_openai_url_download_fallback(generated_image_downloads: list[str]) -> None: +async def test_openai_url_download_fallback( + generated_image_downloads: list[tuple[str, str | None]], +) -> None: fake = FakeClient(FakeResponse({"data": [{"url": "https://cdn.example/image.png"}]})) fake.get_response = FakeResponse({}, content=PNG_BYTES) + proxy = "http://127.0.0.1:23458" client = OpenAIImageGenerationClient( api_key="sk-openai-test", + proxy=proxy, client=fake, # type: ignore[arg-type] ) response = await client.generate(prompt="draw", model="dall-e-3") assert response.images[0].startswith("data:image/png;base64,") - assert generated_image_downloads == ["https://cdn.example/image.png"] + assert generated_image_downloads == [("https://cdn.example/image.png", proxy)] @pytest.mark.asyncio @@ -1211,14 +1217,16 @@ 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( - generated_image_downloads: list[str], + generated_image_downloads: list[tuple[str, str | None]], ) -> None: fake = FakeClient(FakeResponse({"data": [{"url": "https://images.example/cat.png"}]})) fake.get_response = FakeResponse({}, content=PNG_BYTES) + proxy = "http://127.0.0.1:23458" client = CustomImageGenerationClient( api_key="sk-custom-test", api_base="https://custom.example/v1", extra_body={"response_format": "url", "size": "2K"}, + proxy=proxy, client=fake, # type: ignore[arg-type] ) @@ -1229,7 +1237,7 @@ async def test_custom_generate_extra_body_can_override_defaults( ) assert response.images == [PNG_DATA_URL] - assert generated_image_downloads == ["https://images.example/cat.png"] + assert generated_image_downloads == [("https://images.example/cat.png", proxy)] body = fake.calls[0]["json"] assert body["response_format"] == "url" assert body["size"] == "2K" @@ -1636,19 +1644,21 @@ async def test_zhipu_image_generation_with_explicit_size() -> None: @pytest.mark.asyncio async def test_zhipu_image_generation_downloads_url_response( - generated_image_downloads: list[str], + generated_image_downloads: list[tuple[str, str | None]], ) -> None: fake = FakeClient(FakeResponse({"data": [{"url": "https://cdn.example/image.png"}]})) fake.get_response = FakeResponse({}, content=PNG_BYTES) + proxy = "http://127.0.0.1:23458" client = ZhipuImageGenerationClient( api_key="sk-zhipu-test", + proxy=proxy, client=fake, # type: ignore[arg-type] ) response = await client.generate(prompt="draw", model="glm-image") assert response.images[0].startswith("data:image/png;base64,") - assert generated_image_downloads == ["https://cdn.example/image.png"] + assert generated_image_downloads == [("https://cdn.example/image.png", proxy)] @pytest.mark.asyncio @@ -1728,7 +1738,9 @@ def _modelscope_fast_poll(monkeypatch) -> None: @pytest.mark.asyncio -async def test_modelscope_image_generation_submit_and_poll() -> None: +async def test_modelscope_image_generation_submit_and_poll( + generated_image_downloads: list[tuple[str, str | None]], +) -> None: submit = FakeResponse({"task_id": "abc123"}) poll_responses = [ FakeResponse({"task_status": "PENDING"}), @@ -1738,9 +1750,11 @@ async def test_modelscope_image_generation_submit_and_poll() -> None: }), ] fake = ModelScopeFakeClient(submit, poll_responses) + proxy = "http://127.0.0.1:23458" client = ModelScopeImageGenerationClient( api_key="ms-token", api_base="https://api-inference.modelscope.cn/v1", + proxy=proxy, client=fake, # type: ignore[arg-type] ) @@ -1750,6 +1764,7 @@ async def test_modelscope_image_generation_submit_and_poll() -> None: ) assert response.images[0].startswith("data:image/png;base64,") + assert generated_image_downloads == [("https://cdn.example/image.png", proxy)] # Verify POST request post_call = fake.calls[0] @@ -1919,3 +1934,18 @@ async def test_modelscope_image_generation_poll_timeout(monkeypatch) -> None: # Should have polled up to the (patched) attempt limit. assert len(fake.get_calls) == 3 + + + +def test_image_provider_http_client_kwargs_include_explicit_proxy() -> None: + proxy = "http://127.0.0.1:23458" + client = AIHubMixImageGenerationClient( + api_key="sk-ahm-test", + proxy=proxy, + ) + + assert client._http_client_kwargs() == { + "timeout": client.timeout, + "proxy": proxy, + "trust_env": False, + } diff --git a/tests/providers/test_image_generation_security.py b/tests/providers/test_image_generation_security.py index 44b586a7..48bc2e7f 100644 --- a/tests/providers/test_image_generation_security.py +++ b/tests/providers/test_image_generation_security.py @@ -33,8 +33,16 @@ def _resolve_public(host: str, port: int | None, *args, **kwargs): ["http://127.0.0.1/admin", "http://[::]/admin"], ids=["ipv4-loopback", "ipv6-unspecified"], ) +@pytest.mark.parametrize( + "proxy", + [None, "http://127.0.0.1:23458"], + ids=["direct", "explicit-proxy"], +) @pytest.mark.asyncio -async def test_generated_image_download_blocks_unsafe_target(url: str) -> None: +async def test_generated_image_download_blocks_unsafe_target( + url: str, + proxy: str | None, +) -> None: requested = False async def handler(request: httpx.Request) -> httpx.Response: @@ -45,6 +53,7 @@ async def test_generated_image_download_blocks_unsafe_target(url: str) -> None: with pytest.raises(ImageGenerationError, match="blocked unsafe generated image URL"): await _download_image_data_url( url, + proxy=proxy, transport=httpx.MockTransport(handler), ) @@ -116,3 +125,103 @@ async def test_generated_image_download_enforces_streaming_size_limit(monkeypatc "https://cdn.example/image.png", transport=httpx.MockTransport(handler), ) + + +class _StreamContext: + def __init__(self, response: httpx.Response) -> None: + self.response = response + + async def __aenter__(self) -> httpx.Response: + return self.response + + async def __aexit__(self, exc_type, exc, traceback) -> None: + await self.response.aclose() + + +@pytest.mark.asyncio +async def test_generated_image_download_uses_explicit_provider_proxy( + monkeypatch, +) -> None: + monkeypatch.setattr( + "nanobot.security.network.socket.getaddrinfo", + _resolve_public, + ) + captured: dict[str, object] = {} + + class FakeAsyncClient: + def __init__(self, **kwargs) -> None: + captured["kwargs"] = kwargs + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, traceback) -> None: + return None + + def stream(self, method: str, url: str) -> _StreamContext: + captured["request"] = (method, url) + request = httpx.Request(method, url) + return _StreamContext(httpx.Response(200, content=PNG_BYTES, request=request)) + + monkeypatch.setattr(image_generation.httpx, "AsyncClient", FakeAsyncClient) + proxy = "http://127.0.0.1:23458" + + result = await _download_image_data_url( + "https://cdn.example/image.png", + proxy=proxy, + ) + + assert result.startswith("data:image/png;base64,") + assert captured["request"] == ("GET", "https://cdn.example/image.png") + assert captured["kwargs"] == { + "follow_redirects": False, + "timeout": image_generation._DEFAULT_TIMEOUT_S, + "trust_env": False, + "proxy": proxy, + } + + +@pytest.mark.asyncio +async def test_proxied_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] = [] + + class FakeAsyncClient: + def __init__(self, **kwargs) -> None: + pass + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, traceback) -> None: + return None + + def stream(self, method: str, url: str) -> _StreamContext: + requested.append(url) + request = httpx.Request(method, url) + return _StreamContext( + httpx.Response( + 302, + headers={"location": "http://169.254.169.254/latest"}, + request=request, + ) + ) + + monkeypatch.setattr(image_generation.httpx, "AsyncClient", FakeAsyncClient) + + with pytest.raises(ImageGenerationError, match="blocked unsafe generated image URL"): + await _download_image_data_url( + "https://cdn.example/image.png", + proxy="http://127.0.0.1:23458", + ) + + assert requested == ["https://cdn.example/image.png"]