fix: close library upload flow and harden image storage

pigway_frame_flow/renderers/images.py

@@ -1,7 +1,11 @@
 import base64
 import ipaddress
 import json
+import logging
+import os
 import socket
+import threading
+from contextlib import contextmanager
 from io import BytesIO
 from urllib.parse import urlsplit
 from uuid import uuid4
@@ -27,7 +31,7 @@ def convert(
     quality=93,
 ):
     try:
-        with Image.open(BytesIO(data)) as source:
+        with Image.open(BytesIO(data), formats=("JPEG", "PNG", "WEBP")) as source:
             if source.format not in {"JPEG", "PNG", "WEBP"}:
                 raise ValueError()
             if source.width * source.height > MAX_IMAGE_PIXELS:
@@ -82,17 +86,99 @@ def convert(
 class Assets:
     def __init__(self, db, settings):
         self.db, self.settings = db, settings
+        self.file_lock = threading.RLock()
+        db.execute(
+            "CREATE TABLE IF NOT EXISTS asset_file_journal(id TEXT PRIMARY KEY,kind TEXT NOT NULL)"
+        )
+        db.execute(
+            "CREATE TABLE IF NOT EXISTS asset_file_moves(id TEXT PRIMARY KEY,source_kind TEXT NOT NULL,target_kind TEXT NOT NULL)"
+        )
+        self.recover_files()
 
-    def put(self, data, name="image.jpg", kind="images"):
+    @staticmethod
+    def sync_directory(directory):
+        fd = os.open(directory, os.O_RDONLY)
+        try:
+            os.fsync(fd)
+        finally:
+            os.close(fd)
+
+    def recover_files(self):
+        with self.file_lock:
+            self._recover_files()
+
+    def _recover_files(self):
+        for row in self.db.all("SELECT * FROM asset_file_journal"):
+            path = self.settings.data_dir / row["kind"] / (row["id"] + ".jpg")
+            try:
+                if not self.db.one("SELECT id FROM assets WHERE id=?", (row["id"],)):
+                    path.unlink(missing_ok=True)
+                path.with_suffix(".part").unlink(missing_ok=True)
+                if path.parent.exists():
+                    self.sync_directory(path.parent)
+                self.db.execute(
+                    "DELETE FROM asset_file_journal WHERE id=?", (row["id"],)
+                )
+            except OSError:
+                logging.getLogger(__name__).warning("Asset file cleanup pending")
+
+        for row in self.db.all("SELECT * FROM asset_file_moves"):
+            asset = self.db.one("SELECT kind FROM assets WHERE id=?", (row["id"],))
+            obsolete = (
+                row["source_kind"]
+                if asset and asset["kind"] == row["target_kind"]
+                else row["target_kind"]
+            )
+            path = self.settings.data_dir / obsolete / (row["id"] + ".jpg")
+            try:
+                path.unlink(missing_ok=True)
+                self.sync_directory(path.parent)
+                self.db.execute("DELETE FROM asset_file_moves WHERE id=?", (row["id"],))
+            except OSError:
+                logging.getLogger(__name__).warning("Asset promotion cleanup pending")
+
+    @contextmanager
+    def creating(self, data, name="image.jpg", kind="images"):
+        with self.file_lock:
+            with self._creating(data, name, kind) as asset:
+                yield asset
+
+    @contextmanager
+    def _creating(self, data, name="image.jpg", kind="images"):
         from pathlib import Path
 
+        if kind not in {"images", "generated", "cache", "source-images", "renditions"}:
+            raise ValueError("Invalid asset kind")
         name = Path(name).stem[:110] + ".jpg"
         asset_id = uuid4().hex
-        (self.settings.data_dir / kind / (asset_id + ".jpg")).write_bytes(data)
-        self.db.execute(
-            "INSERT INTO assets VALUES(?,?,?,?)", (asset_id, name, kind, now_iso())
-        )
-        return self.info(asset_id)
+        path = self.settings.data_dir / kind / (asset_id + ".jpg")
+        path.parent.mkdir(exist_ok=True)
+        self.db.execute("INSERT INTO asset_file_journal VALUES(?,?)", (asset_id, kind))
+        temporary = path.with_suffix(".part")
+        try:
+            with temporary.open("xb") as handle:
+                handle.write(data)
+                handle.flush()
+                os.fsync(handle.fileno())
+            temporary.replace(path)
+            self.sync_directory(path.parent)
+            with self.db.transaction():
+                self.db.execute(
+                    "INSERT INTO assets VALUES(?,?,?,?)",
+                    (asset_id, name, kind, now_iso()),
+                )
+                yield self.info(asset_id)
+                self.db.execute(
+                    "DELETE FROM asset_file_journal WHERE id=?", (asset_id,)
+                )
+        except BaseException:
+            # The journal remains durable if filesystem cleanup also fails.
+            self.recover_files()
+            raise
+
+    def put(self, data, name="image.jpg", kind="images"):
+        with self.creating(data, name, kind) as asset:
+            return asset
 
     def record_origin(self, asset_id, operation):
         job = self.db.one(
@@ -123,20 +209,36 @@ class Assets:
         )
 
     def promote(self, asset_id):
-        row = self.db.one("SELECT * FROM assets WHERE id=?", (asset_id,))
-        if row and row["kind"] == "cache":
-            source = self.db.one(
-                "SELECT payload FROM asset_sources WHERE asset_id=?", (asset_id,)
-            )
-            kind = (
-                "generated"
-                if source and json.loads(source["payload"]).get("type") == "template"
-                else "images"
-            )
-            target = self.settings.data_dir / kind / (asset_id + ".jpg")
-            self.path(asset_id).replace(target)
-            self.db.execute("UPDATE assets SET kind=? WHERE id=?", (kind, asset_id))
-        return self.info(asset_id)
+        with self.file_lock:
+            row = self.db.one("SELECT * FROM assets WHERE id=?", (asset_id,))
+            if row and row["kind"] == "cache":
+                source = self.db.one(
+                    "SELECT payload FROM asset_sources WHERE asset_id=?", (asset_id,)
+                )
+                kind = (
+                    "generated"
+                    if source
+                    and json.loads(source["payload"]).get("type") == "template"
+                    else "images"
+                )
+                target = self.settings.data_dir / kind / (asset_id + ".jpg")
+                self.recover_files()
+                self.db.execute(
+                    "INSERT INTO asset_file_moves VALUES(?,?,?)",
+                    (asset_id, "cache", kind),
+                )
+                try:
+                    # Keep the source valid until the database points to the durable target.
+                    os.link(self.path(asset_id), target)
+                    self.sync_directory(target.parent)
+                    self.db.execute(
+                        "UPDATE assets SET kind=? WHERE id=?", (kind, asset_id)
+                    )
+                except BaseException:
+                    self.recover_files()
+                    raise
+                self.recover_files()
+            return self.info(asset_id)
 
     def info(self, asset_id):
         row = self.db.one("SELECT * FROM assets WHERE id=?", (asset_id,))