"""Fail-closed local upload storage with bounded atomic writes.""" from __future__ import annotations import asyncio import hashlib import os import re import stat import uuid from collections.abc import AsyncIterable from dataclasses import dataclass from enum import StrEnum from pathlib import Path from typing import Final _SHA256_PATTERN: Final = re.compile(r"^[0-9a-f]{64}$") _WRITE_MODE: Final = 0o600 _FINAL_MODE: Final = 0o440 class StorageErrorCode(StrEnum): INVALID_CONTRACT = "INVALID_CONTRACT" TOO_LARGE = "TOO_LARGE" SIZE_MISMATCH = "SIZE_MISMATCH" HASH_MISMATCH = "HASH_MISMATCH" OBJECT_CONFLICT = "OBJECT_CONFLICT" ROOT_UNSAFE = "ROOT_UNSAFE" IO_UNAVAILABLE = "IO_UNAVAILABLE" _SAFE_MESSAGES: Final[dict[StorageErrorCode, str]] = { StorageErrorCode.INVALID_CONTRACT: "The upload storage contract is invalid.", StorageErrorCode.TOO_LARGE: "The upload exceeds its configured size limit.", StorageErrorCode.SIZE_MISMATCH: "The uploaded byte count does not match its declaration.", StorageErrorCode.HASH_MISMATCH: "The uploaded content digest does not match its declaration.", StorageErrorCode.OBJECT_CONFLICT: "The upload object already exists with different content.", StorageErrorCode.ROOT_UNSAFE: "The configured upload storage root is not safe.", StorageErrorCode.IO_UNAVAILABLE: "The upload storage is temporarily unavailable.", } class LocalStorageError(RuntimeError): """A sanitized storage error that never contains paths or upload bytes.""" def __init__(self, code: StorageErrorCode) -> None: self.code = code super().__init__(_SAFE_MESSAGES[code]) def __repr__(self) -> str: return f"{type(self).__name__}(code={self.code.value!r})" @dataclass(frozen=True, slots=True) class StoredUpload: storage_key: uuid.UUID byte_size: int sha256: str class LocalUploadStorage: """Store UUID-keyed objects below one absolute, non-symlink root. Client filenames are intentionally absent from this adapter. Writes use a same-directory temporary file, fsync, read-only permissions, and atomic replacement. A retry for an already stored identical object is idempotent. """ def __init__(self, root: Path, *, max_bytes: int) -> None: if not root.is_absolute() or isinstance(max_bytes, bool) or max_bytes <= 0: raise LocalStorageError(StorageErrorCode.INVALID_CONTRACT) self._root = root self._max_bytes = max_bytes async def store( self, *, storage_key: uuid.UUID, chunks: AsyncIterable[bytes], expected_size: int, expected_sha256: str, ) -> StoredUpload: self._validate_expectation(expected_size, expected_sha256) try: directory, target = await asyncio.to_thread(self._prepare_target, storage_key) existing = await asyncio.to_thread( self._existing_object, target, storage_key, expected_size, expected_sha256, ) except LocalStorageError: raise except OSError: raise LocalStorageError(StorageErrorCode.IO_UNAVAILABLE) from None if existing is not None: return existing temporary = directory / f".{storage_key.hex}.{uuid.uuid4().hex}.upload" descriptor: int | None = None total = 0 digest = hashlib.sha256() installed = False try: descriptor = await asyncio.to_thread(self._open_temporary, temporary) async for chunk in chunks: if not isinstance(chunk, bytes): raise LocalStorageError(StorageErrorCode.INVALID_CONTRACT) if not chunk: continue total += len(chunk) if total > self._max_bytes or total > expected_size: raise LocalStorageError(StorageErrorCode.TOO_LARGE) digest.update(chunk) await asyncio.to_thread(_write_all, descriptor, chunk) if total != expected_size: raise LocalStorageError(StorageErrorCode.SIZE_MISMATCH) actual_sha256 = digest.hexdigest() if actual_sha256 != expected_sha256: raise LocalStorageError(StorageErrorCode.HASH_MISMATCH) await asyncio.to_thread(os.fsync, descriptor) await asyncio.to_thread(os.close, descriptor) descriptor = None await asyncio.to_thread(os.chmod, temporary, _FINAL_MODE, follow_symlinks=False) await asyncio.to_thread(os.replace, temporary, target) installed = True await asyncio.to_thread(_fsync_directory, directory) return StoredUpload(storage_key, total, actual_sha256) except LocalStorageError: raise except OSError: raise LocalStorageError(StorageErrorCode.IO_UNAVAILABLE) from None finally: if descriptor is not None: try: os.close(descriptor) except OSError: pass if not installed: try: temporary.unlink(missing_ok=True) except OSError: pass async def remove(self, storage_key: uuid.UUID) -> None: """Remove an unreferenced deduplicated upload without exposing its path.""" try: directory, target = await asyncio.to_thread(self._prepare_target, storage_key) await asyncio.to_thread(target.unlink, missing_ok=True) await asyncio.to_thread(_fsync_directory, directory) except LocalStorageError: raise except OSError: raise LocalStorageError(StorageErrorCode.IO_UNAVAILABLE) from None async def read_verified( self, *, storage_key: uuid.UUID, expected_size: int, expected_sha256: str, ) -> bytes: """Read one UUID object while rechecking its immutable size and digest.""" self._validate_expectation(expected_size, expected_sha256) try: _, target = await asyncio.to_thread(self._prepare_target, storage_key) return await asyncio.to_thread( self._read_verified_object, target, expected_size, expected_sha256, ) except LocalStorageError: raise except OSError: raise LocalStorageError(StorageErrorCode.IO_UNAVAILABLE) from None def _validate_expectation(self, expected_size: int, expected_sha256: str) -> None: if ( isinstance(expected_size, bool) or not 1 <= expected_size <= self._max_bytes or _SHA256_PATTERN.fullmatch(expected_sha256) is None ): raise LocalStorageError(StorageErrorCode.INVALID_CONTRACT) def _prepare_target(self, storage_key: uuid.UUID) -> tuple[Path, Path]: if not isinstance(storage_key, uuid.UUID): raise LocalStorageError(StorageErrorCode.INVALID_CONTRACT) if self._root.exists() and self._root.is_symlink(): raise LocalStorageError(StorageErrorCode.ROOT_UNSAFE) self._root.mkdir(mode=0o750, parents=True, exist_ok=True) if self._root.is_symlink() or not self._root.is_dir(): raise LocalStorageError(StorageErrorCode.ROOT_UNSAFE) directory = self._root / storage_key.hex[:2] if directory.exists() and directory.is_symlink(): raise LocalStorageError(StorageErrorCode.ROOT_UNSAFE) directory.mkdir(mode=0o750, exist_ok=True) if directory.is_symlink() or not directory.is_dir(): raise LocalStorageError(StorageErrorCode.ROOT_UNSAFE) target = directory / storage_key.hex if target.parent != directory or directory.parent != self._root: raise LocalStorageError(StorageErrorCode.ROOT_UNSAFE) return directory, target @staticmethod def _open_temporary(path: Path) -> int: flags = os.O_CREAT | os.O_EXCL | os.O_WRONLY if hasattr(os, "O_NOFOLLOW"): flags |= os.O_NOFOLLOW return os.open(path, flags, _WRITE_MODE) @staticmethod def _existing_object( target: Path, storage_key: uuid.UUID, expected_size: int, expected_sha256: str, ) -> StoredUpload | None: try: metadata = target.lstat() except FileNotFoundError: return None if stat.S_ISLNK(metadata.st_mode) or not stat.S_ISREG(metadata.st_mode): raise LocalStorageError(StorageErrorCode.ROOT_UNSAFE) if metadata.st_size != expected_size: raise LocalStorageError(StorageErrorCode.OBJECT_CONFLICT) digest = hashlib.sha256() try: with target.open("rb") as source: for chunk in iter(lambda: source.read(1024 * 1024), b""): digest.update(chunk) except OSError: raise LocalStorageError(StorageErrorCode.IO_UNAVAILABLE) from None if digest.hexdigest() != expected_sha256: raise LocalStorageError(StorageErrorCode.OBJECT_CONFLICT) return StoredUpload(storage_key, expected_size, expected_sha256) @staticmethod def _read_verified_object( target: Path, expected_size: int, expected_sha256: str, ) -> bytes: flags = os.O_RDONLY if hasattr(os, "O_NOFOLLOW"): flags |= os.O_NOFOLLOW try: descriptor = os.open(target, flags) except OSError: raise LocalStorageError(StorageErrorCode.IO_UNAVAILABLE) from None try: metadata = os.fstat(descriptor) if not stat.S_ISREG(metadata.st_mode): raise LocalStorageError(StorageErrorCode.ROOT_UNSAFE) if metadata.st_size != expected_size: raise LocalStorageError(StorageErrorCode.OBJECT_CONFLICT) value = bytearray() digest = hashlib.sha256() while len(value) < expected_size: chunk = os.read(descriptor, min(1024 * 1024, expected_size - len(value))) if not chunk: break value.extend(chunk) digest.update(chunk) if len(value) != expected_size or os.read(descriptor, 1): raise LocalStorageError(StorageErrorCode.OBJECT_CONFLICT) if digest.hexdigest() != expected_sha256: raise LocalStorageError(StorageErrorCode.OBJECT_CONFLICT) return bytes(value) finally: os.close(descriptor) def _write_all(descriptor: int, value: bytes) -> None: view = memoryview(value) while view: written = os.write(descriptor, view) if written <= 0: raise OSError("bounded upload write made no progress") view = view[written:] def _fsync_directory(directory: Path) -> None: flags = os.O_RDONLY if hasattr(os, "O_DIRECTORY"): flags |= os.O_DIRECTORY descriptor = os.open(directory, flags) try: os.fsync(descriptor) finally: os.close(descriptor)