Artur Mukhamadiev
e65fbf4b67
FastAPI app, upstream Confluence/model adapters, authoritative history, rootless container lifecycle, artifact storage and downloads, fake peers under backend/dev, tests under tests/backend. Root pytest.ini deselects the live marker by default; requirements gain the backend dependencies.
445 lines
16 KiB
Python
445 lines
16 KiB
Python
"""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()
|