2026-08-06 22:01:42 +08:00

593 lines
22 KiB
Python

"""Shared upload management logic.
Pure business logic — no FastAPI/HTTP dependencies.
Both Gateway and Client delegate to these functions.
"""
import errno
import logging
import os
import secrets
import shutil
import stat
from collections.abc import Iterator
from dataclasses import dataclass
from pathlib import Path
from typing import BinaryIO
from deerflow.config.paths import get_paths
from deerflow.runtime.user_context import get_effective_user_id
from deerflow.uploads.errors import AtomicUploadPublishError, PathTraversalError, UnsafeUploadPathError
from deerflow.uploads.layout import (
UPLOAD_CONVERSIONS_DIRNAME,
_truncate_utf8,
artifact_url_for_virtual_path,
existing_conversion_path_for_upload,
upload_virtual_path,
)
from deerflow.uploads.lease import UploadIdentity, UploadNameLease, UploadStageLease
from deerflow.utils.thread_id import validate_thread_id
logger = logging.getLogger(__name__)
UPLOAD_STAGING_PREFIX = ".upload-"
UPLOAD_STAGING_SUFFIX = ".part"
@dataclass(slots=True)
class StagedUpload:
"""A complete-or-in-progress upload stored under a hidden temporary name."""
base_dir: Path
path: Path
handle: BinaryIO
lease: UploadStageLease
@dataclass(slots=True)
class PublishedUpload:
"""A published upload whose actual filename remains exclusively leased."""
path: Path
identity: UploadIdentity
lease: UploadNameLease
@property
def is_active(self) -> bool:
return self.lease.is_active
def release(self) -> None:
"""Release the publication's name lease."""
self.lease.release()
def get_uploads_dir(thread_id: str, *, user_id: str | None = None) -> Path:
"""Return the uploads directory path for a thread (no side effects)."""
validate_thread_id(thread_id)
return get_paths().sandbox_uploads_dir(thread_id, user_id=user_id or get_effective_user_id())
def ensure_uploads_dir(thread_id: str, *, user_id: str | None = None) -> Path:
"""Return the uploads directory for a thread, creating it if needed."""
base = get_uploads_dir(thread_id, user_id=user_id)
base.mkdir(parents=True, exist_ok=True)
return base
def normalize_filename(filename: str) -> str:
"""Sanitize a filename by extracting its basename.
Strips any directory components and rejects traversal patterns.
Args:
filename: Raw filename from user input (may contain path components).
Returns:
Safe filename (basename only).
Raises:
ValueError: If filename is empty or resolves to a traversal pattern.
"""
if not filename:
raise ValueError("Filename is empty")
if "\0" in filename:
raise ValueError(f"Filename contains NUL byte: {filename!r}")
safe = Path(filename).name
if not safe or safe in {".", ".."}:
raise ValueError(f"Filename is unsafe: {filename!r}")
# Reject backslashes — on Linux Path.name keeps them as literal chars,
# but they indicate a Windows-style path that should be stripped or rejected.
if "\\" in safe:
raise ValueError(f"Filename contains backslash: {filename!r}")
if "<" in safe or ">" in safe or "--- BEGIN USER INPUT ---" in safe or "--- END USER INPUT ---" in safe:
raise ValueError(f"Filename contains reserved model-context token: {filename!r}")
if len(safe.encode("utf-8")) > 255:
raise ValueError(f"Filename too long: {len(safe)} chars")
if is_upload_staging_file(safe):
raise ValueError(f"Filename uses reserved upload staging pattern: {filename!r}")
return safe
def claim_unique_filename(name: str, seen: set[str]) -> str:
"""Generate a unique filename by appending ``_N`` suffix on collision.
Automatically adds the returned name to *seen* so callers don't need to.
Args:
name: Candidate filename.
seen: Set of filenames already claimed (mutated in place).
Returns:
A filename not present in *seen* (already added to *seen*).
"""
if name not in seen:
seen.add(name)
return name
stem, suffix = Path(name).stem, Path(name).suffix
counter = 1
candidate = f"{stem}_{counter}{suffix}"
while candidate in seen:
counter += 1
candidate = f"{stem}_{counter}{suffix}"
seen.add(candidate)
return candidate
def is_upload_staging_file(filename: str) -> bool:
"""Return whether *filename* is a transient Gateway upload staging file."""
return filename.startswith(UPLOAD_STAGING_PREFIX) and filename.endswith(UPLOAD_STAGING_SUFFIX)
def _validate_upload_directory(base_dir: Path) -> Path:
"""Return a real upload directory without following a directory symlink."""
try:
st = os.lstat(base_dir)
except FileNotFoundError as exc:
raise UnsafeUploadPathError(f"Upload directory does not exist: {base_dir}") from exc
if stat.S_ISLNK(st.st_mode) or not stat.S_ISDIR(st.st_mode):
raise UnsafeUploadPathError(f"Upload directory is unsafe: {base_dir}")
return base_dir
def create_upload_staging_file(base_dir: Path) -> StagedUpload:
"""Create a hidden same-directory staging file for a complete payload."""
base_dir = _validate_upload_directory(Path(base_dir))
while True:
temp_path = base_dir / f"{UPLOAD_STAGING_PREFIX}{secrets.token_hex(16)}{UPLOAD_STAGING_SUFFIX}"
lease = UploadStageLease.acquire(base_dir, temp_path.name)
flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL
if hasattr(os, "O_BINARY"):
flags |= os.O_BINARY
try:
fd = os.open(temp_path, flags, 0o600)
except FileExistsError:
lease.release()
continue
except BaseException:
lease.release()
raise
try:
handle = os.fdopen(fd, "wb")
except Exception:
os.close(fd)
temp_path.unlink(missing_ok=True)
lease.release()
raise
return StagedUpload(base_dir=base_dir, path=temp_path, handle=handle, lease=lease)
def abort_staged_upload(staged: StagedUpload) -> None:
"""Close and remove a staging file, tolerating repeated cleanup."""
close_error: BaseException | None = None
unlink_error: BaseException | None = None
lease_error: BaseException | None = None
try:
if not staged.handle.closed:
staged.handle.close()
except BaseException as exc:
close_error = exc
try:
staged.path.unlink(missing_ok=True)
except BaseException as exc:
unlink_error = exc
try:
staged.lease.release()
except BaseException as exc:
lease_error = exc
if close_error is not None:
if unlink_error is not None:
raise close_error from unlink_error
if lease_error is not None:
raise close_error from lease_error
raise close_error
if unlink_error is not None:
if lease_error is not None:
raise unlink_error from lease_error
raise unlink_error
if lease_error is not None:
raise lease_error
def _validate_staged_upload(staged: StagedUpload) -> None:
"""Reject a staging path that was replaced or moved outside its directory."""
_validate_upload_directory(staged.base_dir)
if not staged.lease.is_active or staged.lease.stage_dir != staged.base_dir or staged.lease.stage_filename != staged.path.name:
raise UnsafeUploadPathError("Upload staging liveness lease is missing")
if staged.path.parent.resolve() != staged.base_dir.resolve():
raise UnsafeUploadPathError("Upload staging path escaped its directory")
try:
staged_stat = os.lstat(staged.path)
except FileNotFoundError as exc:
raise UnsafeUploadPathError("Upload staging file disappeared") from exc
if not stat.S_ISREG(staged_stat.st_mode) or staged_stat.st_nlink != 1:
raise UnsafeUploadPathError("Upload staging path is not an exclusive regular file")
def _filename_candidates(name: str) -> Iterator[str]:
"""Yield collision candidates that stay within the filename byte limit."""
yield name
stem, suffix = Path(name).stem, Path(name).suffix
counter = 1
while True:
marker = f"_{counter}"
max_stem_bytes = 255 - len(marker.encode("utf-8")) - len(suffix.encode("utf-8"))
if max_stem_bytes < 1:
raise AtomicUploadPublishError("Filename suffix leaves no room for collision marker")
candidate_stem = _truncate_utf8(stem, max_stem_bytes)
if not candidate_stem:
raise AtomicUploadPublishError("Filename stem leaves no room for collision marker")
yield f"{candidate_stem}{marker}{suffix}"
counter += 1
def _unlink_matching_upload(path: Path, identity: UploadIdentity) -> None:
"""Remove *path* only while it still names *identity*."""
if identity.matches(path):
path.unlink(missing_ok=True)
def _rollback_link_without_masking(path: Path, identity: UploadIdentity) -> None:
try:
_unlink_matching_upload(path, identity)
except BaseException:
logger.warning("Failed to roll back partially published upload: %s", path, exc_info=True)
def _release_lease_without_masking(lease: UploadNameLease) -> None:
try:
lease.release()
except BaseException:
logger.warning("Failed to release upload name lease: %s", lease.lock_path, exc_info=True)
def publish_staged_upload_leased(staged: StagedUpload, preferred_filename: str) -> PublishedUpload:
"""Atomically publish a staging file and retain its actual-name lease."""
safe_name = normalize_filename(preferred_filename)
if not staged.handle.closed:
staged.handle.close()
_validate_staged_upload(staged)
staged_identity = UploadIdentity.from_path(staged.path)
for candidate_name in _filename_candidates(safe_name):
candidate = staged.base_dir / candidate_name
try:
os.lstat(candidate)
except FileNotFoundError:
pass
else:
# A visible entry cannot be published to, so skip it before taking
# its lease. This also lets one multi-file request retain the first
# generation's lease while choosing a suffix for a duplicate name.
continue
lease = UploadNameLease.acquire(staged.base_dir, candidate_name)
linked = False
try:
try:
os.link(staged.path, candidate, follow_symlinks=False)
linked = True
except FileExistsError:
lease.release()
continue
except (NotImplementedError, TypeError) as exc:
raise AtomicUploadPublishError("Storage does not support atomic no-replace publication") from exc
except OSError as exc:
if exc.errno == errno.EEXIST:
lease.release()
continue
raise AtomicUploadPublishError(f"Storage does not support atomic no-replace publication: {exc}") from exc
if not staged_identity.matches(candidate):
raise AtomicUploadPublishError("Published upload identity changed during publication")
try:
staged.path.unlink()
except OSError as exc:
_rollback_link_without_masking(candidate, staged_identity)
raise AtomicUploadPublishError("Failed to remove upload staging link after publication") from exc
staged.lease.release()
try:
candidate_stat = os.lstat(candidate)
except FileNotFoundError as exc:
raise AtomicUploadPublishError("Published upload disappeared") from exc
if not stat.S_ISREG(candidate_stat.st_mode) or candidate_stat.st_nlink != 1 or not staged_identity.matches(candidate):
_rollback_link_without_masking(candidate, staged_identity)
raise AtomicUploadPublishError("Published upload did not become an exclusive regular file")
return PublishedUpload(path=candidate, identity=staged_identity, lease=lease)
except BaseException:
if linked:
_rollback_link_without_masking(candidate, staged_identity)
_release_lease_without_masking(lease)
raise
def publish_staged_upload(staged: StagedUpload, preferred_filename: str) -> Path:
"""Atomically publish a complete staging file without replacing an entry."""
publication = publish_staged_upload_leased(staged, preferred_filename)
try:
return publication.path
finally:
publication.release()
def _abort_staged_upload_without_masking(staged: StagedUpload) -> None:
try:
abort_staged_upload(staged)
except BaseException:
logger.warning("Failed to clean up upload staging file: %s", staged.path, exc_info=True)
def publish_upload_bytes_leased(base_dir: Path, preferred_filename: str, data: bytes) -> PublishedUpload:
"""Stage bytes, publish them atomically, and retain the actual-name lease."""
safe_name = normalize_filename(preferred_filename)
staged = create_upload_staging_file(base_dir)
try:
staged.handle.write(data)
return publish_staged_upload_leased(staged, safe_name)
except BaseException:
_abort_staged_upload_without_masking(staged)
raise
def publish_upload_bytes(base_dir: Path, preferred_filename: str, data: bytes) -> Path:
"""Stage and atomically publish an in-memory upload payload."""
publication = publish_upload_bytes_leased(base_dir, preferred_filename, data)
try:
return publication.path
finally:
publication.release()
def publish_upload_copy_leased(base_dir: Path, preferred_filename: str, source_path: Path) -> PublishedUpload:
"""Copy a source into staging, publish it, and retain the actual-name lease."""
safe_name = normalize_filename(preferred_filename)
staged = create_upload_staging_file(base_dir)
try:
with Path(source_path).open("rb") as source:
shutil.copyfileobj(source, staged.handle)
return publish_staged_upload_leased(staged, safe_name)
except BaseException:
_abort_staged_upload_without_masking(staged)
raise
def publish_upload_copy(base_dir: Path, preferred_filename: str, source_path: Path) -> Path:
"""Copy a local source into staging and atomically publish it."""
publication = publish_upload_copy_leased(base_dir, preferred_filename, source_path)
try:
return publication.path
finally:
publication.release()
def make_upload_file_sandbox_readable(file_path: Path) -> None:
"""Add group/other read bits to one verified regular upload artifact."""
file_path = Path(file_path)
file_stat = os.lstat(file_path)
if not stat.S_ISREG(file_stat.st_mode) or file_stat.st_nlink != 1:
raise UnsafeUploadPathError(f"Unsafe upload file: {file_path.name}")
readable_mode = stat.S_IMODE(file_stat.st_mode) | stat.S_IRGRP | stat.S_IROTH
chmod_kwargs = {"follow_symlinks": False} if os.chmod in os.supports_follow_symlinks else {}
os.chmod(file_path, readable_mode, **chmod_kwargs)
def rollback_published_upload(publication: PublishedUpload) -> None:
"""Remove only the still-leased upload generation represented by *publication*."""
if not publication.is_active:
raise RuntimeError("Cannot roll back a publication after releasing its lease")
if publication.lease.filename != publication.path.name:
raise UnsafeUploadPathError("Publication lease does not match its upload path")
if not publication.identity.matches(publication.path):
return
owned_conversion = existing_conversion_path_for_upload(publication.path)
if owned_conversion is not None:
owned_conversion.unlink(missing_ok=True)
publication.path.unlink()
def replace_system_owned_staged_file(staged: StagedUpload, filename: str) -> Path:
"""Atomically replace one generated file inside the owned namespace."""
if staged.base_dir.name != UPLOAD_CONVERSIONS_DIRNAME:
raise UnsafeUploadPathError("System-owned replace requires conversion directory")
if not staged.handle.closed:
staged.handle.close()
_validate_staged_upload(staged)
target = staged.base_dir / normalize_filename(filename)
os.replace(staged.path, target)
staged.lease.release()
return target
def validate_path_traversal(path: Path, base: Path) -> None:
"""Verify that *path* is inside *base*.
Raises:
PathTraversalError: If a path traversal is detected.
"""
try:
path.resolve().relative_to(base.resolve())
except ValueError:
raise PathTraversalError("Path traversal detected") from None
def validate_upload_destination(base_dir: Path, filename: str) -> Path:
"""Validate an upload destination without mutating an existing file."""
safe_name = normalize_filename(filename)
dest = base_dir / safe_name
try:
st = os.lstat(dest)
except FileNotFoundError:
st = None
if st is not None and not stat.S_ISREG(st.st_mode):
raise UnsafeUploadPathError(f"Upload destination is not a regular file: {safe_name}")
if st is not None and st.st_nlink > 1:
raise UnsafeUploadPathError(f"Upload destination has multiple links: {safe_name}")
validate_path_traversal(dest, base_dir)
return dest
def _iter_upload_storage_dirs(base_dir: Path):
for user_data_dir in base_dir.glob("threads/*/user-data"):
yield user_data_dir / "uploads"
yield user_data_dir / UPLOAD_CONVERSIONS_DIRNAME
for user_data_dir in base_dir.glob("users/*/threads/*/user-data"):
yield user_data_dir / "uploads"
yield user_data_dir / UPLOAD_CONVERSIONS_DIRNAME
def cleanup_stale_upload_staging_files(base_dir: Path | str | None = None) -> int:
"""Remove orphaned Gateway upload staging files left by a hard crash."""
root = Path(base_dir) if base_dir is not None else get_paths().base_dir
removed = 0
for uploads_dir in _iter_upload_storage_dirs(root):
try:
directory_stat = os.lstat(uploads_dir)
except FileNotFoundError:
continue
if stat.S_ISLNK(directory_stat.st_mode) or not stat.S_ISDIR(directory_stat.st_mode):
continue
try:
with os.scandir(uploads_dir) as entries:
for entry in entries:
if not is_upload_staging_file(entry.name) or not entry.is_file(follow_symlinks=False):
continue
stage_lease = UploadStageLease.try_acquire(uploads_dir, entry.name)
if stage_lease is None:
continue
try:
os.unlink(entry.path)
removed += 1
except FileNotFoundError:
pass
except OSError:
logger.warning("Failed to remove stale upload staging file: %s", entry.path, exc_info=True)
finally:
stage_lease.release()
except FileNotFoundError:
continue
except OSError:
logger.warning("Failed to scan uploads directory for stale staging files: %s", uploads_dir, exc_info=True)
return removed
def write_upload_file_no_symlink(base_dir: Path, filename: str, data: bytes) -> Path:
"""Compatibility wrapper for collision-safe upload publication."""
return publish_upload_bytes(base_dir, filename, data)
def list_files_in_dir(directory: Path) -> dict:
"""List files (not directories) in *directory*.
Args:
directory: Directory to scan.
Returns:
Dict with "files" list (sorted by name) and "count".
Each file entry has ``size`` as *int* (bytes). Call
:func:`enrich_file_listing` to add virtual / artifact URLs.
"""
if not directory.is_dir():
return {"files": [], "count": 0}
files = []
with os.scandir(directory) as entries:
for entry in sorted(entries, key=lambda e: e.name):
if is_upload_staging_file(entry.name):
continue
if not entry.is_file(follow_symlinks=False):
continue
st = entry.stat(follow_symlinks=False)
files.append(
{
"filename": entry.name,
"size": st.st_size,
"path": entry.path,
"extension": Path(entry.name).suffix,
"modified": st.st_mtime,
}
)
return {"files": files, "count": len(files)}
def delete_file_safe(base_dir: Path, filename: str) -> dict:
"""Delete a primary upload and only its exact owned conversion.
Args:
base_dir: Directory containing the file.
filename: Name of file to delete.
Returns:
Dict with success and message.
Raises:
FileNotFoundError: If the file does not exist.
PathTraversalError: If path traversal is detected.
"""
safe_name = normalize_filename(filename)
if safe_name != filename:
raise PathTraversalError("Path traversal detected")
try:
base_dir = _validate_upload_directory(Path(base_dir))
except UnsafeUploadPathError as exc:
if isinstance(exc.__cause__, FileNotFoundError):
raise FileNotFoundError(f"File not found: {filename}") from exc
raise
lease = UploadNameLease.acquire(base_dir, safe_name)
try:
file_path = base_dir / safe_name
try:
file_stat = os.lstat(file_path)
except FileNotFoundError:
raise FileNotFoundError(f"File not found: {filename}") from None
if not stat.S_ISREG(file_stat.st_mode) or file_stat.st_nlink != 1:
raise UnsafeUploadPathError(f"Unsafe upload file: {safe_name}")
owned_conversion = existing_conversion_path_for_upload(file_path)
if owned_conversion is not None:
owned_conversion.unlink(missing_ok=True)
file_path.unlink()
finally:
lease.release()
return {"success": True, "message": f"Deleted {filename}"}
def upload_artifact_url(thread_id: str, filename: str) -> str:
"""Build the artifact URL for a file in a thread's uploads directory.
*filename* is percent-encoded so that spaces, ``#``, ``?`` etc. are safe.
"""
return artifact_url_for_virtual_path(thread_id, upload_virtual_path(filename))
def enrich_file_listing(result: dict, thread_id: str) -> dict:
"""Add virtual paths and artifact URLs on a listing result.
Mutates *result* in place and returns it for convenience.
"""
for f in result["files"]:
filename = f["filename"]
f["virtual_path"] = upload_virtual_path(filename)
f["artifact_url"] = upload_artifact_url(thread_id, filename)
return result