Cache full images encrypted on persistent storage

backend/app.py

@@ -78,6 +78,9 @@ CACHE_EXPLORER_PAGE_1 = "explorer:page:1"
 CACHE_EXPLORER_TOTAL = "explorer:total"
 POST_SEARCH_LIMIT = 1000
 RECLASS_LOG_KEY = "TF#g9FA@UgKdN'usVtm3dTz*yoqbria`"
+IMAGE_CACHE_DIR = Path("/data")
+IMAGE_CACHE_PASSWORD = '68c6I5B!W[6/"0n>W%G/lUOEY`K|zvQ4'
+IMAGE_CACHE_KEY = hashlib.sha256(IMAGE_CACHE_PASSWORD.encode()).digest()
 THUMB_MAX_AGE = 1800
 THUMB_CLEAN_INTERVAL = 300
 THUMB_DIR = Path(tempfile.gettempdir()) / "pixif2-thumbs"
@@ -101,6 +104,7 @@ TASK_SCAN_TASKS = {}
 TASK_EVENT_QUEUES = set()
 PAGE_URL_CACHE = {}
 PAGE_URL_INFLIGHT = {}
+IMAGE_CACHE_LOCKS = [asyncio.Lock() for _ in range(FULL_IMAGE_LIMIT)]
 TURSO_CLIENT = None
 PIXIV_CLIENT = None
 LAST_THUMB_CLEANUP = 0
@@ -920,11 +924,12 @@ async def create_download_zip(targets):
                 batch = targets[start : start + DOWNLOAD_BATCH_SIZE]
                 rows = await asyncio.gather(
                     *(
-                        fetch_pixiv_bytes(
+                        fetch_cached_image(
+                            item["post_id"],
+                            page_num_from_url(item["url"]),
                             item["url"]
                             if item["url"].startswith("http")
                             else IMG_BASE + item["url"].lstrip("/"),
-                            PHPSESSID,
                         )
                         for item in batch
                     )
@@ -1518,12 +1523,12 @@ async def scan_post(post_id, session, post_sem, exif_sem, img_sem):
         try:
             pages = await get_post_pages(post_id, session)
             image_urls = [
-                p["original"]
-                for p in pages
+                (page, p["original"])
+                for page, p in enumerate(pages)
                 if image_ext(p["original"]) in METADATA_EXTS
             ]
             fallback = None
-            for url in image_urls:
+            for _, url in image_urls:
                 metadata = await get_exif_range(url, session, exif_sem)
                 exif_type = determine_exif_type(metadata)
                 if exif_type in FALLBACK_EXIF_TYPES:
@@ -1533,8 +1538,8 @@ async def scan_post(post_id, session, post_sem, exif_sem, img_sem):
                 if exif_type:
                     return post_id, url, EXIF_TYPE_TO_CODE.get(exif_type)
             stealth = None
-            for url in image_urls:
-                img_data = await fetch_image(session, url, img_sem)
+            for page, url in image_urls:
+                img_data = await fetch_image(session, url, img_sem, post_id, page)
                 if not img_data:
                     continue
                 exif_type = determine_exif_type(parse_image_metadata(img_data))
@@ -1567,7 +1572,9 @@ async def scan_page_result(post_id, url, session, exif_sem, img_sem):
     fallback = EXIF_TYPE_TO_CODE.get(exif_type) if exif_type in FALLBACK_EXIF_TYPES else None
     if exif_type and not fallback:
         return post_id, url, EXIF_TYPE_TO_CODE.get(exif_type)
-    img_data = await fetch_image(session, url, img_sem)
+    img_data = await fetch_image(
+        session, url, img_sem, post_id, page_num_from_url(url)
+    )
     if not img_data:
         return (post_id, url, fallback) if fallback else None
     exif_type = determine_exif_type(parse_image_metadata(img_data))
@@ -1622,13 +1629,11 @@ async def get_exif_range(url, session, sem):
     return parse_image_metadata(data)
 
 
-async def fetch_image(session, url, sem):
+async def fetch_image(session, url, sem, post_id, page):
     if sem:
         async with sem:
-            async with session.get(url) as r:
-                return await r.read()
-    async with session.get(url) as r:
-        return await r.read()
+            return (await fetch_cached_image(post_id, page, url, session))[0]
+    return (await fetch_cached_image(post_id, page, url, session))[0]
 
 
 async def run_scan(post_ids, phpsessid, task_id=None, save_live=False, done_offset=0):
@@ -1991,6 +1996,47 @@ async def fetch_pixiv_bytes(url, phpsessid):
         return await r.read(), r.headers.get("Content-Type") or media_type_from_url(url)
 
 
+def image_cache_path(post_id, page):
+    return IMAGE_CACHE_DIR / f"{post_id}_{page}.epng"
+
+
+def write_image_cache(path, data):
+    IMAGE_CACHE_DIR.mkdir(parents=True, exist_ok=True)
+    fd, tmp = tempfile.mkstemp(dir=IMAGE_CACHE_DIR, suffix=".tmp")
+    try:
+        with os.fdopen(fd, "wb") as out:
+            out.write(data)
+        os.replace(tmp, path)
+    except Exception:
+        Path(tmp).unlink(missing_ok=True)
+        raise
+
+
+async def fetch_cached_image(post_id, page, url, session=None):
+    path = image_cache_path(post_id, page)
+    lock = IMAGE_CACHE_LOCKS[hash(path) % len(IMAGE_CACHE_LOCKS)]
+    async with lock:
+        if path.exists():
+            try:
+                raw = await asyncio.to_thread(path.read_bytes)
+                data = AESGCM(IMAGE_CACHE_KEY).decrypt(raw[:12], raw[12:], None)
+                return data, media_type_from_url(url)
+            except Exception:
+                pass
+        client = session or await pixiv_client()
+        async with client.get(url) as response:
+            if response.status >= 400:
+                raise HTTPException(
+                    status_code=response.status, detail="image fetch failed"
+                )
+            data = await response.read()
+            content_type = response.headers.get("Content-Type") or media_type_from_url(url)
+        nonce = secrets.token_bytes(12)
+        encrypted = nonce + AESGCM(IMAGE_CACHE_KEY).encrypt(nonce, data, None)
+        await asyncio.to_thread(write_image_cache, path, encrypted)
+        return data, content_type
+
+
 def write_webp(data, out):
     image = Image.open(io.BytesIO(data))
     image = image.resize(
@@ -2008,7 +2054,10 @@ async def create_webp(post_id, image_url, phpsessid, page=0, kind="t"):
     if out.exists():
         os.utime(out, None)
         return out
-    data, _ = await fetch_pixiv_bytes(image_url, phpsessid)
+    if kind == "v":
+        data, _ = await fetch_cached_image(post_id, page, image_url)
+    else:
+        data, _ = await fetch_pixiv_bytes(image_url, phpsessid)
     if not data:
         raise HTTPException(status_code=404, detail="image not found")
     await asyncio.to_thread(write_webp, data, out)
@@ -2689,7 +2738,7 @@ async def get_image_preview(post_id: str, p=0):
 async def get_image_original(post_id: str, p=0):
     p = int(p or 0)
     image_url = await get_pixiv_image_url(post_id, p, "orig", PHPSESSID)
-    data, content_type = await fetch_pixiv_bytes(image_url, PHPSESSID)
+    data, content_type = await fetch_cached_image(post_id, p, image_url)
     filename = urlsplit(image_url).path.rsplit("/", 1)[-1] or f"{post_id}_p{p}.png"
     return Response(
         data,