From a3d5c16318465a65e11087b57f753dd619dda143 Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Thu, 22 Oct 2026 05:14:03 -0700 Subject: [PATCH] feat(annotations): add bounded private media capture store Add utils/media_capture.py: CaptureStore copies one rooted media file into an owner-only prompt-manager-captures-* temp directory, streaming it in 1 MiB chunks while computing the whole-file SHA-256, and compares the opened descriptor with the path (dev, ino where meaningful, size, mtime_ns, ctime_ns) before and after the copy; any disagreement raises CaptureChanged and removes the copy. Symlinked files or directories, root escapes, non-media names and directories raise CaptureDenied; the locator is normcase(join(realpath(root), relative)), the listing's identity, so a symlinked root matches it. Limits are the plan's: 256 MiB per file on stat and stream (CaptureTooLarge), 8 completed captures, 512 MiB reserved and 2 capture/verify workers with nonblocking admission (CaptureCapacity), 900 s idle and 1800 s hard lifetimes on an injected clock. acquire and release_use keep an in-use capture past expiry, release and invalidate_all until its last use ends; close removes only owned files and the owned directory. A header-only probe of the copy marks Generate unsupported above 4096 px per side or 16,777,216 px, for unreadable headers and for video; verify re-hashes the current file without a copy. Every handle closes before unlink. tests/test_media_capture.py covers exact bytes and owner-only modes, alias and escape denial, the symlinked-root locator, change detection during the copy (size, same-size rewrite with restored mtime through the POSIX change time, inode replacement), both size caps, count, byte and worker admission, idle and hard expiry with in-use retention, cleanup scope, handle release at unlink, dimension and video readiness, and the allowlist equality with the inspection service. TESTING.md documents the store and the Windows change-time limit. --- TESTING.md | 21 ++ tests/test_media_capture.py | 734 ++++++++++++++++++++++++++++++++++++ utils/media_capture.py | 560 +++++++++++++++++++++++++++ 3 files changed, 1315 insertions(+) create mode 100644 tests/test_media_capture.py create mode 100644 utils/media_capture.py diff --git a/TESTING.md b/TESTING.md index 975fce0c..638b02c8 100644 --- a/TESTING.md +++ b/TESTING.md @@ -1369,6 +1369,27 @@ preserves both keeps the old set matching until the file is next observed. The browser fixtures clear `media_image_tags`, `media_annotation_bytes` and `media_annotation_files` on reset and never the metadata row. +Capture store (#249, `utils/media_capture.py`, `tests/test_media_capture.py`): +one explicit action captures one rooted media file into an owner-only +`prompt-manager-captures-*` temp directory, streaming it in 1 MiB chunks while +hashing, with the opened descriptor compared to the path (`dev`, `ino` where +meaningful, `size`, `mtime_ns`, `ctime_ns`) before and after the copy; any +disagreement is `CaptureChanged` and the copy is removed. Aliases (a symlinked +file or directory, even inside the root), escapes, non-media names and +directories are `CaptureDenied`; the locator is +`normcase(join(realpath(root), relative))`, so a symlinked root shares the +listing's identity. Fixed limits: 256 MiB per file (stat and stream), 8 +completed captures, 512 MiB reserved, 2 capture/verify workers, 900 s idle, +1800 s hard; admission is nonblocking (`CaptureCapacity`) and lifetimes are +tested with an injected clock. An acquired capture outlives expiry and release +until its last use ends. Header-only dimensions (4096 px per side, +16,777,216 px) and video mark Generate unsupported; `verify` re-hashes the +current file without a copy. Every handle closes before unlink (a patched +`os.unlink` asserts it, plus `assert_closed`). Known limit: Windows exposes no +change time through `stat`, so a same-size rewrite that restores `mtime_ns` +during the copy window completes there with a digest of the copied bytes; POSIX +detects it through `ctime`. + Budgets: `test_filter_work_is_bounded` always runs. It asserts that a plain first page makes 1 link SELECT and a plain cursor page none, that every filter shape makes exactly 1, the covering-index plan, and that the per-file values are built diff --git a/tests/test_media_capture.py b/tests/test_media_capture.py new file mode 100644 index 00000000..206f416e --- /dev/null +++ b/tests/test_media_capture.py @@ -0,0 +1,734 @@ +"""Bounded private capture store (#249 A3): exact bytes, hard limits, cleanup.""" + +import hashlib +import os +import stat +import struct +import tempfile +import threading +import time +import unittest +import zlib +from pathlib import Path +from unittest.mock import patch + +try: + from tests.open_handles import assert_closed, open_handles_under +except ImportError: # discovered with tests/ as the top-level directory + from open_handles import assert_closed, open_handles_under + +try: # guarded so a missing module fails assertions, not test loading + from utils import media_capture as capture_module + from utils.media_capture import ( + CAPTURE_EXTENSIONS, + CaptureCapacity, + CaptureChanged, + CaptureDenied, + CaptureStore, + CaptureTooLarge, + ) +except ImportError: + capture_module = None + CAPTURE_EXTENSIONS = None + CaptureCapacity = CaptureChanged = CaptureDenied = CaptureTooLarge = None + CaptureStore = None + +EPOCH = "epoch-1" +SCOPE = "a" * 64 +WAIT = 30.0 # generous ceiling for cross-thread handoffs, never a short clock + + +def _chunk(kind: bytes, data: bytes) -> bytes: + return ( + struct.pack(">I", len(data)) + + kind + + data + + struct.pack(">I", zlib.crc32(kind + data)) + ) + + +def png_bytes(width: int, height: int, pixels: bytes = b"\x00" * 16) -> bytes: + """A structurally valid PNG whose header declares *width* x *height*.""" + ihdr = struct.pack(">IIBBBBB", width, height, 8, 2, 0, 0, 0) + return ( + b"\x89PNG\r\n\x1a\n" + + _chunk(b"IHDR", ihdr) + + _chunk(b"IDAT", pixels) + + _chunk(b"IEND", b"") + ) + + +class FakeClock: + def __init__(self, now: float = 1_000.0): + self.now = now + + def __call__(self) -> float: + return self.now + + def advance(self, seconds: float) -> None: + self.now += seconds + + +class CaptureStoreTestCase(unittest.TestCase): + def setUp(self): + self.assertIsNotNone(capture_module, "utils.media_capture is not implemented") + parent = tempfile.TemporaryDirectory() + self.addCleanup(parent.cleanup) + self.addCleanup(assert_closed, self, parent.name) + self.parent = Path(parent.name) + self.root = self.parent / "output" + self.root.mkdir() + self.clock = FakeClock() + + def make_store(self, **options) -> "CaptureStore": + options.setdefault("clock", self.clock) + store = CaptureStore(self.parent / "captures", **options) + self.addCleanup(store.close) + return store + + def write(self, relative: str, data: bytes) -> Path: + path = self.root / relative + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(data) + return path + + def symlink(self, target, link): + try: + os.symlink(str(target), str(link)) + except (OSError, NotImplementedError) as error: + self.skipTest(f"symlinks unavailable here: {error}") + + def wait_for_change_time_after(self, ns: int) -> None: + """Block until the filesystem clock passes *ns* (coarse ticks on Linux).""" + probe = self.parent / "tick" + deadline = time.monotonic() + WAIT + while time.monotonic() < deadline: + probe.write_bytes(b"x") + if os.stat(probe).st_ctime_ns > ns: + return + time.sleep(0.001) + self.fail("the filesystem clock did not advance") # pragma: no cover + + def capture(self, store, relative="a.png", **options): + options.setdefault("epoch", EPOCH) + options.setdefault("scope", SCOPE) + return store.capture(self.root, relative, **options) + + def assert_empty(self, store): + self.assertFalse(any(Path(store.directory).iterdir())) + self.assertEqual( + store.stats(), {"captures": 0, "reserved_bytes": 0, "workers": 0} + ) + + # -- exact bytes and identity ------------------------------------------- + + def test_capture_hashes_and_copies_exact_bytes_owner_only(self): + source_bytes = png_bytes(3, 2, os.urandom(40)) + source = self.write("a.png", source_bytes) + store = self.make_store() + + cap = self.capture(store, root_index=1) + + self.assertEqual(Path(cap.path).read_bytes(), source_bytes) + self.assertEqual(cap.sha256, hashlib.sha256(source_bytes).hexdigest()) + self.assertEqual(cap.mtime_ns, os.stat(source).st_mtime_ns) + self.assertEqual(cap.size, len(source_bytes)) + self.assertEqual( + cap.locator, + os.path.normcase(os.path.join(os.path.realpath(self.root), "a.png")), + ) + self.assertEqual((cap.root_index, cap.scope, cap.epoch), (1, SCOPE, EPOCH)) + self.assertEqual((cap.media_type, cap.width, cap.height), ("image", 3, 2)) + self.assertTrue(cap.inferable) + self.assertIsNone(cap.inferable_reason) + self.assertIsInstance(cap.token, str) + self.assertGreaterEqual(len(cap.token), 32) + self.assertNotIn(cap.token, str(cap.path)) + self.assertNotIn(str(self.root), str(cap.path)) + capture_path = Path(cap.path) + self.assertEqual(capture_path.parent, Path(store.directory)) + self.assertEqual(capture_path.suffix, ".png") + self.assertTrue( + Path(store.directory).name.startswith("prompt-manager-captures-") + ) + if os.name == "posix": + self.assertEqual(stat.S_IMODE(os.stat(store.directory).st_mode) & 0o077, 0) + self.assertEqual(stat.S_IMODE(os.stat(cap.path).st_mode) & 0o077, 0) + self.assertIs(store.get(cap.token), cap) + self.assertEqual( + store.stats(), + {"captures": 1, "reserved_bytes": len(source_bytes), "workers": 0}, + ) + self.assertTrue(store.release(cap.token)) + self.assertFalse(store.release(cap.token)) + self.assertIsNone(store.get(cap.token)) + self.assert_empty(store) + assert_closed(self, self.parent) + + def test_streams_in_bounded_chunks_never_the_whole_file(self): + data = png_bytes(1, 1, os.urandom(5 * capture_module.CHUNK_SIZE + 7)) + self.write("big.png", data) + store = self.make_store(max_file=len(data)) + sizes = [] + original = CaptureStore._read + + def hook(store_, stream, size): + sizes.append(size) + return original(store_, stream, size) + + with patch.object(CaptureStore, "_read", hook): + cap = self.capture(store, "big.png") + self.assertTrue(sizes) + self.assertLessEqual(max(sizes), capture_module.CHUNK_SIZE) + self.assertEqual(cap.sha256, hashlib.sha256(data).hexdigest()) + + def test_symlinked_file_or_root_escape_is_denied(self): + outside = self.parent / "outside" + outside.mkdir() + (outside / "real.png").write_bytes(png_bytes(1, 1)) + self.write("a.png", png_bytes(1, 1)) + self.write("notes.txt", b"text") + (self.root / "folder").mkdir() + store = self.make_store() + + denied = [ + "../outside/real.png", + "a/../a.png", + "./a.png", + "a.png/", + "", + "/a.png", + "C:/a.png", + "a\\b.png", + "a\x00.png", + "notes.txt", + "folder", + "x" * 1021 + ".png", + ] + for relative in denied: + with self.subTest(relative=relative): + with self.assertRaises(CaptureDenied): + self.capture(store, relative) + with self.assertRaises(CaptureDenied): + self.capture(store, "a.png", root_index=True) + with self.assertRaises(CaptureChanged): + self.capture(store, "missing.png") + self.assert_empty(store) + + self.symlink(outside / "real.png", self.root / "link.png") + self.symlink(outside, self.root / "linked-dir") + with self.assertRaises(CaptureDenied): + self.capture(store, "link.png") + with self.assertRaises(CaptureDenied): + self.capture(store, "linked-dir/real.png") + self.assert_empty(store) + # A symlink to a file inside the root is still an alias, never a locator. + self.symlink(self.root / "a.png", self.root / "alias.png") + with self.assertRaises(CaptureDenied): + self.capture(store, "alias.png") + self.assert_empty(store) + + def test_locator_matches_listing_identity_for_symlinked_root(self): + real_root = self.parent / "real-output" + real_root.mkdir() + (real_root / "sub").mkdir() + source_bytes = png_bytes(2, 2) + (real_root / "sub" / "a.png").write_bytes(source_bytes) + link_root = self.parent / "link-output" + self.symlink(real_root, link_root) + store = self.make_store() + + cap = store.capture(link_root, "sub/a.png", epoch=EPOCH, scope=SCOPE) + + expected = os.path.normcase( + os.path.join(os.path.realpath(link_root), "sub", "a.png") + ) + self.assertEqual(cap.locator, expected) + self.assertEqual( + cap.locator, os.path.normcase(os.path.realpath(link_root / "sub/a.png")) + ) + self.assertTrue(cap.locator.startswith(os.path.normcase(str(real_root)))) + self.assertEqual(cap.sha256, hashlib.sha256(source_bytes).hexdigest()) + + # -- change detection ---------------------------------------------------- + + def test_change_during_copy_is_detected_and_temp_removed(self): + source_bytes = png_bytes(1, 1, b"\x11" * 40) + source = self.write("a.png", source_bytes) + store = self.make_store(max_file=4096) + original = CaptureStore._read + + def appending(store_, stream, size): + data = original(store_, stream, 16) + if data and not getattr(appending, "done", False): + appending.done = True + with open(source, "ab") as handle: + handle.write(b"\x22" * 8) + return data + + with patch.object(CaptureStore, "_read", appending): + with self.assertRaises(CaptureChanged): + store.capture(self.root, "a.png", epoch=EPOCH, scope=SCOPE) + self.assert_empty(store) + + # Same size, different bytes, mtime restored: only the change time + # (POSIX ctime) can tell. Windows exposes no change time through stat, + # so there the torn copy completes and its digest describes the copy. + source.write_bytes(source_bytes) + before = os.stat(source) + self.wait_for_change_time_after(before.st_ctime_ns) + rewritten = bytes(b ^ 0xFF for b in source_bytes) + + def rewriting(store_, stream, size): + data = original(store_, stream, 16) + if data and not getattr(rewriting, "done", False): + rewriting.done = True + with open(source, "r+b") as handle: + handle.write(rewritten) + os.utime(source, ns=(before.st_atime_ns, before.st_mtime_ns)) + return data + + with patch.object(CaptureStore, "_read", rewriting): + if os.name == "posix": + with self.assertRaises(CaptureChanged): + store.capture(self.root, "a.png", epoch=EPOCH, scope=SCOPE) + else: + cap = store.capture(self.root, "a.png", epoch=EPOCH, scope=SCOPE) + copied = Path(cap.path).read_bytes() + self.assertEqual(cap.sha256, hashlib.sha256(copied).hexdigest()) + store.release(cap.token) + self.assert_empty(store) + self.assertEqual(os.stat(source).st_mtime_ns, before.st_mtime_ns) + + if os.name == "posix": + # The path now names another file (new inode) once the copy ends. + source.write_bytes(source_bytes) + + def replacing(store_, stream, size): + data = original(store_, stream, size) + if not data and not getattr(replacing, "done", False): + replacing.done = True + other = self.root / "other.png" + other.write_bytes(source_bytes) + os.replace(other, source) + return data + + with patch.object(CaptureStore, "_read", replacing): + with self.assertRaises(CaptureChanged): + store.capture(self.root, "a.png", epoch=EPOCH, scope=SCOPE) + self.assert_empty(store) + assert_closed(self, self.parent) + + def test_size_cap_on_stat_and_stream(self): + store = self.make_store(max_file=64) + self.write("large.png", bytes(65)) + self.write("fits.png", bytes(64)) + + def never(store_, stream, size): # pragma: no cover - guard only + raise AssertionError("an over-cap file must not be read") + + with patch.object(CaptureStore, "_read", never): + with self.assertRaises(CaptureTooLarge): + self.capture(store, "large.png") + self.assert_empty(store) + + cap = self.capture(store, "fits.png") + self.assertEqual(cap.size, 64) + store.release(cap.token) + + growing = self.write("growing.png", bytes(40)) + original = CaptureStore._read + + def grow(store_, stream, size): + data = original(store_, stream, 16) + if data and not getattr(grow, "done", False): + grow.done = True + with open(growing, "ab") as handle: + handle.write(bytes(40)) + return data + + with patch.object(CaptureStore, "_read", grow): + with self.assertRaises(CaptureTooLarge): + self.capture(store, "growing.png") + self.assert_empty(store) + + # -- admission ----------------------------------------------------------- + + def test_count_byte_and_worker_admission_is_nonblocking(self): + for name in ("a", "b", "c"): + self.write(f"{name}.png", bytes(60)) + store = self.make_store(max_captures=2, max_bytes=1000) + first = self.capture(store, "a.png") + self.capture(store, "b.png") + with self.assertRaises(CaptureCapacity): + self.capture(store, "c.png") + self.assertTrue(store.release(first.token)) + third = self.capture(store, "c.png") + self.assertEqual(store.stats()["captures"], 2) + store.invalidate_all() + self.assertIsNone(store.get(third.token)) + self.assert_empty(store) + + bytes_store = self.make_store(max_bytes=100) + kept = self.capture(bytes_store, "a.png") + with self.assertRaises(CaptureCapacity): + self.capture(bytes_store, "b.png") + self.assertEqual(bytes_store.stats()["reserved_bytes"], 60) + bytes_store.release(kept.token) + self.capture(bytes_store, "b.png") + + workers = self.make_store(max_workers=1) + entered, proceed = threading.Event(), threading.Event() + original = CaptureStore._read + results = {} + + def blocking(store_, stream, size): + if not entered.is_set(): + entered.set() + self.assertTrue(proceed.wait(WAIT)) + return original(store_, stream, size) + + def worker(): + try: + results["capture"] = self.capture(workers, "a.png") + except Exception as error: # pragma: no cover - reported below + results["error"] = error + + with patch.object(CaptureStore, "_read", blocking): + thread = threading.Thread(target=worker) + thread.start() + try: + self.assertTrue(entered.wait(WAIT)) + self.assertEqual(workers.stats()["workers"], 1) + with self.assertRaises(CaptureCapacity): + self.capture(workers, "b.png") + finally: + proceed.set() + thread.join(WAIT) + self.assertFalse(thread.is_alive()) + self.assertNotIn("error", results) + self.assertEqual(workers.stats()["workers"], 0) + self.assertEqual(results["capture"].size, 60) + self.capture(workers, "b.png") + + # -- lifetime ------------------------------------------------------------ + + def test_idle_and_hard_expiry_and_in_use_retention(self): + self.write("a.png", png_bytes(1, 1)) + self.write("b.png", png_bytes(1, 1)) + store = self.make_store(idle=900, hard=1800) + cap = self.capture(store, "a.png") + + self.clock.advance(899) + self.assertIs(store.get(cap.token), cap) # touches idle + self.assertAlmostEqual(store.expires_in(cap.token), 900) + self.clock.advance(899) + self.assertIs(store.get(cap.token), cap) + self.assertAlmostEqual(store.expires_in(cap.token), 2) + self.clock.advance(2) + self.assertIsNone(store.get(cap.token)) # hard limit despite touches + self.assertIsNone(store.expires_in(cap.token)) + self.assertFalse(Path(cap.path).exists()) + self.assert_empty(store) + + idle = self.capture(store, "a.png") + self.clock.advance(900) + store.sweep() + self.assertIsNone(store.get(idle.token)) + self.assert_empty(store) + + used = self.capture(store, "b.png") + self.assertIs(store.acquire(used.token), used) + self.assertIs(store.acquire(used.token), used) + self.clock.advance(901) + store.sweep() + self.assertIs(store.get(used.token), used) # in use is not idle + self.clock.advance(900) + store.sweep() + self.assertIsNone(store.get(used.token)) # expired for new requests + self.assertIsNone(store.acquire(used.token)) + self.assertTrue(Path(used.path).exists()) # retained while in use + self.assertEqual(store.stats()["captures"], 1) + store.release_use(used.token) + self.assertTrue(Path(used.path).exists()) + store.release_use(used.token) + self.assertFalse(Path(used.path).exists()) + self.assert_empty(store) + store.release_use(used.token) # unknown token: no error, nothing freed + self.assertIsNone(store.acquire("unknown")) + + def test_release_while_in_use_keeps_file_until_last_use_ends(self): + self.write("a.png", png_bytes(1, 1)) + store = self.make_store() + cap = self.capture(store, "a.png") + self.assertIs(store.acquire(cap.token), cap) + self.assertTrue(store.release(cap.token)) + self.assertIsNone(store.get(cap.token)) + self.assertTrue(Path(cap.path).exists()) + store.invalidate_all() + self.assertTrue(Path(cap.path).exists()) + store.release_use(cap.token) + self.assertFalse(Path(cap.path).exists()) + self.assert_empty(store) + + def test_invalidate_all_and_close_remove_owned_directory_only(self): + self.write("a.png", png_bytes(1, 1)) + self.write("b.png", png_bytes(1, 1)) + sibling = self.parent / "captures" / "sibling.txt" + sibling.parent.mkdir() + sibling.write_bytes(b"keep") + store = CaptureStore(self.parent / "captures", clock=self.clock) + try: + a = self.capture(store, "a.png") + b = self.capture(store, "b.png") + self.assertEqual(Path(store.directory).parent, sibling.parent) + store.invalidate_all() + self.assertIsNone(store.get(a.token)) + self.assertIsNone(store.get(b.token)) + self.assert_empty(store) + self.assertTrue(Path(store.directory).is_dir()) + c = self.capture(store, "a.png") + self.assertTrue(Path(c.path).exists()) + finally: + store.close() + self.assertFalse(Path(store.directory).exists()) + self.assertEqual(sibling.read_bytes(), b"keep") + self.assertTrue((self.root / "a.png").exists()) + store.close() # idempotent + with self.assertRaises(RuntimeError): + self.capture(store, "a.png") + self.assertIsNone(store.get(c.token)) + assert_closed(self, self.parent) + + def test_windows_handles_closed_before_unlink(self): + source_bytes = png_bytes(1, 1, b"\x33" * 40) + source = self.write("a.png", source_bytes) + store = self.make_store() + original_unlink = os.unlink + unlinked = [] + + def checked_unlink(path, *args, **kwargs): + self.assertEqual(open_handles_under(path), []) + unlinked.append(os.fspath(path)) + return original_unlink(path, *args, **kwargs) + + with patch("os.unlink", checked_unlink): + cap = self.capture(store, "a.png") + self.assertEqual(open_handles_under(cap.path), []) + self.assertEqual(open_handles_under(source), []) + store.release(cap.token) + original = CaptureStore._read + + def appending(store_, stream, size): + data = original(store_, stream, 16) + if data and not getattr(appending, "done", False): + appending.done = True + with open(source, "ab") as handle: + handle.write(b"\x44") + return data + + with patch.object(CaptureStore, "_read", appending): + with self.assertRaises(CaptureChanged): + self.capture(store, "a.png") + self.assertEqual(open_handles_under(source), []) + self.assertEqual(len(unlinked), 2) + self.assertTrue(all(p.startswith(str(store.directory)) for p in unlinked)) + self.assert_empty(store) + assert_closed(self, self.parent) + + # -- inference readiness ------------------------------------------------- + + def test_dimension_limits_mark_generate_unsupported(self): + cases = { + "edge.png": (png_bytes(4096, 4096), True), + "wide.png": (png_bytes(4097, 1), False), + "tall.png": (png_bytes(1, 4097), False), + } + for name, (data, _) in cases.items(): + self.write(name, data) + self.write("garbage.jpg", b"\xff\xd8not really a jpeg") + store = self.make_store() + + for name, (data, supported) in cases.items(): + with self.subTest(name=name): + cap = self.capture(store, name) + self.assertEqual(cap.inferable, supported) + self.assertEqual((cap.width, cap.height) != (None, None), True) + if supported: + self.assertIsNone(cap.inferable_reason) + else: + self.assertIn("4096", cap.inferable_reason) + store.release(cap.token) + self.assertEqual( + (capture_module.MAX_INFERENCE_SIDE, capture_module.MAX_INFERENCE_PIXELS), + (4096, 16_777_216), + ) + cap = self.capture(store, "garbage.jpg") + self.assertEqual((cap.media_type, cap.width, cap.height), ("image", None, None)) + self.assertFalse(cap.inferable) + self.assertIn("header", cap.inferable_reason) + self.assertEqual( + cap.sha256, hashlib.sha256(b"\xff\xd8not really a jpeg").hexdigest() + ) + + def test_video_capture_is_not_inferable(self): + data = b"\x00\x00\x00\x18ftypmp42" + os.urandom(32) + self.write("clip.mp4", data) + store = self.make_store() + + def never(path): # pragma: no cover - guard only + raise AssertionError("video captures are never probed") + + with patch.object(capture_module, "read_inspection_file", never): + cap = self.capture(store, "clip.mp4") + self.assertEqual((cap.media_type, cap.width, cap.height), ("video", None, None)) + self.assertFalse(cap.inferable) + self.assertEqual( + cap.inferable_reason, + "Auto-tagging is available for supported still images only.", + ) + self.assertEqual(Path(cap.path).suffix, ".mp4") + self.assertEqual(cap.sha256, hashlib.sha256(data).hexdigest()) + + # -- verification without a copy ------------------------------------- + + def test_verify_hashes_current_file_without_retaining_a_copy(self): + source_bytes = png_bytes(1, 1, b"\x55" * 30) + source = self.write("a.png", source_bytes) + store = self.make_store(max_file=128, max_workers=1) + + facts = store.verify(self.root, "a.png") + + self.assertEqual(facts.sha256, hashlib.sha256(source_bytes).hexdigest()) + self.assertEqual(facts.size, len(source_bytes)) + self.assertEqual(facts.mtime_ns, os.stat(source).st_mtime_ns) + self.assertEqual( + facts.locator, + os.path.normcase(os.path.join(os.path.realpath(self.root), "a.png")), + ) + self.assert_empty(store) + self.write("large.png", bytes(129)) + with self.assertRaises(CaptureTooLarge): + store.verify(self.root, "large.png") + with self.assertRaises(CaptureDenied): + store.verify(self.root, "../a.png") + self.assert_empty(store) + + # -- arguments and failure paths --------------------------------------- + + def test_constructor_and_arguments_are_validated(self): + for options in ( + {"max_file": 0}, + {"max_captures": True}, + {"max_bytes": -1}, + {"max_workers": "2"}, + {"idle": 0}, + {"hard": -5}, + ): + with self.subTest(options=options): + with self.assertRaises(ValueError): + CaptureStore(self.parent / "captures", **options) + self.write("a.png", png_bytes(1, 1)) + store = self.make_store() + with self.assertRaises(ValueError): + store.capture(self.root, "a.png", epoch="", scope=SCOPE) + with self.assertRaises(ValueError): + store.capture(self.root, "a.png", epoch=EPOCH, scope=None) + with self.assertRaises(CaptureDenied): + self.capture(store, root_index=-1) + with self.assertRaises(CaptureDenied): + store.capture( + self.parent / "missing-root", "a.png", epoch=EPOCH, scope=SCOPE + ) + self.assert_empty(store) + + default = CaptureStore(clock=self.clock) + try: + self.assertEqual( + Path(default.directory).parent, Path(tempfile.gettempdir()) + ) + self.assertTrue(Path(default.directory).is_dir()) + finally: + default.close() + self.assertFalse(Path(default.directory).exists()) + + def test_directory_target_and_unreadable_file(self): + (self.root / "dir.png").mkdir() + store = self.make_store() + with self.assertRaises(CaptureDenied): + self.capture(store, "dir.png") + self.assert_empty(store) + if os.name != "posix" or os.geteuid() == 0: + return # read permission cannot be withdrawn here + locked = self.write("locked.png", png_bytes(1, 1)) + os.chmod(locked, 0) + try: + with self.assertRaises(CaptureChanged): + self.capture(store, "locked.png") + finally: + os.chmod(locked, 0o600) + self.assert_empty(store) + + def test_removal_failures_are_logged_and_never_touch_other_files(self): + self.write("a.png", png_bytes(1, 1)) + store = CaptureStore(self.parent / "captures", clock=self.clock) + stray = Path(store.directory) / "stray.bin" + try: + cap = self.capture(store, "a.png") + self.assertIs(store.acquire(cap.token), cap) + self.assertAlmostEqual(store.expires_in(cap.token), 1800) # in use + store.release_use(cap.token) + original_unlink = os.unlink + + def failing(path, *args, **kwargs): + if os.fspath(path) == cap.path: + raise PermissionError("busy") + return original_unlink(path, *args, **kwargs) + + with patch("os.unlink", failing): + with self.assertLogs("utils.media_capture", "WARNING") as logs: + self.assertTrue(store.release(cap.token)) + self.assertIn("could not be removed", logs.output[0]) + self.assertIsNone(store.get(cap.token)) + self.assertTrue(Path(cap.path).exists()) + os.unlink(cap.path) + stray.write_bytes(b"not ours") + finally: + with self.assertLogs("utils.media_capture", "WARNING") as logs: + store.close() + self.assertIn("could not be removed", logs.output[0]) + self.assertEqual(stray.read_bytes(), b"not ours") + stray.unlink() + os.rmdir(store.directory) + assert_closed(self, self.parent) + + # -- constants --------------------------------------------------------- + + def test_defaults_match_the_plan(self): + store = self.make_store() + self.assertEqual( + ( + store.max_file, + store.max_captures, + store.max_bytes, + store.max_workers, + store.idle, + store.hard, + ), + (256 * 2**20, 8, 512 * 2**20, 2, 900, 1800), + ) + from py.api.inspection_service import SERVABLE_EXTENSIONS + + self.assertEqual(CAPTURE_EXTENSIONS, SERVABLE_EXTENSIONS) + + def test_exceptions_carry_codes(self): + self.assertEqual(CaptureCapacity.code, "review-capacity") + self.assertEqual(CaptureChanged.code, "file-changed") + self.assertEqual(CaptureTooLarge.code, "unsupported-media") + self.assertEqual(CaptureDenied.code, "source-changed") + for cls in (CaptureCapacity, CaptureChanged, CaptureTooLarge, CaptureDenied): + self.assertTrue(issubclass(cls, capture_module.CaptureError)) + + +if __name__ == "__main__": + unittest.main() diff --git a/utils/media_capture.py b/utils/media_capture.py new file mode 100644 index 00000000..cf5874bc --- /dev/null +++ b/utils/media_capture.py @@ -0,0 +1,560 @@ +"""Bounded private capture store for image annotations (#249 A3). + +One explicit Tag or review action captures exactly one selected media file: +the rooted file is streamed to a private, owner-only temporary copy while its +whole-file SHA-256 is computed, and the identity of the opened descriptor is +checked against the path before and after the copy. Preview, inference and the +metadata read use only that copy, so what the user reviewed is what the model +saw and what Save re-checks. Everything is bounded: bytes per file, completed +captures, reserved bytes, concurrent capture/verify workers, idle and hard +lifetimes. + +The contract is "these bytes at this location, last checked at this time". +Stat equality is never content proof: a same-size rewrite that restores +``mtime_ns`` is visible only through the POSIX change time, which Windows does +not expose through ``stat``. +""" + +from __future__ import annotations + +import hashlib +import logging +import os +import secrets +import stat as stat_module +import tempfile +import threading +import time +import warnings +from dataclasses import dataclass +from pathlib import Path +from typing import Callable + +from PIL import Image + +try: + from .prompt_inspector import MetadataReadError, read_inspection_file + from .video_metadata import VIDEO_EXTENSIONS, is_video_path +except ImportError: # imported with the repo root on sys.path + from utils.prompt_inspector import MetadataReadError, read_inspection_file + from utils.video_metadata import VIDEO_EXTENSIONS, is_video_path + +logger = logging.getLogger(__name__) + +CHUNK_SIZE = 1024 * 1024 +DEFAULT_MAX_FILE = 256 * 2**20 +DEFAULT_MAX_CAPTURES = 8 +DEFAULT_MAX_BYTES = 512 * 2**20 +DEFAULT_MAX_WORKERS = 2 +DEFAULT_IDLE_SECONDS = 900 +DEFAULT_HARD_SECONDS = 1800 +MAX_INFERENCE_SIDE = 4096 +MAX_INFERENCE_PIXELS = 16_777_216 +MAX_RELATIVE_LENGTH = 1024 + +VIDEO_ONLY_REASON = "Auto-tagging is available for supported still images only." +HEADER_REASON = "The image header could not be read." +DIMENSION_REASON = ( + "Images larger than 4096 px per side or 16,777,216 px are not supported." +) + +# The listing's media allowlist (``py.api.inspection_service.SERVABLE_EXTENSIONS``), +# restated here so utilities never import the aiohttp API package; a test pins +# the two sets equal. +_STILL_EXTENSIONS = frozenset( + {".png", ".jpg", ".jpeg", ".webp", ".gif", ".bmp", ".tiff", ".tif"} +) +CAPTURE_EXTENSIONS = _STILL_EXTENSIONS | frozenset(VIDEO_EXTENSIONS) + +_DIRECTORY_PREFIX = "prompt-manager-captures-" +_OPEN_FLAGS = ( + os.O_RDONLY + | getattr(os, "O_BINARY", 0) + | getattr(os, "O_NOFOLLOW", 0) + | getattr(os, "O_CLOEXEC", 0) +) + + +class CaptureError(Exception): + """Base of the capture outcomes; ``code`` is the API error code.""" + + code = "capture-error" + + +class CaptureCapacity(CaptureError): + """Count, byte or worker admission failed (nonblocking).""" + + code = "review-capacity" + + +class CaptureTooLarge(CaptureError): + """The file exceeds the per-file byte cap on stat or on the stream.""" + + code = "unsupported-media" + + +class CaptureChanged(CaptureError): + """The file changed, vanished or was replaced across the copy (retryable).""" + + code = "file-changed" + + +class CaptureDenied(CaptureError): + """The target is not a listed rooted media file (alias, escape, non-media).""" + + code = "source-changed" + + +@dataclass(frozen=True) +class FileFacts: + """Whole-file facts of the current rooted file, without a retained copy.""" + + locator: str + sha256: str + size: int + mtime_ns: int + + +@dataclass(frozen=True) +class Capture: + """One completed private copy; ``path`` is the owner-only temporary file.""" + + token: str + locator: str + root_index: int + scope: str + epoch: str + sha256: str + size: int + mtime_ns: int + path: str + media_type: str + width: int | None + height: int | None + inferable: bool + inferable_reason: str | None + + +class _Entry: + __slots__ = ("capture", "created_at", "last_used_at", "uses", "doomed") + + def __init__(self, capture: Capture, now: float): + self.capture = capture + self.created_at = now + self.last_used_at = now + self.uses = 0 + self.doomed = False + + +def _same_file(before: os.stat_result, after: os.stat_result) -> bool: + """Size, mtime and change time must agree; device/inode where meaningful.""" + if ( + before.st_size != after.st_size + or before.st_mtime_ns != after.st_mtime_ns + or before.st_ctime_ns != after.st_ctime_ns + ): + return False + if before.st_ino and after.st_ino and before.st_ino != after.st_ino: + return False + if before.st_dev and after.st_dev and before.st_dev != after.st_dev: + return False + return True + + +def _unlink_quietly(path: str) -> None: + try: + os.unlink(path) + except FileNotFoundError: + pass + except OSError as error: + logger.warning("Capture file %s could not be removed: %s", path, error) + + +class CaptureStore: + """Bounded, lock-protected registry of private media captures.""" + + def __init__( + self, + temp_parent: os.PathLike | str | None = None, + *, + max_file: int = DEFAULT_MAX_FILE, + max_captures: int = DEFAULT_MAX_CAPTURES, + max_bytes: int = DEFAULT_MAX_BYTES, + max_workers: int = DEFAULT_MAX_WORKERS, + idle: float = DEFAULT_IDLE_SECONDS, + hard: float = DEFAULT_HARD_SECONDS, + clock: Callable[[], float] = time.monotonic, + ): + for name, value in ( + ("max_file", max_file), + ("max_captures", max_captures), + ("max_bytes", max_bytes), + ("max_workers", max_workers), + ): + if isinstance(value, bool) or not isinstance(value, int) or value < 1: + raise ValueError(f"{name} must be a positive integer") + if not (idle > 0 and hard > 0): + raise ValueError("idle and hard must be positive") + self.max_file = max_file + self.max_captures = max_captures + self.max_bytes = max_bytes + self.max_workers = max_workers + self.idle = idle + self.hard = hard + self._clock = clock + parent = None if temp_parent is None else os.fspath(temp_parent) + if parent is not None: + os.makedirs(parent, exist_ok=True) + # mkdtemp creates the directory with owner-only permissions. + self.directory = tempfile.mkdtemp(prefix=_DIRECTORY_PREFIX, dir=parent) + self._lock = threading.Lock() + self._entries: dict[str, _Entry] = {} + self._completed_bytes = 0 + self._pending = 0 + self._pending_bytes = 0 + self._workers = 0 + self._closed = False + + # -- public API ------------------------------------------------------------ + + def capture( + self, + root: os.PathLike | str, + relative: str, + *, + epoch: str, + scope: str, + root_index: int = 0, + ) -> Capture: + """Copy and hash one rooted media file into a private capture.""" + if isinstance(root_index, bool) or not isinstance(root_index, int): + raise CaptureDenied("Invalid gallery root") + if root_index < 0: + raise CaptureDenied("Invalid gallery root") + for name, value in (("epoch", epoch), ("scope", scope)): + if not isinstance(value, str) or not value: + raise ValueError(f"{name} must be a non-empty string") + parts, extension = _relative_parts(relative) + candidate, locator = _resolve(root, parts) + before = self._pre_stat(candidate) + self._admit(before.st_size, reserve=True) + try: + capture = self._copy( + candidate, locator, extension, before, root_index, epoch, scope + ) + finally: + self._release_admission(before.st_size, reserve=True) + return capture + + def verify(self, root: os.PathLike | str, relative: str) -> FileFacts: + """Hash the current rooted file under the same bounds, keeping no copy.""" + parts, _extension = _relative_parts(relative) + candidate, locator = _resolve(root, parts) + before = self._pre_stat(candidate) + self._admit(0, reserve=False) + try: + sha256 = self._stream(candidate, locator, before, None) + finally: + self._release_admission(0, reserve=False) + return FileFacts(locator, sha256, before.st_size, before.st_mtime_ns) + + def get(self, token: str) -> Capture | None: + """The live capture for *token*, touching its idle clock; else None.""" + with self._lock: + entry = self._live_locked(token) + if entry is None: + return None + entry.last_used_at = self._clock() + return entry.capture + + def acquire(self, token: str) -> Capture | None: + """Mark the capture in use (inference); it outlives expiry until released.""" + with self._lock: + entry = self._live_locked(token) + if entry is None: + return None + entry.uses += 1 + entry.last_used_at = self._clock() + return entry.capture + + def release_use(self, token: str) -> None: + """End one ``acquire``; a doomed capture is removed with its last use.""" + with self._lock: + entry = self._entries.get(token) + if entry is None: + return + entry.uses = max(0, entry.uses - 1) + entry.last_used_at = self._clock() + if entry.doomed and entry.uses == 0: + self._remove_locked(entry) + + def release(self, token: str) -> bool: + """Explicit release; True when the token was live.""" + with self._lock: + entry = self._entries.get(token) + if entry is None or entry.doomed: + return False + self._doom_locked(entry) + return True + + def invalidate_all(self) -> None: + """Expire every capture now (storage epoch changed, restore, shutdown).""" + with self._lock: + for entry in list(self._entries.values()): + self._doom_locked(entry) + + def sweep(self) -> None: + """Remove idle and hard-expired captures that are not in use.""" + with self._lock: + self._sweep_locked(self._clock()) + + def expires_in(self, token: str) -> float | None: + """Seconds until *token* expires without further use; None when gone.""" + with self._lock: + entry = self._live_locked(token) + if entry is None: + return None + now = self._clock() + remaining = self.hard - (now - entry.created_at) + if entry.uses == 0: + remaining = min(remaining, self.idle - (now - entry.last_used_at)) + return max(0.0, remaining) + + def stats(self) -> dict: + with self._lock: + return { + "captures": len(self._entries), + "reserved_bytes": self._completed_bytes + self._pending_bytes, + "workers": self._workers, + } + + def close(self) -> None: + """Remove every owned capture file and the owned directory. Idempotent.""" + with self._lock: + self._closed = True + for entry in list(self._entries.values()): + self._remove_locked(entry) + try: + os.rmdir(self.directory) + except FileNotFoundError: + pass + except OSError as error: + logger.warning( + "Capture directory %s could not be removed: %s", self.directory, error + ) + + # -- admission ------------------------------------------------------------- + + def _admit(self, size: int, *, reserve: bool) -> None: + with self._lock: + if self._closed: + raise RuntimeError("The capture store is closed") + self._sweep_locked(self._clock()) + if self._workers >= self.max_workers: + raise CaptureCapacity("Too many captures are in progress") + if reserve: + if len(self._entries) + self._pending >= self.max_captures: + raise CaptureCapacity("Too many captures are held") + reserved = self._completed_bytes + self._pending_bytes + if reserved + size > self.max_bytes: + raise CaptureCapacity("Capture byte budget exhausted") + self._pending += 1 + self._pending_bytes += size + self._workers += 1 + + def _release_admission(self, size: int, *, reserve: bool) -> None: + with self._lock: + self._workers -= 1 + if reserve: + self._pending -= 1 + self._pending_bytes -= size + + def _register(self, capture: Capture) -> None: + with self._lock: + if self._closed: + raise RuntimeError("The capture store is closed") + self._completed_bytes += capture.size + self._entries[capture.token] = _Entry(capture, self._clock()) + + # -- lifetime (all under the lock) ----------------------------------------- + + def _expired(self, entry: _Entry, now: float) -> bool: + if now - entry.created_at >= self.hard: + return True + return entry.uses == 0 and now - entry.last_used_at >= self.idle + + def _live_locked(self, token: str) -> _Entry | None: + entry = self._entries.get(token) + if entry is None or entry.doomed: + return None + if self._expired(entry, self._clock()): + self._doom_locked(entry) + return None + return entry + + def _sweep_locked(self, now: float) -> None: + for entry in list(self._entries.values()): + if entry.doomed or self._expired(entry, now): + self._doom_locked(entry) + + def _doom_locked(self, entry: _Entry) -> None: + entry.doomed = True + if entry.uses == 0: + self._remove_locked(entry) + + def _remove_locked(self, entry: _Entry) -> None: + _unlink_quietly(entry.capture.path) + if self._entries.pop(entry.capture.token, None) is not None: + self._completed_bytes -= entry.capture.size + + # -- copy and hash ----------------------------------------------------------- + + def _pre_stat(self, candidate: str) -> os.stat_result: + try: + facts = os.lstat(candidate) + except FileNotFoundError: + raise CaptureChanged("The file is no longer there") from None + except OSError as error: + raise CaptureChanged(f"The file could not be read: {error}") from None + if stat_module.S_ISLNK(facts.st_mode): + raise CaptureDenied("Symbolic links are not captured") + if not stat_module.S_ISREG(facts.st_mode): + raise CaptureDenied("Only regular files are captured") + if facts.st_size > self.max_file: + raise CaptureTooLarge("The file is larger than the capture limit") + return facts + + def _copy( + self, + candidate: str, + locator: str, + extension: str, + before: os.stat_result, + root_index: int, + epoch: str, + scope: str, + ) -> Capture: + fd, temp_path = tempfile.mkstemp( + prefix="capture-", suffix=extension, dir=self.directory + ) + sink = os.fdopen(fd, "wb") + completed = False + try: + sha256 = self._stream(candidate, locator, before, sink) + sink.close() + media_type = "video" if is_video_path(candidate) else "image" + width, height, inferable, reason = self._probe(temp_path, media_type) + capture = Capture( + token=secrets.token_urlsafe(32), + locator=locator, + root_index=root_index, + scope=scope, + epoch=epoch, + sha256=sha256, + size=before.st_size, + mtime_ns=before.st_mtime_ns, + path=temp_path, + media_type=media_type, + width=width, + height=height, + inferable=inferable, + inferable_reason=reason, + ) + self._register(capture) + completed = True + return capture + finally: + if not sink.closed: + sink.close() + if not completed: + _unlink_quietly(temp_path) + + def _read(self, stream, size: int) -> bytes: + """One bounded read of the source; a seam for change-detection tests.""" + return stream.read(size) + + def _stream(self, candidate: str, locator: str, before, sink) -> str: + """Hash (and copy into *sink*) the file; every handle closes here.""" + digest = hashlib.sha256() + total = 0 + try: + fd = os.open(candidate, _OPEN_FLAGS) + except OSError as error: + raise CaptureChanged(f"The file could not be opened: {error}") from None + with os.fdopen(fd, "rb") as stream: + opened = os.fstat(stream.fileno()) + if not _same_file(before, opened): + raise CaptureChanged("The file changed before it was read") + while True: + chunk = self._read(stream, CHUNK_SIZE) + if not chunk: + break + total += len(chunk) + if total > self.max_file: + raise CaptureTooLarge("The file grew past the capture limit") + digest.update(chunk) + if sink is not None: + sink.write(chunk) + after = os.fstat(stream.fileno()) + if total != before.st_size or not _same_file(before, after): + raise CaptureChanged("The file changed while it was read") + try: + current = os.stat(candidate) + except OSError: + raise CaptureChanged("The file is no longer there") from None + if not _same_file(before, current): + raise CaptureChanged("The path names another file now") + if os.path.normcase(os.path.realpath(candidate)) != locator: + raise CaptureChanged("The path resolves elsewhere now") + return digest.hexdigest() + + def _probe(self, path: str, media_type: str): + """Header-only dimensions of the capture; never decodes pixels.""" + if media_type == "video": + return None, None, False, VIDEO_ONLY_REASON + try: + with warnings.catch_warnings(): + warnings.simplefilter("ignore", Image.DecompressionBombWarning) + read = read_inspection_file(Path(path)) + except (MetadataReadError, Image.DecompressionBombError): + return None, None, False, HEADER_REASON + width, height = read.width, read.height + if width is None or height is None: + return None, None, False, HEADER_REASON + if ( + width > MAX_INFERENCE_SIDE + or height > MAX_INFERENCE_SIDE + or width * height > MAX_INFERENCE_PIXELS + ): + return width, height, False, DIMENSION_REASON + return width, height, True, None + + +def _relative_parts(relative) -> tuple[list[str], str]: + """Split a safe relative POSIX media path; anything else is denied.""" + if not isinstance(relative, str) or not relative: + raise CaptureDenied("Invalid file path") + if len(relative) > MAX_RELATIVE_LENGTH or "\x00" in relative: + raise CaptureDenied("Invalid file path") + if "\\" in relative or relative.startswith("/"): + raise CaptureDenied("Invalid file path") + parts = relative.split("/") + for part in parts: + if part in ("", ".", "..") or ":" in part: + raise CaptureDenied("Invalid file path") + extension = os.path.splitext(parts[-1])[1].lower() + if extension not in CAPTURE_EXTENSIONS: + raise CaptureDenied("Unsupported media type") + return parts, extension + + +def _resolve(root, parts: list[str]) -> tuple[str, str]: + """The candidate path and its locator; aliases and escapes are denied.""" + real_root = os.path.realpath(os.fspath(root)) + if not os.path.isdir(real_root): + raise CaptureDenied("Unknown gallery root") + candidate = os.path.join(real_root, *parts) + locator = os.path.normcase(candidate) + if os.path.normcase(os.path.realpath(candidate)) != locator: + raise CaptureDenied("The file is not a listed rooted file") + return candidate, locator