feat(products): image-fetch phase — SSRF guard, bounds, renditions, resume (SD-0002 §6.5.4, §6.9, INV-18)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
@@ -8,6 +8,7 @@ from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from .imagefetch import recover_incomplete_runs, run_image_phase
|
||||
from .errors import (
|
||||
DraftExpired,
|
||||
DraftNotFound,
|
||||
@@ -40,4 +41,5 @@ __all__ = [
|
||||
"MAX_DATA_ROWS", "MAX_FILE_BYTES", "SAMPLE_CSV_PATH",
|
||||
"import_validate", "get_draft", "get_draft_records", "discard_draft",
|
||||
"confirm_draft", "list_runs", "get_run", "summary", "export_catalog",
|
||||
"run_image_phase", "recover_incomplete_runs",
|
||||
]
|
||||
|
||||
@@ -0,0 +1,141 @@
|
||||
"""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
|
||||
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()
|
||||
|
||||
|
||||
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
|
||||
@@ -0,0 +1,137 @@
|
||||
"""image-fetch phase: SSRF guard, fetch+process+store, run completion, resume (§6.5.4/§6.9)."""
|
||||
import io
|
||||
import threading
|
||||
from http.server import BaseHTTPRequestHandler, HTTPServer
|
||||
|
||||
import psycopg
|
||||
import pytest
|
||||
from PIL import Image
|
||||
|
||||
from app.domains.products import imagefetch, repo
|
||||
from app.platform import db, objectstore
|
||||
|
||||
|
||||
def _png(w, h):
|
||||
buf = io.BytesIO(); Image.new("RGB", (w, h), (1, 2, 3)).save(buf, "PNG"); return buf.getvalue()
|
||||
|
||||
|
||||
class _Host(BaseHTTPRequestHandler):
|
||||
GOOD = _png(800, 800)
|
||||
TINY = _png(64, 64)
|
||||
def log_message(self, *a): pass
|
||||
def do_GET(self):
|
||||
if self.path.startswith("/good.png"):
|
||||
self.send_response(200); self.send_header("Content-Type", "image/png")
|
||||
self.end_headers(); self.wfile.write(self.GOOD)
|
||||
elif self.path.startswith("/tiny.png"):
|
||||
self.send_response(200); self.send_header("Content-Type", "image/png")
|
||||
self.end_headers(); self.wfile.write(self.TINY)
|
||||
else:
|
||||
self.send_response(404); self.end_headers()
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def host():
|
||||
srv = HTTPServer(("127.0.0.1", 0), _Host)
|
||||
t = threading.Thread(target=srv.serve_forever, daemon=True); t.start()
|
||||
yield f"http://127.0.0.1:{srv.server_address[1]}"
|
||||
srv.shutdown()
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def pool(fresh_db_url):
|
||||
with psycopg.connect(fresh_db_url) as conn:
|
||||
db.migrate(conn); conn.commit()
|
||||
p = db.open_pool(fresh_db_url, max_size=6)
|
||||
yield p
|
||||
p.close()
|
||||
|
||||
|
||||
def _seed(pool):
|
||||
"""account+storefront+run; returns (storefront_id, account_id, run_id)."""
|
||||
with pool.connection() as conn:
|
||||
acct = conn.execute("INSERT INTO account (email) VALUES ('m@example.com') RETURNING id").fetchone()[0]
|
||||
sf = conn.execute("INSERT INTO storefront (name) VALUES ('Shop') RETURNING id").fetchone()[0]
|
||||
conn.execute("INSERT INTO storefront_membership (account_id, storefront_id) VALUES (%s,%s)", (acct, sf))
|
||||
rid = conn.execute(
|
||||
"INSERT INTO import_run (storefront_id, account_id, file_name, dialect,"
|
||||
" products_added, products_updated, rows_errored, status)"
|
||||
" VALUES (%s,%s,'c.csv','canonical',0,0,0,'fetching_images') RETURNING id", (sf, acct)).fetchone()[0]
|
||||
conn.commit()
|
||||
return sf, acct, rid
|
||||
|
||||
|
||||
def _product(pool, sf, handle):
|
||||
with pool.connection() as conn:
|
||||
pid = conn.execute("INSERT INTO product (storefront_id, handle, title) VALUES (%s,%s,%s) RETURNING id",
|
||||
(sf, handle, handle.title())).fetchone()[0]
|
||||
conn.commit()
|
||||
return pid
|
||||
|
||||
|
||||
def _image(pool, pid, url, rid, position=1):
|
||||
with pool.connection() as conn:
|
||||
iid = conn.execute("INSERT INTO product_image (product_id, source_url, position, status, import_run_id)"
|
||||
" VALUES (%s,%s,%s,'pending',%s) RETURNING id", (pid, url, position, rid)).fetchone()[0]
|
||||
conn.commit()
|
||||
return iid
|
||||
|
||||
|
||||
def _img_row(pool, iid):
|
||||
with pool.connection() as conn:
|
||||
r = conn.execute("SELECT status, key_original, key_thumb, key_card, key_detail FROM product_image WHERE id=%s",
|
||||
(iid,)).fetchone()
|
||||
return {"status": r[0], "key_original": r[1], "key_thumb": r[2], "key_card": r[3], "key_detail": r[4]}
|
||||
|
||||
|
||||
def _run_status(pool, rid):
|
||||
with pool.connection() as conn:
|
||||
return conn.execute("SELECT status FROM import_run WHERE id=%s", (rid,)).fetchone()[0]
|
||||
|
||||
|
||||
def test_ssrf_guard_rejects_private_when_disallowed():
|
||||
with pytest.raises(imagefetch.FetchBlocked):
|
||||
imagefetch.fetch_bytes("http://127.0.0.1:1/x.png", allow_private=False)
|
||||
with pytest.raises(imagefetch.FetchBlocked):
|
||||
imagefetch.fetch_bytes("http://169.254.169.254/latest/meta-data", allow_private=False)
|
||||
with pytest.raises(imagefetch.FetchBlocked):
|
||||
imagefetch.fetch_bytes("ftp://example.com/x", allow_private=False)
|
||||
|
||||
|
||||
def test_ssrf_guard_allows_private_when_enabled(host):
|
||||
data, ctype = imagefetch.fetch_bytes(f"{host}/good.png", allow_private=True)
|
||||
assert ctype.startswith("image/") and len(data) > 0
|
||||
|
||||
|
||||
def test_run_phase_classifies_each_image(pool, host, tmp_path):
|
||||
store = objectstore.LocalObjectStore(str(tmp_path))
|
||||
sf, _acct, rid = _seed(pool)
|
||||
p = _product(pool, sf, "lamp")
|
||||
ok = _image(pool, p, f"{host}/good.png", rid, position=1)
|
||||
_image(pool, p, f"{host}/tiny.png", rid, position=2)
|
||||
_image(pool, p, f"{host}/missing.png", rid, position=3)
|
||||
imagefetch.run_image_phase(pool, store, rid, allow_private=True)
|
||||
with pool.connection() as conn:
|
||||
counts = repo.run_image_counts(conn, rid)
|
||||
assert counts["fetched"] == 1 and counts["rejected"] == 1 and counts["failed"] == 1
|
||||
row = _img_row(pool, ok)
|
||||
assert row["status"] == "fetched"
|
||||
for k in ("key_original", "key_thumb", "key_card", "key_detail"):
|
||||
assert row[k] and store.get(row[k])
|
||||
assert _run_status(pool, rid) == "complete_with_problems"
|
||||
|
||||
|
||||
def test_resume_after_kill_completes_remaining(pool, host, tmp_path):
|
||||
store = objectstore.LocalObjectStore(str(tmp_path))
|
||||
sf, _acct, rid = _seed(pool)
|
||||
p = _product(pool, sf, "lamp")
|
||||
a = _image(pool, p, f"{host}/good.png", rid, position=1)
|
||||
with pool.connection() as conn: # simulate crash: one already fetched, run still fetching_images
|
||||
repo.mark_image_fetched(conn, a, {"original": "o", "thumb": "t", "card": "c", "detail": "d"})
|
||||
conn.commit()
|
||||
_image(pool, p, f"{host}/good.png?2", rid, position=2) # still pending
|
||||
resumed = imagefetch.recover_incomplete_runs(pool, store, allow_private=True)
|
||||
assert rid in resumed
|
||||
with pool.connection() as conn:
|
||||
assert repo.run_image_counts(conn, rid)["pending"] == 0
|
||||
assert _run_status(pool, rid) == "complete"
|
||||
Reference in New Issue
Block a user