Files
saas_backend/app/core/document_storage.py
T
2026-08-31 20:39:41 -04:00

128 lines
4.1 KiB
Python

"""Where the bytes live.
A narrow interface over a local directory. There is no object-store credential in
this deployment, and inventing one would be inventing an infrastructure decision
that is not mine to make — but the surface here is four functions, so putting S3
behind it later is one class rather than a rewrite.
## The two rules that matter
**A storage key is generated, never derived.** Not from the filename, not from
the description, not from anything a person typed. Deriving a path from user
input is how `../../etc/passwd` gets written, how two people uploading
`report.pdf` overwrite each other's work, and how a URL becomes something worth
guessing at.
**Every path is re-checked against the root before it is opened.** The key is
generated here so it cannot escape, and it is checked anyway — because "this
value is safe because of where it came from" is an argument that survives exactly
until somebody adds a second caller.
"""
from __future__ import annotations
import hashlib
import logging
import secrets
from pathlib import Path
from typing import BinaryIO, Iterator
logger = logging.getLogger(__name__)
CHUNK_BYTES = 64 * 1024
def _root() -> Path:
from app.config.settings import settings
return Path(settings.DOCUMENT_STORAGE_PATH).resolve()
def new_key(tenant_id) -> str:
"""A fresh key for one document.
The workspace id is the first segment so a filesystem listing is navigable
and so a whole workspace's files can be removed with it; the rest is random,
which is what makes the key unguessable and collision-free.
"""
return f"{tenant_id}/{secrets.token_hex(16)}"
def _path_for(key: str) -> Path:
"""The absolute path for a key, or a refusal.
The key is generated by `new_key` and cannot contain a traversal — and this
checks anyway. "Safe because of where it came from" holds until the day
somebody adds a second caller, and this is a filesystem write.
"""
root = _root()
candidate = (root / key).resolve()
if not candidate.is_relative_to(root):
raise ValueError("storage key escapes the storage root")
return candidate
def write(key: str, source: BinaryIO, *, max_bytes: int) -> tuple[int, str]:
"""Stream a file to disk. Returns its size and SHA-256.
Refuses past `max_bytes` **while writing** rather than after: a limit checked
on `Content-Length` alone trusts a header, and one checked after the write
has already spent the disk.
A partial file is removed on the way out. Leaving one behind would leave a
row pointing at truncated bytes, or worse, no row and bytes nobody can find.
"""
path = _path_for(key)
path.parent.mkdir(parents=True, exist_ok=True)
digest = hashlib.sha256()
written = 0
try:
with path.open("wb") as sink:
while True:
chunk = source.read(CHUNK_BYTES)
if not chunk:
break
written += len(chunk)
if written > max_bytes:
raise ValueError("file is larger than the limit")
digest.update(chunk)
sink.write(chunk)
except Exception:
path.unlink(missing_ok=True)
raise
return written, digest.hexdigest()
def read(key: str) -> Iterator[bytes]:
"""Stream a file back. Raises `FileNotFoundError` if the bytes have gone."""
path = _path_for(key)
with path.open("rb") as source:
while True:
chunk = source.read(CHUNK_BYTES)
if not chunk:
return
yield chunk
def delete(key: str) -> None:
"""Remove the bytes. Missing is success — the point is that they are gone.
Failure is logged rather than raised: the caller has already decided the
document is deleted, and a file left on disk is a cost, not a correctness
problem. Raising here would leave the row undeleted *and* the file present.
"""
try:
_path_for(key).unlink(missing_ok=True)
except Exception:
logger.warning("could not remove stored file %s", key, exc_info=True)
def exists(key: str) -> bool:
try:
return _path_for(key).is_file()
except ValueError:
return False