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.
This commit is contained in:
Vito Sansevero
+21
@@ -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
|
||||
|
||||
@@ -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()
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user