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>
2026-06-12 00:26:26 -07:00
parent 13e74f4c61
commit 23267c4d4c
3 changed files with 280 additions and 0 deletions
@@ -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"