"""Artifact validation, private staging, session binding, and downloads.""" from __future__ import annotations import base64 import threading from functools import wraps import os import posixpath import shutil import time import unicodedata import uuid from pathlib import Path from typing import Any, Dict, List, Optional, Tuple, Set from backend.errors import InvalidInputError from backend.transport import ID_REGEX MAX_FILES_PER_QUERY = 20 MAX_BYTES_PER_FILE = 10 * 1024 * 1024 # 10 MiB MAX_BYTES_PER_QUERY = 50 * 1024 * 1024 # 50 MiB GLOBAL_STORAGE_LIMIT = 500 * 1024 * 1024 # 500 MiB MAX_CHUNK_BYTES = 64 * 1024 # 64 KiB decoded DEFAULT_TTL_SECONDS = 900 # 15 minutes def validate_artifact_name(name: str) -> str: """Validate and normalize an artifact display name. Must be NFC normalized, POSIX relative path components, no backslash, no control characters, no dot traversal, maximum 1,024 UTF-8 bytes. """ if not isinstance(name, str) or not name.strip(): raise InvalidInputError("Artifact name must be a non-empty string") normalized = unicodedata.normalize("NFC", name.strip()) encoded = normalized.encode("utf-8") if len(encoded) > 1024: raise InvalidInputError("Artifact name exceeds 1,024 bytes") if "\\" in normalized: raise InvalidInputError("Artifact name must not contain backslashes") # Reject control characters for ch in normalized: if unicodedata.category(ch) == "Cc": raise InvalidInputError("Artifact name must not contain control characters") # Reject absolute path if normalized.startswith("/"): raise InvalidInputError("Artifact name must be a relative path") # Reject dot traversal segments = normalized.split("/") if ".." in segments or "." in segments or "" in segments: raise InvalidInputError("Artifact name contains directory traversal or empty segments") clean = posixpath.normpath(normalized) if clean.startswith("..") or clean == "." or clean != normalized: raise InvalidInputError("Artifact name contains invalid path segments") return normalized def synchronized(method): @wraps(method) def wrapped(self, *args, **kwargs): store = self if isinstance(self, ArtifactStore) else self.store with store._lock: return method(self, *args, **kwargs) return wrapped class StagedArtifact: def __init__(self, transfer_id: str, name: str, declared_size: int, staging_path: Path): self.transfer_id = transfer_id self.name = name self.declared_size = declared_size self.staging_path = staging_path self.written_bytes = 0 self.expected_chunk_index = 0 self.chunks_received = 0 # Exclusive file creation with restrictive 0600 permissions fd = os.open(staging_path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600) self.file_handle = os.fdopen(fd, "wb") def write_chunk(self, index: int, data: bytes) -> None: if index != self.expected_chunk_index: raise InvalidInputError( f"Chunk index mismatch: expected {self.expected_chunk_index}, got {index}" ) if len(data) > MAX_CHUNK_BYTES: raise InvalidInputError( f"Chunk size ({len(data)} bytes) exceeds 64 KiB limit" ) if self.written_bytes + len(data) > self.declared_size: raise InvalidInputError("Chunk data exceeds declared artifact size") self.file_handle.write(data) self.written_bytes += len(data) self.expected_chunk_index += 1 self.chunks_received += 1 def close(self, expected_size: int, expected_chunks: int) -> None: self.file_handle.flush() self.file_handle.close() if self.written_bytes != expected_size or self.written_bytes != self.declared_size: raise InvalidInputError( f"Artifact size mismatch: declared {self.declared_size}, actual {self.written_bytes}" ) if self.chunks_received != expected_chunks: raise InvalidInputError( f"Chunk count mismatch: expected {expected_chunks}, got {self.chunks_received}" ) class CommittedArtifact: def __init__( self, artifact_id: str, name: str, file_path: Path, size_bytes: int, expires_at_ts: float, session_id: str, ): self.artifact_id = artifact_id self.name = name self.file_path = file_path self.size_bytes = size_bytes self.expires_at_ts = expires_at_ts self.session_id = session_id self.active_readers = 0 self.revoked = False class ArtifactStore: """Manages private temporary storage, global capacity, retention, and downloads.""" def __init__(self, root_dir: Path): self._lock = threading.RLock() self.root_dir = root_dir.resolve() self.staging_dir = self.root_dir / "staging" self.committed_dir = self.root_dir / "committed" # Initialize directories with 0700 permissions self._ensure_dirs() self.global_reserved_bytes = 0 self._committed: Dict[str, CommittedArtifact] = {} def _ensure_dirs(self) -> None: self.root_dir.mkdir(parents=True, exist_ok=True) os.chmod(self.root_dir, 0o700) self.staging_dir.mkdir(exist_ok=True) os.chmod(self.staging_dir, 0o700) self.committed_dir.mkdir(exist_ok=True) os.chmod(self.committed_dir, 0o700) @synchronized def purge_all(self) -> None: """Delete prior-run stored files on server restart.""" for d in (self.staging_dir, self.committed_dir): if d.exists(): for item in d.iterdir(): if item.is_file(): item.unlink(missing_ok=True) elif item.is_dir(): shutil.rmtree(item, ignore_errors=True) self._committed.clear() self.global_reserved_bytes = 0 @synchronized def expire_artifacts(self) -> None: """Remove expired artifacts whose active reader count is zero.""" now = time.time() to_delete = [] for aid, art in self._committed.items(): if now >= art.expires_at_ts and art.active_readers == 0: to_delete.append(aid) for aid in to_delete: art = self._committed[aid] art.file_path.unlink(missing_ok=True) self._committed.pop(aid) self.global_reserved_bytes = max(0, self.global_reserved_bytes - art.size_bytes) @synchronized def release_reader(self, artifact_id: str) -> None: """Decrement active reader count for an artifact.""" art = self._committed.get(artifact_id) if art: art.active_readers = max(0, art.active_readers - 1) @synchronized def get_artifact_for_download( self, artifact_id: str, session_id: str ) -> Optional[Tuple[Path, str, int]]: """Retrieve artifact path and sanitized name if session matches and not expired. Increments active_readers. Caller must invoke release_reader(artifact_id) when finished. Returns (file_path, display_name, size_bytes) or None. """ self.expire_artifacts() art = self._committed.get(artifact_id) if not art or art.revoked: return None # Check session ownership if art.session_id != session_id: return None # Check expiry if time.time() >= art.expires_at_ts: return None # Track active reader to prevent race condition during download art.active_readers += 1 return art.file_path, art.name, art.size_bytes @synchronized def create_staging_session(self, query_id: str) -> QueryArtifactStaging: query_staging_dir = self.staging_dir / query_id if not ID_REGEX.fullmatch(query_id): raise InvalidInputError("Invalid query ID") query_staging_dir.mkdir(mode=0o700) os.chmod(query_staging_dir, 0o700) return QueryArtifactStaging(self, query_id, query_staging_dir) class QueryArtifactStaging: """Handles sequential artifact transfers during query collection.""" def __init__(self, store: ArtifactStore, query_id: str, staging_dir: Path): self.store = store self.query_id = query_id self.staging_dir = staging_dir self.accepted_count = 0 self.query_reserved_bytes = 0 self.open_transfer: Optional[StagedArtifact] = None self.staged_artifacts: Dict[str, StagedArtifact] = {} self.seen_transfer_ids: Set[str] = set() self.staged_names: Set[str] = set() self.committed_artifact_ids: List[str] = [] self._disposed = False @synchronized def handle_begin( self, transfer_id: str, name: str, size_bytes: int ) -> Tuple[str, Optional[Dict[str, Any]]]: """Handle artifact_begin message. Returns (decision, warning_dict).""" if self.open_transfer is not None: raise InvalidInputError("Another artifact transfer is already open") if not isinstance(transfer_id, str) or not ID_REGEX.fullmatch(transfer_id): raise InvalidInputError("Invalid transfer ID") if type(size_bytes) is not int: raise InvalidInputError("Artifact size must be an integer") # Reject duplicate transfer IDs if transfer_id in self.seen_transfer_ids: raise InvalidInputError(f"Duplicate transfer ID: '{transfer_id}'") self.seen_transfer_ids.add(transfer_id) try: norm_name = validate_artifact_name(name) except InvalidInputError as e: return "skip", {"code": "artifact_invalid_name", "message": str(e), "name": name} # Reject duplicate normalized display names across transfers in the same query if norm_name in self.staged_names: return "skip", { "code": "artifact_duplicate_name", "message": f"Artifact name '{norm_name}' was already exported in this query", "name": norm_name, } if size_bytes < 0: return "skip", {"code": "artifact_invalid_size", "message": "Size cannot be negative", "name": norm_name} # Check per-query count if self.accepted_count >= MAX_FILES_PER_QUERY: return "skip", { "code": "artifact_count_exceeded", "message": f"Query artifact count limit ({MAX_FILES_PER_QUERY}) reached", "name": norm_name, } # Check per-file size if size_bytes > MAX_BYTES_PER_FILE: return "skip", { "code": "artifact_size_exceeded", "message": f"Artifact size {size_bytes} exceeds 10 MiB limit", "name": norm_name, } # Check per-query storage if self.query_reserved_bytes + size_bytes > MAX_BYTES_PER_QUERY: return "skip", { "code": "artifact_query_storage_exceeded", "message": "Query artifact storage limit (50 MiB) exceeded", "name": norm_name, } # Check global storage if self.store.global_reserved_bytes + size_bytes > GLOBAL_STORAGE_LIMIT: return "skip", { "code": "artifact_global_storage_exceeded", "message": "Global artifact storage limit (500 MiB) exceeded", "name": norm_name, } # Accept transfer self.query_reserved_bytes += size_bytes self.store.global_reserved_bytes += size_bytes self.accepted_count += 1 self.seen_transfer_ids.add(transfer_id) self.staged_names.add(norm_name) file_path = self.staging_dir / f"{transfer_id}.tmp" self.open_transfer = StagedArtifact( transfer_id=transfer_id, name=norm_name, declared_size=size_bytes, staging_path=file_path, ) return "accept", None @synchronized def handle_chunk(self, transfer_id: str, index: int, data_base64: str) -> None: """Handle artifact_chunk message.""" if self.open_transfer is None or self.open_transfer.transfer_id != transfer_id: raise InvalidInputError(f"No open transfer matching ID '{transfer_id}'") try: raw_bytes = base64.b64decode(data_base64, validate=True) except Exception as e: raise InvalidInputError("Invalid base64 in artifact chunk") from e self.open_transfer.write_chunk(index, raw_bytes) @synchronized def handle_end( self, transfer_id: str, size_bytes: int, chunks: int ) -> Tuple[str, Optional[Dict[str, Any]]]: """Handle artifact_end message. Returns ('stored', None).""" if self.open_transfer is None or self.open_transfer.transfer_id != transfer_id: raise InvalidInputError(f"No open transfer matching ID '{transfer_id}'") staged = self.open_transfer staged.close(size_bytes, chunks) self.staged_artifacts[transfer_id] = staged self.open_transfer = None return "stored", None @synchronized def commit(self, session_id: str, ttl_seconds: int = DEFAULT_TTL_SECONDS) -> List[Dict[str, Any]]: """Commit stored artifacts into the store. Returns metadata list for HTTP response.""" if self.open_transfer is not None: raise InvalidInputError("Cannot commit while an artifact transfer is open") now = time.time() expires_at_ts = now + ttl_seconds expires_at_str = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime(expires_at_ts)) results = [] try: for transfer_id, staged in self.staged_artifacts.items(): artifact_id = uuid.uuid4().hex dest_path = self.store.committed_dir / f"{artifact_id}.dat" committed = CommittedArtifact( artifact_id=artifact_id, name=staged.name, file_path=dest_path, size_bytes=staged.written_bytes, expires_at_ts=expires_at_ts, session_id=session_id, ) self.store._committed[artifact_id] = committed self.committed_artifact_ids.append(artifact_id) shutil.move(str(staged.staging_path), str(dest_path)) results.append({ "id": artifact_id, "name": staged.name, "size_bytes": staged.written_bytes, "expires_at": expires_at_str, }) except Exception: # On commit failure, purge partially committed artifacts and restore reservations self.purge_committed_and_discard() raise # Cleanup staging directory self.discard(release_committed=False) return results @synchronized def purge_committed_and_discard(self) -> None: """Purge all committed artifacts created by this query and discard staging.""" failure = None for aid in list(self.committed_artifact_ids): art = self.store._committed.get(aid) if art: art.revoked = True art.expires_at_ts = 0 if not self._disposed: self.query_reserved_bytes -= art.size_bytes try: art.file_path.unlink(missing_ok=True) self.store._committed.pop(aid) self.store.global_reserved_bytes -= art.size_bytes except OSError as exc: failure = exc self.committed_artifact_ids.clear() self.discard(release_committed=True) if failure: raise failure @synchronized def discard(self, release_committed: bool = True) -> None: """Discard all staged artifacts and release reservations.""" if self._disposed: return if self.open_transfer is not None: try: self.open_transfer.file_handle.close() except Exception: pass self.open_transfer = None # Remove query staging directory if self.staging_dir.exists(): shutil.rmtree(self.staging_dir) self._disposed = True if release_committed: # Release query reservations from global self.store.global_reserved_bytes = max( 0, self.store.global_reserved_bytes - self.query_reserved_bytes ) self.query_reserved_bytes = 0 self.staged_artifacts.clear()