147 lines
6.2 KiB
Python
147 lines
6.2 KiB
Python
"""The post-commit image-fetch phase (SD-0002 §6.5.4, §6.9).
|
|
|
|
In-process, bounded-concurrency, resumable. For a run's pending images: fetch
|
|
each URL behind an SSRF guard with INV-18 bounds, run platform/images, store
|
|
renditions via the objectstore, and move the image pending -> fetched |
|
|
rejected_* | failed. Per-image idempotent (claim guard), so the phase can die
|
|
and resume. Lives in the domains layer (imports repo + platform); main.py
|
|
schedules it post-confirm and runs recovery at startup.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import ipaddress
|
|
import socket
|
|
import time
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from urllib.parse import urlsplit
|
|
|
|
import httpx
|
|
|
|
from app.platform import images as images_mod
|
|
from app.platform import telemetry
|
|
|
|
from . import repo
|
|
|
|
MAX_FETCH_BYTES = 20 * 1024 * 1024 # INV-18
|
|
FETCH_TIMEOUT_S = 30 # INV-18
|
|
FETCH_CONCURRENCY = 4 # §6.6
|
|
_CONTENT_PREFIX = "image/"
|
|
|
|
|
|
class FetchBlocked(Exception):
|
|
"""SSRF guard or bounds refusal — maps to image status 'failed' with reason."""
|
|
|
|
|
|
def _guard_host(url: str, allow_private: bool) -> None:
|
|
parts = urlsplit(url)
|
|
if parts.scheme not in {"http", "https"}:
|
|
raise FetchBlocked(f"unsupported scheme: {parts.scheme!r}")
|
|
host = parts.hostname
|
|
if not host:
|
|
raise FetchBlocked("no host")
|
|
if allow_private:
|
|
return
|
|
try:
|
|
infos = socket.getaddrinfo(host, parts.port or (443 if parts.scheme == "https" else 80))
|
|
except socket.gaierror as exc:
|
|
raise FetchBlocked(f"dns failure: {exc}") from exc
|
|
for info in infos:
|
|
ip = ipaddress.ip_address(info[4][0])
|
|
if (ip.is_private or ip.is_loopback or ip.is_link_local
|
|
or ip.is_reserved or ip.is_multicast or ip.is_unspecified):
|
|
raise FetchBlocked(f"non-public address: {ip}")
|
|
|
|
|
|
def fetch_bytes(url: str, allow_private: bool) -> tuple[bytes, str]:
|
|
"""SSRF-guarded, bounded GET. Returns (bytes, content_type). Raises FetchBlocked."""
|
|
_guard_host(url, allow_private)
|
|
with httpx.Client(follow_redirects=False, timeout=FETCH_TIMEOUT_S) as client:
|
|
with client.stream("GET", url) as resp:
|
|
if resp.status_code != 200:
|
|
raise FetchBlocked(f"http {resp.status_code}")
|
|
ctype = resp.headers.get("content-type", "").split(";")[0].strip().lower()
|
|
if not ctype.startswith(_CONTENT_PREFIX):
|
|
raise FetchBlocked(f"not an image content-type: {ctype!r}")
|
|
chunks: list[bytes] = []
|
|
total = 0
|
|
for chunk in resp.iter_bytes():
|
|
total += len(chunk)
|
|
if total > MAX_FETCH_BYTES:
|
|
raise FetchBlocked("oversize")
|
|
chunks.append(chunk)
|
|
return b"".join(chunks), ctype
|
|
|
|
|
|
def _key(storefront_id: int, image_id: int, name: str) -> str:
|
|
return f"storefronts/{storefront_id}/product-images/{image_id}/{name}"
|
|
|
|
|
|
def _process_one(pool, store, image: dict, allow_private: bool) -> None:
|
|
with pool.connection() as conn:
|
|
claimed = repo.claim_image_for_fetch(conn, image["id"])
|
|
conn.commit()
|
|
if not claimed:
|
|
return
|
|
image_id, sf = image["id"], image["storefront_id"]
|
|
try:
|
|
data, ctype = fetch_bytes(image["source_url"], allow_private)
|
|
except (FetchBlocked, httpx.HTTPError) as exc:
|
|
with pool.connection() as conn:
|
|
repo.mark_image_failed(conn, image_id, str(exc)[:200]); conn.commit()
|
|
return
|
|
try:
|
|
result = images_mod.process(data)
|
|
if isinstance(result, images_mod.Rejected):
|
|
reason = ("below the resolution bar" if result.reason == "rejected_low_res"
|
|
else "not a supported image")
|
|
with pool.connection() as conn:
|
|
repo.mark_image_rejected(conn, image_id, result.reason, reason); conn.commit()
|
|
return
|
|
keys = {"original": _key(sf, image_id, "original"),
|
|
"thumb": _key(sf, image_id, "thumb.webp"),
|
|
"card": _key(sf, image_id, "card.webp"),
|
|
"detail": _key(sf, image_id, "detail.webp")}
|
|
store.put(keys["original"], data, ctype)
|
|
store.put(keys["thumb"], result.renditions["thumb"], "image/webp")
|
|
store.put(keys["card"], result.renditions["card"], "image/webp")
|
|
store.put(keys["detail"], result.renditions["detail"], "image/webp")
|
|
with pool.connection() as conn:
|
|
repo.mark_image_fetched(conn, image_id, keys); conn.commit()
|
|
except Exception as exc: # noqa: BLE001 — one bad image must never wedge the run (review fix)
|
|
with pool.connection() as conn:
|
|
repo.mark_image_failed(conn, image_id, f"processing error: {type(exc).__name__}")
|
|
conn.commit()
|
|
|
|
|
|
def run_image_phase(pool, store, run_id: int, allow_private: bool) -> None:
|
|
"""Fetch all pending images of one run, then move the run to a terminal status."""
|
|
started = time.monotonic()
|
|
with pool.connection() as conn:
|
|
pending = repo.pending_images_for_run(conn, run_id)
|
|
if pending:
|
|
with ThreadPoolExecutor(max_workers=FETCH_CONCURRENCY) as pool_x:
|
|
list(pool_x.map(lambda img: _process_one(pool, store, img, allow_private), pending))
|
|
with pool.connection() as conn:
|
|
counts = repo.run_image_counts(conn, run_id)
|
|
status = ("complete_with_problems"
|
|
if counts["rejected"] + counts["failed"] > 0 else "complete")
|
|
repo.set_run_status(conn, run_id, status)
|
|
conn.commit()
|
|
telemetry.emit("image_phase_completed", run_id=run_id, fetched=counts["fetched"],
|
|
rejected=counts["rejected"], failed=counts["failed"],
|
|
duration_ms=int((time.monotonic() - started) * 1000))
|
|
|
|
|
|
def recover_incomplete_runs(pool, store, allow_private: bool) -> list[int]:
|
|
"""Startup scan (§6.9): resume any run stuck in fetching_images. Returns the ids."""
|
|
with pool.connection() as conn:
|
|
runs = repo.incomplete_runs(conn)
|
|
resumed: list[int] = []
|
|
for run in runs:
|
|
with pool.connection() as conn:
|
|
pending = len(repo.pending_images_for_run(conn, run["id"]))
|
|
telemetry.emit("image_phase_recovered", run_id=run["id"], pending_resumed=pending)
|
|
run_image_phase(pool, store, run["id"], allow_private)
|
|
resumed.append(run["id"])
|
|
return resumed
|