fix(database): validate restored pair ownership and provenance (#193)
This commit is contained in:
Vito Sansevero
@@ -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:
|
||||
|
||||
+5
-2
@@ -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()
|
||||
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user