From b7ed652bf4f38fa623313a99e93dfb8fa44bcacb Mon Sep 17 00:00:00 2001 From: wunianze666-netizen Date: Sat, 22 Aug 2026 14:18:14 +0800 Subject: [PATCH] fix(external): guard AlibabaCloud image downloads --- .../providers/alibabacloud.py | 53 +++++++++++-------- .../test_alibabacloud_provider.py | 42 +++++++++------ 2 files changed, 57 insertions(+), 38 deletions(-) diff --git a/invokeai/app/services/external_generation/providers/alibabacloud.py b/invokeai/app/services/external_generation/providers/alibabacloud.py index 6a1d01baefa..91f2bce7394 100644 --- a/invokeai/app/services/external_generation/providers/alibabacloud.py +++ b/invokeai/app/services/external_generation/providers/alibabacloud.py @@ -15,6 +15,7 @@ ExternalGenerationResult, ) from invokeai.app.services.external_generation.image_utils import decode_image_base64, encode_image_base64 +from invokeai.app.util.ssrf import UnsafeDownloadURLException, build_guarded_session # Models that support the synchronous multimodal-generation endpoint with messages format _SYNC_MODELS = { @@ -290,33 +291,39 @@ def _parse_async_response( def _download_image(self, url: str) -> PILImageType: """Download an image from a URL and return it as a PIL Image, with a size cap.""" try: - response = requests.get(url, timeout=_DOWNLOAD_TIMEOUT, stream=True) - except requests.RequestException as exc: - raise ExternalProviderRequestError(f"Failed to download image from DashScope: {exc}") from exc + with build_guarded_session() as session: + response = session.get(url, timeout=_DOWNLOAD_TIMEOUT, stream=True) - with response: - if not response.ok: - raise ExternalProviderRequestError( - f"Failed to download image from DashScope (status {response.status_code})" - ) - - content_length = response.headers.get("Content-Length") - if content_length is not None: - try: - if int(content_length) > _DOWNLOAD_MAX_BYTES: + with response: + if not response.ok: raise ExternalProviderRequestError( - f"DashScope image exceeds {_DOWNLOAD_MAX_BYTES} byte cap (Content-Length={content_length})" + f"Failed to download image from DashScope (status {response.status_code})" ) - except ValueError: - pass - buffer = bytearray() - for chunk in response.iter_content(chunk_size=64 * 1024): - if not chunk: - continue - buffer.extend(chunk) - if len(buffer) > _DOWNLOAD_MAX_BYTES: - raise ExternalProviderRequestError(f"DashScope image exceeds {_DOWNLOAD_MAX_BYTES} byte cap") + content_length = response.headers.get("Content-Length") + if content_length is not None: + try: + if int(content_length) > _DOWNLOAD_MAX_BYTES: + raise ExternalProviderRequestError( + f"DashScope image exceeds {_DOWNLOAD_MAX_BYTES} byte cap " + f"(Content-Length={content_length})" + ) + except ValueError: + pass + + buffer = bytearray() + for chunk in response.iter_content(chunk_size=64 * 1024): + if not chunk: + continue + buffer.extend(chunk) + if len(buffer) > _DOWNLOAD_MAX_BYTES: + raise ExternalProviderRequestError( + f"DashScope image exceeds {_DOWNLOAD_MAX_BYTES} byte cap" + ) + except UnsafeDownloadURLException as exc: + raise ExternalProviderRequestError("DashScope returned an unsafe image URL") from exc + except requests.RequestException as exc: + raise ExternalProviderRequestError(f"Failed to download image from DashScope: {exc}") from exc return Image.open(io.BytesIO(bytes(buffer))).convert("RGB") diff --git a/tests/app/services/external_generation/test_alibabacloud_provider.py b/tests/app/services/external_generation/test_alibabacloud_provider.py index d8c9015daf6..eb982aa5054 100644 --- a/tests/app/services/external_generation/test_alibabacloud_provider.py +++ b/tests/app/services/external_generation/test_alibabacloud_provider.py @@ -14,6 +14,7 @@ from invokeai.app.services.external_generation.image_utils import encode_image_base64 from invokeai.app.services.external_generation.providers import alibabacloud as alibabacloud_module from invokeai.app.services.external_generation.providers.alibabacloud import AlibabaCloudProvider +from invokeai.app.util.ssrf import UnsafeDownloadURLException from invokeai.backend.model_manager.configs.external_api import ExternalApiModelConfig, ExternalModelCapabilities @@ -137,7 +138,7 @@ def fake_post(url: str, headers: dict, json: dict, timeout: int) -> DummyRespons }, ) - def fake_get(url: str, timeout: int, stream: bool = False) -> DummyResponse: + def fake_get(_session: Any, url: str, timeout: int, stream: bool = False) -> DummyResponse: assert url == image_url return DummyResponse( ok=True, @@ -146,7 +147,7 @@ def fake_get(url: str, timeout: int, stream: bool = False) -> DummyResponse: ) monkeypatch.setattr("requests.post", fake_post) - monkeypatch.setattr("requests.get", fake_get) + monkeypatch.setattr("requests.Session.get", fake_get) result = provider.generate(request) @@ -198,11 +199,11 @@ def fake_post(url: str, headers: dict, json: dict, timeout: int) -> DummyRespons }, ) - def fake_get(url: str, timeout: int, stream: bool = False) -> DummyResponse: + def fake_get(_session: Any, url: str, timeout: int, stream: bool = False) -> DummyResponse: return DummyResponse(ok=True, content=image_bytes, headers={"Content-Length": str(len(image_bytes))}) monkeypatch.setattr("requests.post", fake_post) - monkeypatch.setattr("requests.get", fake_get) + monkeypatch.setattr("requests.Session.get", fake_get) monkeypatch.setattr("time.sleep", lambda _s: None) result = provider.generate(request) @@ -217,10 +218,10 @@ def test_async_parser_does_not_double_count(monkeypatch: pytest.MonkeyPatch) -> image_bytes = _png_bytes(_make_image("magenta")) image_url = "https://example.invalid/x.png" - def fake_get(url: str, timeout: int, stream: bool = False) -> DummyResponse: + def fake_get(_session: Any, url: str, timeout: int, stream: bool = False) -> DummyResponse: return DummyResponse(ok=True, content=image_bytes, headers={"Content-Length": str(len(image_bytes))}) - monkeypatch.setattr("requests.get", fake_get) + monkeypatch.setattr("requests.Session.get", fake_get) output: dict[str, Any] = { "results": [ @@ -250,19 +251,35 @@ def test_download_image_size_cap(monkeypatch: pytest.MonkeyPatch) -> None: provider = _provider() too_big = alibabacloud_module._DOWNLOAD_MAX_BYTES + 1 - def fake_get(url: str, timeout: int, stream: bool = False) -> DummyResponse: + def fake_get(_session: Any, url: str, timeout: int, stream: bool = False) -> DummyResponse: return DummyResponse( ok=True, content=b"\x00" * 16, # body itself is small; we trip the Content-Length check first headers={"Content-Length": str(too_big)}, ) - monkeypatch.setattr("requests.get", fake_get) + monkeypatch.setattr("requests.Session.get", fake_get) with pytest.raises(ExternalProviderRequestError, match="exceeds"): provider._download_image("https://example.invalid/big.png") +def test_download_image_rejects_unsafe_provider_url(monkeypatch: pytest.MonkeyPatch) -> None: + provider = _provider() + + def fail_unguarded_get(*_args: Any, **_kwargs: Any) -> DummyResponse: + pytest.fail("provider response URLs must not use the unguarded requests.get path") + + def reject_unsafe_url(*_args: Any, **_kwargs: Any) -> DummyResponse: + raise UnsafeDownloadURLException("non-public address") + + monkeypatch.setattr("requests.get", fail_unguarded_get) + monkeypatch.setattr("requests.Session.get", reject_unsafe_url) + + with pytest.raises(ExternalProviderRequestError, match="unsafe image URL"): + provider._download_image("http://127.0.0.1/internal.png") + + def test_poll_task_first_call_no_initial_sleep(monkeypatch: pytest.MonkeyPatch) -> None: """First poll must not be preceded by a sleep — fast tasks should not pay the poll interval.""" provider = _provider() @@ -289,13 +306,8 @@ def fake_get(url: str, headers: dict, timeout: int) -> DummyResponse: def fake_download_get(url: str, timeout: int, stream: bool = False) -> DummyResponse: return DummyResponse(ok=True, content=image_bytes, headers={"Content-Length": str(len(image_bytes))}) - # Single requests.get is shared by polling (with headers kwarg) and download (no kwarg). - def dispatch_get(*args: Any, **kwargs: Any) -> DummyResponse: - if "headers" in kwargs and "task" in args[0]: - return fake_get(*args, **kwargs) - return fake_download_get(*args, **kwargs) - - monkeypatch.setattr("requests.get", dispatch_get) + monkeypatch.setattr("requests.get", fake_get) + monkeypatch.setattr("requests.Session.get", lambda _session, *args, **kwargs: fake_download_get(*args, **kwargs)) monkeypatch.setattr("time.sleep", fake_sleep) result = provider._poll_task(