fix(security): harden gallery endpoint URL checks

Replace substring OpenAI endpoint detection with exact parsed-host matching.

Route gallery image endpoint construction through a constant path allowlist.

Remove client-visible exception and upstream response body leaks from gallery image flows while preserving diagnostics in server logs.

Add focused regression tests for OpenAI host matching, checked endpoint joining, harmonize SSRF hardening, and sanitized client errors.
This commit is contained in:
Alexandre Teixeira
2026-06-28 13:23:19 +01:00
parent a7fc1343a3
commit 0cf7eddde6
2 changed files with 454 additions and 88 deletions
+128 -88
View File
@@ -77,6 +77,39 @@ def _normalize_image_endpoint_base(url: str) -> str:
return base
def _is_openai_api_base(url: str) -> bool:
"""Return True only when url's hostname is exactly api.openai.com."""
from urllib.parse import urlsplit
try:
candidate = url if "://" in url else f"https://{url}"
return urlsplit(candidate).hostname == "api.openai.com"
except Exception:
return False
_GALLERY_ENDPOINT_PATHS = frozenset({
"/images/edits",
"/images/generations",
"/images/harmonize",
"/images/img2img",
"/images/inpaint",
"/images/upscale",
"/images/variations",
"/sdapi/v1/img2img",
})
def _join_checked_gallery_endpoint(base: str, path: str) -> str:
"""Append a known-constant gallery path suffix to a validated base URL.
Rejects paths not in the pre-approved list so arbitrary strings can never
be spliced into the URL passed to httpx.
"""
if path not in _GALLERY_ENDPOINT_PATHS:
raise ValueError(f"Unexpected gallery path: {path!r}")
return base + path
def _visible_image_endpoint_query(db, owner: str | None):
from src.auth_helpers import owner_filter
q = db.query(ModelEndpoint).filter(
@@ -255,9 +288,10 @@ def setup_gallery_routes() -> APIRouter:
pass
try:
db.commit()
except Exception as e:
except Exception:
db.rollback()
raise HTTPException(500, f"DB commit failed: {e}")
logger.exception("gallery_replace: DB commit failed")
raise HTTPException(500, "Image update failed")
return {"ok": True, "width": img.width, "height": img.height}
finally:
db.close()
@@ -385,8 +419,9 @@ def setup_gallery_routes() -> APIRouter:
return {"image": data.get("data", [{}])[0].get("b64_json", "")}
# Fallback: no upscale endpoint — return error
return {"error": f"Upscale endpoint not available ({resp.status_code})"}
except Exception as e:
return {"error": str(e)}
except Exception:
logger.exception("ai_upscale: request failed")
return {"error": "Upscale request failed"}
# ---- POST /api/gallery/style-transfer ----
@router.post("/api/gallery/style-transfer")
@@ -431,8 +466,9 @@ def setup_gallery_routes() -> APIRouter:
if img_data:
return {"image": img_data}
return {"error": f"Style transfer failed ({resp.status_code})"}
except Exception as e:
return {"error": str(e)}
except Exception:
logger.exception("style_transfer: request failed")
return {"error": "Style transfer failed"}
# ---- GET /api/gallery/tags ----
@router.get("/api/gallery/tags")
@@ -588,9 +624,9 @@ def setup_gallery_routes() -> APIRouter:
"tags": sorted(all_tags),
"models": all_models,
}
except Exception as e:
logger.error(f"Failed to fetch gallery library: {e}")
raise HTTPException(500, f"Failed to fetch gallery library: {e}")
except Exception:
logger.exception("Failed to fetch gallery library")
raise HTTPException(500, "Failed to fetch gallery library")
finally:
db.close()
@@ -766,9 +802,10 @@ def setup_gallery_routes() -> APIRouter:
return _image_to_dict(img)
except HTTPException:
raise
except Exception as e:
except Exception:
db.rollback()
raise HTTPException(500, str(e))
logger.exception("patch_gallery_image: update failed")
raise HTTPException(500, "Image update failed")
finally:
db.close()
@@ -845,9 +882,10 @@ def setup_gallery_routes() -> APIRouter:
cleared += 1
db.commit()
return {"ok": True, "cleared": cleared}
except Exception as e:
except Exception:
db.rollback()
raise HTTPException(500, str(e))
logger.exception("clear_gallery_user_tags: failed")
raise HTTPException(500, "Tag update failed")
finally:
db.close()
@@ -871,9 +909,10 @@ def setup_gallery_routes() -> APIRouter:
cleared += 1
db.commit()
return {"ok": True, "cleared": cleared}
except Exception as e:
except Exception:
db.rollback()
raise HTTPException(500, str(e))
logger.exception("clear_gallery_ai_tags: failed")
raise HTTPException(500, "Tag update failed")
finally:
db.close()
@@ -909,9 +948,10 @@ def setup_gallery_routes() -> APIRouter:
img.tags = ', '.join(cleaned)
db.commit()
return {"ok": True, "rows_touched": rows_touched, "tags_removed": tags_removed}
except Exception as e:
except Exception:
db.rollback()
raise HTTPException(500, str(e))
logger.exception("dedupe_gallery_tags: failed")
raise HTTPException(500, "Tag deduplication failed")
finally:
db.close()
@@ -1029,9 +1069,10 @@ def setup_gallery_routes() -> APIRouter:
return {"status": "deleted", "id": image_id}
except HTTPException:
raise
except Exception as e:
except Exception:
db.rollback()
raise HTTPException(500, str(e))
logger.exception("delete_gallery_image: failed")
raise HTTPException(500, "Image deletion failed")
finally:
db.close()
@@ -1044,21 +1085,22 @@ def setup_gallery_routes() -> APIRouter:
import httpx
user = require_privilege(request, "can_generate_images")
body = await request.json()
# Use endpoint from request body (editor dropdown) or fall back to DB lookup
base = (body.pop("_endpoint", "") or "").rstrip("/")
# Use endpoint from request body (editor dropdown) or fall back to DB lookup.
# Store as requested_base to avoid carrying user input into the outbound request.
requested_base = (body.pop("_endpoint", "") or "").rstrip("/")
# SSRF hardening: validate a client-supplied endpoint before any
# outbound request (mirrors routes/embedding_routes.py).
if base:
if requested_base:
from src.url_safety import check_outbound_url
ok, reason = check_outbound_url(
base,
requested_base,
block_private=os.getenv("IMAGE_BLOCK_PRIVATE_IPS", "false").lower() == "true",
)
if not ok:
raise HTTPException(400, f"Rejected endpoint URL: {reason}")
chosen_model = (body.pop("_model", "") or "").strip()
api_key = None
if not base:
if not requested_base:
db = SessionLocal()
try:
ep = _first_visible_image_endpoint(db, user)
@@ -1069,32 +1111,23 @@ def setup_gallery_routes() -> APIRouter:
finally:
db.close()
else:
# Pull api_key from the matching DB row so OpenAI auth works.
# Users may have stored base_url with/without /v1 suffix and with/without
# trailing slash, so compare normalized forms.
def _norm_url(u: str) -> str:
if not u:
return u
u = u.rstrip("/")
if u.endswith("/v1"):
u = u[:-3]
return u
_target = _norm_url(base)
# Resolve the client-supplied base to a registered visible endpoint.
# Admins are not exempted — gallery proxy routes must use a DB row
# so the outbound URL never depends directly on request-body input.
db = SessionLocal()
try:
ep = _visible_image_endpoint_for_base(db, _target, user)
if ep:
base = (ep.base_url or base).rstrip("/")
api_key = ep.api_key
elif user and not _current_user_is_admin(request, user):
ep = _visible_image_endpoint_for_base(db, requested_base, user)
if not ep:
raise HTTPException(403, "Choose a registered image endpoint")
base = ep.base_url.rstrip("/")
api_key = ep.api_key
finally:
db.close()
if not base.endswith("/v1"):
base += "/v1"
is_openai = "api.openai.com" in base
is_openai = _is_openai_api_base(base)
if is_openai:
# OpenAI path: /v1/images/edits with gpt-image-1.
@@ -1131,8 +1164,9 @@ def setup_gallery_routes() -> APIRouter:
mask_buf.seek(0)
except HTTPException:
raise
except Exception as e:
raise HTTPException(400, f"Failed to prepare OpenAI request: {e}")
except Exception:
logger.exception("inpaint_proxy: failed to prepare OpenAI request")
raise HTTPException(400, "Failed to prepare inpaint request")
width = int(body.get("width") or 1024)
height = int(body.get("height") or 1024)
@@ -1163,9 +1197,10 @@ def setup_gallery_routes() -> APIRouter:
headers = {"Authorization": f"Bearer {api_key}"}
try:
async with httpx.AsyncClient(timeout=120) as client:
r = await client.post(f"{base}/images/edits", headers=headers, data=data, files=files)
r = await client.post(_join_checked_gallery_endpoint(base, "/images/edits"), headers=headers, data=data, files=files)
if r.status_code != 200:
raise HTTPException(r.status_code, f"OpenAI edit failed: {r.text[:300]}")
logger.error("inpaint_proxy OpenAI edit: status %s", r.status_code)
raise HTTPException(r.status_code, "OpenAI edit failed")
result = r.json()
raw_b64 = None
if result.get("data"):
@@ -1212,16 +1247,18 @@ def setup_gallery_routes() -> APIRouter:
if chosen_model:
body["model"] = chosen_model
async with httpx.AsyncClient(timeout=120) as client:
r = await client.post(f"{base}/images/inpaint", json=body)
r = await client.post(_join_checked_gallery_endpoint(base, "/images/inpaint"), json=body)
if r.status_code != 200:
raise HTTPException(r.status_code, f"Inpaint failed: {r.text[:200]}")
logger.error("inpaint_proxy diffusion: status %s", r.status_code)
raise HTTPException(r.status_code, "Inpaint request failed")
return r.json()
except httpx.TimeoutException:
raise HTTPException(504, "Inpaint request timed out (120s)")
except HTTPException:
raise
except Exception as e:
raise HTTPException(502, f"Inpaint error: {str(e)}")
except Exception:
logger.exception("inpaint_proxy: request failed")
raise HTTPException(502, "Inpaint request failed")
# ---- POST /api/image/harmonize — proper img2img call ----
# Earlier version routed through inpaint with a full-white mask, but
@@ -1243,24 +1280,23 @@ def setup_gallery_routes() -> APIRouter:
if not image_b64:
raise HTTPException(400, "No image provided")
endpoint = (body.get("_endpoint") or "").rstrip("/")
requested_base = (body.get("_endpoint") or "").rstrip("/")
# SSRF hardening: a client-supplied endpoint is fetched server-side
# below, so validate it first (mirrors routes/embedding_routes.py).
# Local-first means loopback/LAN is allowed by default; the cloud
# metadata range and non-HTTP(S) schemes are always rejected.
if endpoint:
if requested_base:
from src.url_safety import check_outbound_url
ok, reason = check_outbound_url(
endpoint,
requested_base,
block_private=os.getenv("IMAGE_BLOCK_PRIVATE_IPS", "false").lower() == "true",
)
if not ok:
raise HTTPException(400, f"Rejected endpoint URL: {reason}")
model = (body.get("_model") or "").strip()
base = endpoint
api_key = None
if not base:
if not requested_base:
db = SessionLocal()
try:
ep = _first_visible_image_endpoint(db, user)
@@ -1271,14 +1307,16 @@ def setup_gallery_routes() -> APIRouter:
finally:
db.close()
else:
# Resolve the client-supplied base to a registered visible endpoint.
# Admins are not exempted — gallery proxy routes must use a DB row
# so the outbound URL never depends directly on request-body input.
db = SessionLocal()
try:
ep = _visible_image_endpoint_for_base(db, base, user)
if ep:
base = (ep.base_url or base).rstrip("/")
api_key = ep.api_key
elif user and not _current_user_is_admin(request, user):
ep = _visible_image_endpoint_for_base(db, requested_base, user)
if not ep:
raise HTTPException(403, "Choose a registered image endpoint")
base = ep.base_url.rstrip("/")
api_key = ep.api_key
finally:
db.close()
@@ -1313,7 +1351,7 @@ def setup_gallery_routes() -> APIRouter:
# source. Earlier hack (alpha-blend the regen back at `strength`)
# produced visibly broken results, so we refuse and tell the
# user to spin up a real diffusion endpoint instead.
if "api.openai.com" in base:
if _is_openai_api_base(base):
raise HTTPException(400,
"Harmonize needs a diffusion server that supports img2img "
"(SD WebUI / Forge / Comfy). OpenAI's API doesn't expose "
@@ -1378,14 +1416,16 @@ def setup_gallery_routes() -> APIRouter:
# 1024×1024 inference pass on slower setups.
async with httpx.AsyncClient(timeout=240) as client:
for path, kind, payload in candidates:
target = base_root + path if path.startswith("/sdapi") else base + path
_effective_base = base_root if path.startswith("/sdapi") else base
target = _join_checked_gallery_endpoint(_effective_base, path)
try:
r = await client.post(target, json=payload, headers=headers)
if r.status_code == 404:
last_err = f"{path}: 404"
continue # try next variant
if r.status_code != 200:
last_err = f"{path}: {r.status_code} {r.text[:120]}"
logger.warning("harmonize: %s returned %s", path, r.status_code)
last_err = f"{path}: {r.status_code}"
continue
data = r.json()
# Normalise return shape.
@@ -1394,8 +1434,8 @@ def setup_gallery_routes() -> APIRouter:
# surface it now instead of trying the other routes
# (otherwise the real error gets buried under 404s).
if data.get("error") and not data.get("image"):
raise HTTPException(502,
f"Diffusion server error at {path}: {data['error']}")
logger.warning("harmonize: server error at %s: %s", path, data.get("error"))
raise HTTPException(502, f"Diffusion server error at {path}")
if data.get("image"):
return {"image": data["image"]}
if data.get("images") and isinstance(data["images"], list):
@@ -1415,15 +1455,15 @@ def setup_gallery_routes() -> APIRouter:
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}")
except httpx.ConnectError:
logger.warning("harmonize: can't reach diffusion server at %s", base)
raise HTTPException(502, "Can't reach diffusion server")
except httpx.TimeoutException:
raise HTTPException(504, "Harmonize timed out (240s) — restart the diffusion server or lower Color match / disable Seam fix")
raise HTTPException(502,
f"None of the img2img routes worked on {base}. "
f"Last response: {last_err or 'unknown'}. "
"Your diffusion server needs to expose one of /v1/images/harmonize, "
"/v1/images/img2img, /v1/images/variations, or /sdapi/v1/img2img.")
"No supported img2img route responded. "
"Your diffusion server needs to expose one of: "
"/v1/images/harmonize, /v1/images/img2img, /v1/images/variations, /sdapi/v1/img2img.")
# ---- POST /api/image/sharpen ----
@router.post("/api/image/sharpen")
@@ -1467,8 +1507,8 @@ def setup_gallery_routes() -> APIRouter:
import base64, io
from PIL import Image
import numpy as np
except ImportError as e:
raise HTTPException(500, f"Server missing dependency: {e}")
except ImportError:
raise HTTPException(500, "Server missing a required dependency")
# Decode source image (RGB; Real-ESRGAN doesn't preserve alpha).
img_bytes = base64.b64decode(image_b64)
src = Image.open(io.BytesIO(img_bytes)).convert("RGB")
@@ -1495,9 +1535,9 @@ def setup_gallery_routes() -> APIRouter:
buf = io.BytesIO()
out_img.save(buf, format="PNG")
return {"image": base64.b64encode(buf.getvalue()).decode()}
except Exception as e:
logger.warning(f"Denoise failed: {e}")
return {"error": f"Denoise failed: {e}"}
except Exception:
logger.warning("Denoise failed", exc_info=True)
return {"error": "Denoise failed"}
# ---- POST /api/image/upscale-local ----
# Local Real-ESRGAN upscale (2× or 4×). Self-contained — no diffusion
@@ -1518,8 +1558,8 @@ def setup_gallery_routes() -> APIRouter:
import base64, io
from PIL import Image
import numpy as np
except ImportError as e:
raise HTTPException(500, f"Server missing dependency: {e}")
except ImportError:
raise HTTPException(500, "Server missing a required dependency")
img_bytes = base64.b64decode(image_b64)
src = Image.open(io.BytesIO(img_bytes)).convert("RGB")
try:
@@ -1543,9 +1583,9 @@ def setup_gallery_routes() -> APIRouter:
buf = io.BytesIO()
out_img.save(buf, format="PNG")
return {"image": base64.b64encode(buf.getvalue()).decode()}
except Exception as e:
logger.warning(f"Upscale failed: {e}")
return {"error": f"Upscale failed: {e}"}
except Exception:
logger.warning("AI upscale failed", exc_info=True)
return {"error": "AI upscale failed"}
# ---- POST /api/image/remove-bg ----
@router.post("/api/image/remove-bg")
@@ -1703,8 +1743,9 @@ def setup_gallery_routes() -> APIRouter:
buf = io.BytesIO()
enhanced.save(buf, format="PNG")
return {"image": base64.b64encode(buf.getvalue()).decode(), "method": "pil"}
except Exception as e:
raise HTTPException(500, f"Face enhancement failed: {str(e)}")
except Exception:
logger.exception("enhance_face: failed")
raise HTTPException(500, "Face enhancement failed")
# ---- Album management (path-param routes) ----
@@ -1899,9 +1940,8 @@ def setup_gallery_routes() -> APIRouter:
async with httpx.AsyncClient(timeout=60) as client:
resp = await client.post(chat_url, json=payload, headers=h)
if resp.status_code != 200:
body = resp.text[:500]
logger.error(f"Vision model {resp.status_code}: {body}")
return {"error": f"Vision model returned {resp.status_code}: {body[:200]}"}
logger.error("ai_tag vision model: status %s: %s", resp.status_code, resp.text[:500])
return {"error": "Vision model request failed"}
data = resp.json()
# Anthropic returns content[0].text, OpenAI returns choices[0].message.content
if provider == "anthropic":
@@ -1917,9 +1957,9 @@ def setup_gallery_routes() -> APIRouter:
return {"ok": True, "ai_tags": tag_str}
except HTTPException:
raise
except Exception as e:
logger.error(f"AI tagging failed: {e}")
return {"error": str(e)}
except Exception:
logger.exception("AI tagging failed")
return {"error": "Auto-tagging failed"}
finally:
db.close()
+326
View File
@@ -0,0 +1,326 @@
"""Focused security tests for gallery endpoint URL hardening.
Covers:
- _is_openai_api_base: exact hostname matching (no substring bypass)
- _join_checked_gallery_endpoint: allowlist-only path construction
- No bare str(e) / f"...{e}" in gallery exception handlers
- harmonize validates _endpoint via check_outbound_url
- Target URL construction only appends constant paths to the validated base
"""
import ast
import re
from pathlib import Path
SRC = Path(__file__).resolve().parent.parent / "routes" / "gallery" / "gallery_routes.py"
import routes.gallery_routes as gallery_routes
# ---------------------------------------------------------------------------
# _is_openai_api_base — exact hostname, no substring tricks
# ---------------------------------------------------------------------------
def test_is_openai_api_base_accepts_exact_host():
f = gallery_routes._is_openai_api_base
assert f("https://api.openai.com") is True
assert f("https://api.openai.com/v1") is True
assert f("https://api.openai.com/") is True
assert f("api.openai.com") is True
def test_is_openai_api_base_rejects_path_embed():
# attacker hides api.openai.com in the path, not the hostname
f = gallery_routes._is_openai_api_base
assert f("https://evil.test/api.openai.com/v1") is False
def test_is_openai_api_base_rejects_subdomain_suffix():
# hostname ends with .openai.com but isn't exactly api.openai.com
f = gallery_routes._is_openai_api_base
assert f("https://api.openai.com.evil.test/v1") is False
assert f("https://evil-api.openai.com/v1") is False
assert f("https://notapi.openai.com/v1") is False
def test_is_openai_api_base_rejects_malformed():
f = gallery_routes._is_openai_api_base
assert f("") is False
assert f("not a url at all !!!") is False
# ---------------------------------------------------------------------------
# Source-level: gallery no longer uses substring "api.openai.com" in base
# ---------------------------------------------------------------------------
def test_gallery_does_not_use_openai_substring_check():
src = SRC.read_text()
assert '"api.openai.com" in base' not in src, (
"Substring OpenAI check still present — use _is_openai_api_base instead"
)
assert "'api.openai.com' in base" not in src, (
"Substring OpenAI check still present — use _is_openai_api_base instead"
)
# ---------------------------------------------------------------------------
# _join_checked_gallery_endpoint — allowlist enforcement
# ---------------------------------------------------------------------------
def test_join_checked_accepts_known_paths():
j = gallery_routes._join_checked_gallery_endpoint
assert j("http://localhost:7860/v1", "/images/img2img") == "http://localhost:7860/v1/images/img2img"
assert j("http://localhost:7860", "/sdapi/v1/img2img") == "http://localhost:7860/sdapi/v1/img2img"
assert j("https://api.openai.com/v1", "/images/edits") == "https://api.openai.com/v1/images/edits"
def test_join_checked_rejects_unknown_path():
import pytest
j = gallery_routes._join_checked_gallery_endpoint
with pytest.raises(ValueError):
j("http://localhost/v1", "/arbitrary/user/path")
with pytest.raises(ValueError):
j("http://localhost/v1", "")
with pytest.raises(ValueError):
j("http://localhost/v1", "https://evil.test/steal")
# ---------------------------------------------------------------------------
# Source-level: no raw str(e) / f"...{e}" returned to API clients
# ---------------------------------------------------------------------------
def test_no_raw_exception_string_in_client_responses():
src = SRC.read_text()
# Patterns that indicate exception internals flowing into client-visible values.
# We allow them only in logger calls (checked separately below).
bad_patterns = [
r'return \{"error": str\(e\)\}',
r'return \{"error": f"[^"]*\{e\}[^"]*"\}',
r'HTTPException\(\d+, str\(e\)\)',
r'HTTPException\(\d+, f"[^"]*\{e\}[^"]*"\)',
]
for pattern in bad_patterns:
matches = re.findall(pattern, src)
assert not matches, (
f"Pattern {pattern!r} matched — raw exception string returned to client: {matches}"
)
# ---------------------------------------------------------------------------
# harmonize: validates _endpoint via check_outbound_url before outbound request
# ---------------------------------------------------------------------------
def _function_source(src_text: str, func_name: str) -> str:
tree = ast.parse(src_text)
for node in ast.walk(tree):
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == func_name:
return ast.get_source_segment(src_text, node) or ""
raise AssertionError(f"{func_name} not found in {SRC}")
def test_harmonize_validates_endpoint_before_fetch():
src = SRC.read_text()
body = _function_source(src, "harmonize_image")
assert "check_outbound_url" in body, (
"harmonize_image must validate _endpoint via check_outbound_url before outbound requests"
)
# ---------------------------------------------------------------------------
# harmonize: target URL only appends constant allowed paths
# ---------------------------------------------------------------------------
def test_harmonize_uses_join_checked_for_target_construction():
src = SRC.read_text()
body = _function_source(src, "harmonize_image")
assert "_join_checked_gallery_endpoint" in body, (
"harmonize_image must use _join_checked_gallery_endpoint to build target URLs"
)
# Raw concatenation patterns that bypass the allowlist must not appear in harmonize
assert "base_root + path" not in body, (
"harmonize_image must not concatenate base_root + path directly"
)
assert "base + path" not in body, (
"harmonize_image must not concatenate base + path directly"
)
def test_gallery_endpoint_paths_allowlist_covers_all_harmonize_candidates():
# Every path in the candidates list must be in the pre-approved allowlist.
src = SRC.read_text()
body = _function_source(src, "harmonize_image")
# Extract string literals that look like route paths from candidates
candidate_paths = re.findall(r'"/(?:images|sdapi)/[^"]*"', body)
allowed = gallery_routes._GALLERY_ENDPOINT_PATHS
for p in candidate_paths:
p = p.strip('"')
assert p in allowed, (
f"Path {p!r} used in harmonize candidates but not in _GALLERY_ENDPOINT_PATHS allowlist"
)
# ---------------------------------------------------------------------------
# _is_openai_api_base — userinfo bypass
# ---------------------------------------------------------------------------
def test_is_openai_api_base_rejects_userinfo_bypass():
# userinfo trick: user = api.openai.com, host = evil.test
f = gallery_routes._is_openai_api_base
assert f("https://api.openai.com@evil.test/v1") is False
# ---------------------------------------------------------------------------
# Source-level: no client-visible error leaks upstream body fragments
# ---------------------------------------------------------------------------
def _extract_httpexception_call(src: str, pos: int) -> str:
"""Paren-match from the opening '(' of an HTTPException call."""
start = src.index("(", pos)
depth = 0
for k, ch in enumerate(src[start:]):
if ch == "(":
depth += 1
elif ch == ")":
depth -= 1
if depth == 0:
return src[start : start + k + 1]
return src[start:]
def test_no_upstream_data_in_client_responses():
"""No raise HTTPException or return {"error": ...} may expose upstream body data."""
src = SRC.read_text()
forbidden = [
"r.text",
"body[:",
'data["error"]',
"data['error']",
"last_err",
"{base}",
]
for m in re.finditer(r"\braise\s+HTTPException\s*\(", src):
line_start = src.rfind("\n", 0, m.start()) + 1
if "logger." in src[line_start : m.start()]:
continue
call_text = _extract_httpexception_call(src, m.start())
for frag in forbidden:
assert frag not in call_text, (
f"HTTPException raise at byte {m.start()} exposes {frag!r} to client:\n{call_text[:300]}"
)
for m in re.finditer(r'return\s*\{"error":', src):
line_end = src.find("\n", m.start())
line = src[m.start() : line_end if line_end != -1 else len(src)]
for frag in forbidden:
assert frag not in line, (
f"Error return at byte {m.start()} exposes {frag!r} to client:\n{line}"
)
# ---------------------------------------------------------------------------
# inpaint_proxy: endpoint construction via _join_checked_gallery_endpoint
# ---------------------------------------------------------------------------
def test_inpaint_uses_join_checked_endpoint():
src = SRC.read_text()
body = _function_source(src, "inpaint_proxy")
assert 'f"{base}/images/edits"' not in body, (
"inpaint_proxy must not build /images/edits via raw f-string"
)
assert 'f"{base}/images/inpaint"' not in body, (
"inpaint_proxy must not build /images/inpaint via raw f-string"
)
assert '_join_checked_gallery_endpoint(base, "/images/edits")' in body, (
"inpaint_proxy must use _join_checked_gallery_endpoint for /images/edits"
)
assert '_join_checked_gallery_endpoint(base, "/images/inpaint")' in body, (
"inpaint_proxy must use _join_checked_gallery_endpoint for /images/inpaint"
)
# ---------------------------------------------------------------------------
# harmonize final 502: no base URL or last_err in client message
# ---------------------------------------------------------------------------
def test_harmonize_final_502_omits_base_and_last_err():
src = SRC.read_text()
body = _function_source(src, "harmonize_image")
# Collect all HTTPException raises in harmonize and check the last one (final 502)
raises = list(re.finditer(r"\braise\s+HTTPException\s*\(", body))
assert raises, "harmonize_image must contain at least one raise HTTPException"
last_call = _extract_httpexception_call(body, raises[-1].start())
for forbidden in ("last_err", "{base}", "r.text"):
assert forbidden not in last_call, (
f"harmonize final raise exposes {forbidden!r} to client:\n{last_call}"
)
# ---------------------------------------------------------------------------
# inpaint/harmonize: _endpoint must resolve via DB; no raw admin bypass
# ---------------------------------------------------------------------------
def test_inpaint_endpoint_resolved_via_db_not_raw_input():
"""inpaint_proxy must not use the raw request-body value as the outbound base.
The user-supplied value is stored as requested_base; outbound base comes from DB."""
src = SRC.read_text()
body = _function_source(src, "inpaint_proxy")
# requested_base holds the user input; base is only set from ep.base_url
assert "requested_base" in body, (
"inpaint_proxy must use 'requested_base' for the user-supplied value"
)
# The admin bypass (not _current_user_is_admin) must not appear in inpaint
assert "_current_user_is_admin" not in body, (
"inpaint_proxy must not have an admin bypass for raw endpoint resolution"
)
# If no matching endpoint is found, a 403 must be raised unconditionally
assert 'raise HTTPException(403, "Choose a registered image endpoint")' in body, (
"inpaint_proxy must raise 403 when _endpoint doesn't match a registered endpoint"
)
def test_inpaint_outbound_base_not_from_request_body():
"""Confirm _join_checked_gallery_endpoint is never called with the raw
request-body variable (requested_base) — only with the DB-derived base."""
src = SRC.read_text()
body = _function_source(src, "inpaint_proxy")
assert "_join_checked_gallery_endpoint(requested_base," not in body, (
"inpaint_proxy must not pass requested_base to _join_checked_gallery_endpoint"
)
def test_harmonize_endpoint_resolved_via_db_not_raw_input():
"""harmonize_image must not use the raw request-body value as the outbound base."""
src = SRC.read_text()
body = _function_source(src, "harmonize_image")
assert "requested_base" in body, (
"harmonize_image must use 'requested_base' for the user-supplied value"
)
assert "_current_user_is_admin" not in body, (
"harmonize_image must not have an admin bypass for raw endpoint resolution"
)
assert 'raise HTTPException(403, "Choose a registered image endpoint")' in body, (
"harmonize_image must raise 403 when _endpoint doesn't match a registered endpoint"
)
def test_harmonize_outbound_base_not_from_request_body():
"""Confirm _join_checked_gallery_endpoint is never called with requested_base."""
src = SRC.read_text()
body = _function_source(src, "harmonize_image")
assert "_join_checked_gallery_endpoint(requested_base," not in body, (
"harmonize_image must not pass requested_base to _join_checked_gallery_endpoint"
)
def test_inpaint_and_harmonize_no_base_equals_endpoint():
"""Neither function should assign `base = endpoint` or `base = requested_base`
— the outbound base must come exclusively from DB (ep.base_url)."""
src = SRC.read_text()
for func_name in ("inpaint_proxy", "harmonize_image"):
body = _function_source(src, func_name)
assert "base = endpoint" not in body, (
f"{func_name}: 'base = endpoint' carries request-body input into outbound request"
)
assert "base = requested_base" not in body, (
f"{func_name}: 'base = requested_base' carries request-body input into outbound request"
)