mirror of
https://github.com/pewdiepie-archdaemon/odysseus.git
synced 2026-06-16 09:45:24 -04:00
nopoz
6824fbb729
Validate image URLs returned by upstream diffusion/OpenAI responses before server-side fetches to prevent SSRF through result image retrieval.
70 lines
2.3 KiB
Python
70 lines
2.3 KiB
Python
"""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()
|