mirror of
https://github.com/pewdiepie-archdaemon/odysseus.git
synced 2026-06-15 17:25:26 -04:00
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.
This commit is contained in:
@@ -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}")
|
||||
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user