fix(external): guard AlibabaCloud image downloads

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

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(