RAG/backend/app/adapters/local_storage.py

"""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)