fix(ai): validate generated image result URLs (#4289)

tests/test_ai_image_url_safety.py

@@ -0,0 +1,104 @@
+from src import ai_interaction
+
+
+class _GenerationResponse:
+    status_code = 200
+    text = ""
+
+    def __init__(self, image_url):
+        self._image_url = image_url
+
+    def json(self):
+        return {"data": [{"url": self._image_url}]}
+
+
+class _DownloadResponse:
+    status_code = 503
+    content = b""
+
+
+def _patch_generation(monkeypatch, image_url):
+    async def _post(self, url, json, headers):
+        return _GenerationResponse(image_url)
+
+    class _AsyncClient:
+        def __init__(self, *args, **kwargs):
+            pass
+
+        async def __aenter__(self):
+            return self
+
+        async def __aexit__(self, *exc):
+            return False
+
+        post = _post
+
+    import httpx
+    import src.settings as settings
+
+    monkeypatch.setattr(settings, "load_settings", lambda: {})
+    monkeypatch.setattr(httpx, "AsyncClient", _AsyncClient)
+    monkeypatch.setattr(
+        ai_interaction,
+        "_resolve_model",
+        lambda model_spec, owner=None: (
+            "https://api.openai.example/v1/chat/completions",
+            "dall-e-3",
+            {"Authorization": "Bearer test"},
+        ),
+    )
+
+
+async def test_generate_image_validates_provider_url_before_download(monkeypatch):
+    import httpx
+    import src.url_safety as url_safety
+
+    provider_url = "https://images.example.com/generated.png?sig=abc"
+    events = []
+    _patch_generation(monkeypatch, provider_url)
+
+    def _check_outbound_url(url, *, block_private=False):
+        events.append(("check", url, block_private))
+        return True, "ok"
+
+    def _get(url, *, timeout):
+        events.append(("get", url, timeout))
+        return _DownloadResponse()
+
+    monkeypatch.setattr(url_safety, "check_outbound_url", _check_outbound_url)
+    monkeypatch.setattr(httpx, "get", _get)
+
+    result = await ai_interaction.do_generate_image("draw a chair\ndall-e-3")
+
+    assert result["image_url"] == provider_url
+    assert events == [
+        ("check", provider_url, False),
+        ("get", provider_url, 60),
+    ]
+
+
+async def test_generate_image_rejects_unsafe_provider_url_without_download(monkeypatch):
+    import httpx
+    import src.url_safety as url_safety
+
+    unsafe_url = "http://169.254.169.254/latest/meta-data"
+    events = []
+    _patch_generation(monkeypatch, unsafe_url)
+
+    def _check_outbound_url(url, *, block_private=False):
+        events.append(("check", url, block_private))
+        return False, "link-local address blocked (SSRF metadata risk): 169.254.169.254"
+
+    def _get(url, *, timeout):
+        raise AssertionError("unsafe provider image URL must not be downloaded")
+
+    monkeypatch.setattr(url_safety, "check_outbound_url", _check_outbound_url)
+    monkeypatch.setattr(httpx, "get", _get)
+
+    result = await ai_interaction.do_generate_image("draw a chair\ndall-e-3")
+
+    assert result["error"] == (
+        "Image API returned unsafe image URL: "
+        "link-local address blocked (SSRF metadata risk): 169.254.169.254"
+    )
+    assert events == [("check", unsafe_url, False)]