From 752dff909886fa3f581943e6b8ac7f773c021c9d Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Sat, 17 Oct 2026 06:49:45 -0700 Subject: [PATCH] fix(database): validate restored pair ownership and provenance (#193) --- database/models.py | 48 ++++++++++++ database/pairs.py | 7 +- tests/test_pair_restore_semantics.py | 113 +++++++++++++++++++++++++++ 3 files changed, 166 insertions(+), 2 deletions(-) create mode 100644 tests/test_pair_restore_semantics.py diff --git a/database/models.py b/database/models.py index f0a45fd..37a6dcd 100644 --- a/database/models.py +++ b/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: diff --git a/database/pairs.py b/database/pairs.py index 77f8851..9738549 100644 --- a/database/pairs.py +++ b/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() diff --git a/tests/test_pair_restore_semantics.py b/tests/test_pair_restore_semantics.py new file mode 100644 index 0000000..80c7595 --- /dev/null +++ b/tests/test_pair_restore_semantics.py @@ -0,0 +1,113 @@ +"""Restores preserve pair ownership and bounded, readable JSON provenance.""" + +import json +import os +import tempfile +import unittest.mock + +from database.models import PromptModel +from database.operations import PromptDatabase + + +class TestPairRestoreSemantics(unittest.TestCase): + def setUp(self): + directory = tempfile.TemporaryDirectory() + self.addCleanup(directory.cleanup) + self.path = os.path.join(directory.name, "restore.db") + self.db = PromptDatabase(self.path) + self.addCleanup(self.db.close_all) + self.pair_id = self.db.save_prompt_pair("portrait", "blurry") + pair = self.db.get_prompt_pair(self.pair_id) + self.image_id = self.db.link_image_to_prompt( + pair["positive_prompt_id"], os.path.join(directory.name, "image.png") + ) + + def raw_link(self, provenance, pair_id=None): + with self.db.model.get_connection() as conn: + conn.execute("DELETE FROM image_prompt_pairs") + conn.execute( + "INSERT INTO image_prompt_pairs(image_id,pair_id,provenance) " + "VALUES(?,?,?)", + (self.image_id, pair_id or self.pair_id, provenance), + ) + conn.commit() + + def test_restore_rejects_pair_with_different_image_positive_owner(self): + other = self.db.save_prompt_pair("landscape", "blurry") + self.raw_link("{}", other) + ok, reason = PromptModel.verify_database_file(self.path) + self.assertFalse(ok) + self.assertIn("positive", reason.lower()) + + def test_restore_rejects_unreadable_nonobject_and_nonfinite_provenance(self): + for value in ( + "not JSON", + "[]", + "null", + "false", + "42", + '"text"', + '{"value":NaN}', + '{"value":Infinity}', + '{"value":-Infinity}', + '{"value":1e999}', + b'{"blob":true}', + ): + with self.subTest(value=value): + self.raw_link(value) + ok, reason = PromptModel.verify_database_file(self.path) + self.assertFalse(ok) + self.assertIn("provenance", reason.lower()) + + def test_restore_retains_valid_nested_object_provenance(self): + provenance = {"format": "future", "nested": [{"empty": "", "value": 2.5}]} + self.raw_link(json.dumps(provenance)) + self.assertEqual(PromptModel.verify_database_file(self.path), (True, "ok")) + self.assertEqual( + self.db.get_image_prompt_pairs(self.image_id)[0]["provenance"], provenance + ) + + def test_restore_rejects_excessive_json_nesting_without_raising(self): + self.raw_link('{"nested":' * 1500 + "0" + "}" * 1500) + ok, reason = PromptModel.verify_database_file(self.path) + self.assertFalse(ok) + self.assertIn("provenance", reason.lower()) + + def test_restore_and_runtime_share_the_provenance_size_limit(self): + provenance = {"large": "x" * (4_000_000 + 1)} + with self.assertRaises(ValueError): + self.db.link_image_to_pair(self.pair_id, self.image_id, provenance) + self.raw_link(json.dumps(provenance)) + ok, reason = PromptModel.verify_database_file(self.path) + self.assertFalse(ok) + self.assertIn("provenance", reason.lower()) + + def test_runtime_normalizes_deep_json_encoding_error(self): + provenance = {} + child = provenance + for _ in range(1500): + child["nested"] = {} + child = child["nested"] + with self.assertRaises(ValueError): + self.db.link_image_to_pair(self.pair_id, self.image_id, provenance) + self.assertEqual(self.db.get_image_prompt_pairs(self.image_id), []) + + def test_runtime_converts_encoder_recursion_failure_to_validation_error(self): + with unittest.mock.patch( + "database.pairs.json.dumps", side_effect=RecursionError("depth") + ): + with self.assertRaisesRegex(ValueError, "JSON"): + self.db.link_image_to_pair(self.pair_id, self.image_id, {}) + + def test_restore_converts_parser_recursion_failure_to_rejection(self): + self.raw_link("{}") + with unittest.mock.patch( + "database.models.json.loads", side_effect=RecursionError("depth") + ): + ok, reason = PromptModel.verify_database_file(self.path) + self.assertFalse(ok) + self.assertIn("provenance", reason.lower()) + + +if __name__ == "__main__": + unittest.main()