fix(database): validate restored pair ownership and provenance (#193)

database/models.py

@@ -4,6 +4,8 @@ Database schema and models for KikoTextEncode prompt storage.
 
 import contextlib
 import datetime
+import json
+import math
 import shutil
 import sqlite3
 import os
@@ -30,6 +32,41 @@ except ImportError:
 ISO_TIMESTAMP_FORMAT = "%Y-%m-%dT%H:%M:%f+00:00"
 
 
+MAX_PAIR_PROVENANCE_LENGTH = 4_000_000
+MAX_PAIR_PROVENANCE_DEPTH = 64
+
+
+def pair_provenance_error(encoded: object) -> Optional[str]:
+    """Validate the same bounded JSON object for live writes and restores."""
+    if not isinstance(encoded, str):
+        return "Pair provenance must be JSON text"
+    if len(encoded) > MAX_PAIR_PROVENANCE_LENGTH:
+        return "Pair provenance exceeds the size limit"
+    try:
+        value = json.loads(encoded)
+    except (ValueError, TypeError, RecursionError):
+        return "Pair provenance is not readable JSON"
+    if not isinstance(value, dict):
+        return "Pair provenance must be a JSON object"
+    pending = [(iter((value,)), 0)]
+    while pending:
+        children, depth = pending[-1]
+        try:
+            item = next(children)
+        except StopIteration:
+            pending.pop()
+            continue
+        if depth > MAX_PAIR_PROVENANCE_DEPTH:
+            return "Pair provenance exceeds the nesting limit"
+        if isinstance(item, float) and not math.isfinite(item):
+            return "Pair provenance contains nonfinite numbers"
+        if isinstance(item, dict):
+            pending.append((iter(item.values()), depth + 1))
+        elif isinstance(item, list):
+            pending.append((iter(item), depth + 1))
+    return None
+
+
 def utc_now_iso() -> str:
     """Current UTC time as ISO text with microseconds.
 
@@ -442,6 +479,17 @@ class PromptModel:
                             null_unique = True
                 if not null_unique:
                     return "prompt_pairs is missing empty-negative pair uniqueness"
+        if conn.execute(
+            "SELECT 1 FROM image_prompt_pairs ip "
+            "JOIN prompt_pairs pp ON pp.id=ip.pair_id "
+            "JOIN generated_images gi ON gi.id=ip.image_id "
+            "WHERE pp.positive_prompt_id IS NOT gi.prompt_id LIMIT 1"
+        ).fetchone():
+            return "Image pair association has a different positive prompt owner"
+        for row in conn.execute("SELECT provenance FROM image_prompt_pairs"):
+            error = pair_provenance_error(row[0])
+            if error:
+                return error
         return None
 
     def _create_indexes(self, conn: sqlite3.Connection) -> None:

database/pairs.py

@@ -5,7 +5,7 @@ import sqlite3
 from contextlib import contextmanager
 from typing import Any, Dict, Iterator, List, Optional
 
-from .models import utc_now_iso
+from .models import pair_provenance_error, utc_now_iso
 
 try:
     from ..utils.hashing import generate_prompt_hash
@@ -339,8 +339,11 @@ class PairOperations:
             encoded = json.dumps(
                 provenance, sort_keys=True, separators=(",", ":"), allow_nan=False
             )
-        except (TypeError, ValueError) as error:
+        except (TypeError, ValueError, RecursionError) as error:
             raise ValueError("provenance must contain JSON values") from error
+        error = pair_provenance_error(encoded)
+        if error:
+            raise ValueError(error)
         pair = conn.execute(
             "SELECT positive_prompt_id FROM prompt_pairs WHERE id=?", (pair_id,)
         ).fetchone()