ds-chat/backend/app/storage.py

import io
import pathlib
import uuid

from PIL import Image, UnidentifiedImageError

# backend/app/storage.py -> backend/ -> repo root -- same
# resolve-relative-to-file convention FRONTEND_DIST uses in app/main.py, so
# this lands in the right place in both local dev and the /srv/ds-chat
# production layout with zero new config.
UPLOADS_DIR = pathlib.Path(__file__).resolve().parent.parent.parent / "uploads"

# Seed value for the admin-configurable UploadSettings row (see
# app/services/upload_settings_service.py) -- also the fallback `read_capped`
# default for call sites that don't look up the live setting.
DEFAULT_MAX_UPLOAD_BYTES = 8 * 1024 * 1024
_READ_CHUNK_BYTES = 1024 * 1024
_MAX_DIMENSION = 2000

# (storage extension, Pillow format name)
ALLOWED_IMAGE_CONTENT_TYPES: dict[str, tuple[str, str]] = {
    "image/jpeg": (".jpg", "JPEG"),
    "image/png": (".png", "PNG"),
    "image/gif": (".gif", "GIF"),
    "image/webp": (".webp", "WEBP"),
}

# #65: browser-natively-playable video formats -- used to decide whether a
# stored MessageFile gets served inline (a <video> tag can actually play
# it) or forced to download like every other non-image attachment (see
# rooms.py's file-serve endpoint). Deliberately a strict allowlist, not
# "every video/* type": .mov (video/quicktime) has spotty <video> support
# outside Safari, and more importantly this is the one thing standing
# between "serve with the browser trusting our declared Content-Type" and
# reopening the same-origin-script-execution risk Content-Disposition:
# attachment exists to close off for arbitrary uploads -- it must only
# ever contain types a <video> tag renders as media, never as something
# that could execute script.
INLINE_SAFE_VIDEO_CONTENT_TYPES = frozenset({"video/mp4", "video/webm", "video/ogg"})


class UploadTooLargeError(Exception):
    pass


class InvalidImageError(Exception):
    pass


async def read_capped(file, cap: int = DEFAULT_MAX_UPLOAD_BYTES) -> bytes:
    """Reads an UploadFile-like object in chunks, raising as soon as `cap`
    is exceeded rather than after buffering the whole (potentially huge)
    body first."""
    chunks = []
    total = 0
    while True:
        chunk = await file.read(_READ_CHUNK_BYTES)
        if not chunk:
            break
        total += len(chunk)
        if total > cap:
            raise UploadTooLargeError()
        chunks.append(chunk)
    return b"".join(chunks)


def process_image(
    data: bytes,
    content_type: str,
    *,
    square: bool = False,
    max_dimension: int | None = None,
) -> tuple[bytes, str]:
    """Confirms `data` is a genuinely decodable image (not just a spoofed
    Content-Type header) and downscales it so its longer side is
    <=max_dimension (default 2000px) -- except GIF, left untouched so
    animation isn't collapsed to a single frame. When `square` is set
    (avatars), center-crops to the shorter side first. Returns
    (final_bytes, storage_extension)."""
    ext, pillow_format = ALLOWED_IMAGE_CONTENT_TYPES[content_type]
    dimension_cap = max_dimension or _MAX_DIMENSION

    try:
        with Image.open(io.BytesIO(data)) as probe:
            probe.verify()
    except (UnidentifiedImageError, OSError, ValueError) as exc:
        raise InvalidImageError() from exc

    if content_type == "image/gif":
        return data, ext

    # verify() leaves the image unusable for further processing, so reopen.
    image = Image.open(io.BytesIO(data))
    image.load()
    if pillow_format == "JPEG" and image.mode in ("RGBA", "P"):
        image = image.convert("RGB")
    if square:
        side = min(image.width, image.height)
        left = (image.width - side) // 2
        top = (image.height - side) // 2
        image = image.crop((left, top, left + side, top + side))
    image.thumbnail((dimension_cap, dimension_cap))
    out = io.BytesIO()
    image.save(out, format=pillow_format)
    return out.getvalue(), ext


def save_file(data: bytes, ext: str) -> str:
    UPLOADS_DIR.mkdir(parents=True, exist_ok=True)
    storage_filename = f"{uuid.uuid4()}{ext}"
    (UPLOADS_DIR / storage_filename).write_bytes(data)
    return storage_filename


def delete_file(storage_filename: str) -> None:
    """Best-effort delete -- a missing file (already gone, or never
    written) is not an error."""
    (UPLOADS_DIR / storage_filename).unlink(missing_ok=True)