Vito Sansevero
a3d5c16318
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.
561 lines
20 KiB
Python
561 lines
20 KiB
Python
"""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
|