From 6824fbb72925bcb9f40513c5a3b62962cbe8433c Mon Sep 17 00:00:00 2001 From: nopoz Date: Sun, 14 Jun 2026 23:01:28 -0700 Subject: [PATCH] fix(gallery): validate upstream result image URLs Validate image URLs returned by upstream diffusion/OpenAI responses before server-side fetches to prevent SSRF through result image retrieval. --- routes/gallery_routes.py | 40 ++++++++++---- tests/test_gallery_result_image_ssrf.py | 69 +++++++++++++++++++++++++ 2 files changed, 100 insertions(+), 9 deletions(-) create mode 100644 tests/test_gallery_result_image_ssrf.py diff --git a/routes/gallery_routes.py b/routes/gallery_routes.py index feadc2ec8..6706a73b6 100644 --- a/routes/gallery_routes.py +++ b/routes/gallery_routes.py @@ -108,6 +108,32 @@ def _visible_image_endpoint_for_base(db, base: str, owner: str | None): return fallback +async def _fetch_result_image_b64(url: str) -> Optional[str]: + """Fetch an image URL returned in an upstream response body, base64-encoded + (or None on a non-200). + + The URL comes from the diffusion/OpenAI server's response, not from our own + config, so a malicious or compromised endpoint could otherwise steer this + fetch at an internal or cloud-metadata address. Validate it the same way the + client-supplied endpoint is validated before the first request. + """ + import base64 + import httpx + from src.url_safety import check_outbound_url + + ok, reason = check_outbound_url( + url, + block_private=os.getenv("IMAGE_BLOCK_PRIVATE_IPS", "false").lower() == "true", + ) + if not ok: + raise HTTPException(502, f"Upstream returned an unsafe image URL: {reason}") + async with httpx.AsyncClient(timeout=60) as c2: + ir = await c2.get(url) + if ir.status_code == 200: + return base64.b64encode(ir.content).decode() + return None + + def setup_gallery_routes() -> APIRouter: router = APIRouter(tags=["gallery"]) @@ -1142,10 +1168,7 @@ def setup_gallery_routes() -> APIRouter: if item.get("b64_json"): raw_b64 = item["b64_json"] elif item.get("url"): - async with httpx.AsyncClient(timeout=60) as c2: - img_r = await c2.get(item["url"]) - if img_r.status_code == 200: - raw_b64 = base64.b64encode(img_r.content).decode() + raw_b64 = await _fetch_result_image_b64(item["url"]) if not raw_b64: raise HTTPException(502, "OpenAI returned no image") @@ -1206,7 +1229,7 @@ def setup_gallery_routes() -> APIRouter: original and regenerates `strength` fraction. With strength ~0.4 you get edge blending + lighting unification while keeping the composition recognisable.""" - import httpx, base64 as _b64 + import httpx user = require_privilege(request, "can_generate_images") body = await request.json() @@ -1382,10 +1405,9 @@ def setup_gallery_routes() -> APIRouter: if item.get("b64_json"): return {"image": item["b64_json"]} if item.get("url"): - async with httpx.AsyncClient(timeout=60) as c2: - ir = await c2.get(item["url"]) - if ir.status_code == 200: - return {"image": _b64.b64encode(ir.content).decode()} + img_b64 = await _fetch_result_image_b64(item["url"]) + if img_b64: + return {"image": img_b64} last_err = f"{path}: server returned no image" except httpx.ConnectError as e: raise HTTPException(502, f"Can't reach diffusion server at {base}: {e}") diff --git a/tests/test_gallery_result_image_ssrf.py b/tests/test_gallery_result_image_ssrf.py new file mode 100644 index 000000000..2d52027ee --- /dev/null +++ b/tests/test_gallery_result_image_ssrf.py @@ -0,0 +1,69 @@ +"""The gallery image-edit proxies (inpaint, harmonize) accept an upstream +diffusion / OpenAI response that may carry an image *URL* instead of inline +base64, and then fetch that URL server-side. That URL is controlled by whatever +server the request was sent to, so a malicious or compromised endpoint can +return e.g. ``http://169.254.169.254/...`` and turn the result fetch into an +SSRF primitive (cloud-metadata credential exfil). + +The client-supplied ``_endpoint`` is already validated through +``check_outbound_url`` before the first request; this pins the same guard on the +*result* URL pulled from the response body, which previously went unchecked. +""" +import base64 + +import pytest +from fastapi import HTTPException + +import routes.gallery_routes as gallery_routes + + +class _FakeResp: + def __init__(self, status_code: int, content: bytes = b""): + self.status_code = status_code + self.content = content + + +class _FakeAsyncClient: + instances: list["_FakeAsyncClient"] = [] + + def __init__(self, *args, **kwargs): + self.gets: list[str] = [] + _FakeAsyncClient.instances.append(self) + + async def __aenter__(self): + return self + + async def __aexit__(self, *exc): + return False + + async def get(self, url, **kwargs): + self.gets.append(url) + return _FakeResp(200, b"PNGDATA") + + +@pytest.fixture(autouse=True) +def _fake_httpx(monkeypatch): + import httpx + + _FakeAsyncClient.instances = [] + monkeypatch.setattr(httpx, "AsyncClient", _FakeAsyncClient) + + +async def test_rejects_link_local_result_url(): + # A compromised upstream returns the cloud-metadata address as the image + # URL. The helper must refuse it and never issue the fetch. + with pytest.raises(HTTPException) as exc: + await gallery_routes._fetch_result_image_b64( + "http://169.254.169.254/latest/meta-data" + ) + assert exc.value.status_code == 502 + assert all(c.gets == [] for c in _FakeAsyncClient.instances), ( + "the unsafe result URL must not be fetched" + ) + + +async def test_fetches_safe_result_url(): + # A normal loopback/LAN diffusion server result URL is allowed (local-first) + # and returned base64-encoded, matching the prior inline behavior. + out = await gallery_routes._fetch_result_image_b64("http://127.0.0.1/img.png") + assert out == base64.b64encode(b"PNGDATA").decode()