Files
ds-chat/backend/app/storage.py
T
ksmithandClaude Sonnet 5 62e4760c8a Add admin-configurable upload size limits
The 8MB image/file/avatar cap is now a site setting (UploadSettings,
single-row table like SmtpSettings) editable from the Admin Settings tab,
instead of a hardcoded constant. All three upload endpoints read the live
value and interpolate it into their 413 messages. A new GET
/api/uploads/limit endpoint (open to any authenticated user, unlike the
admin-only settings endpoints) lets the composer reject an oversized file
client-side before it ever hits the network, though the server still
enforces the same cap independently.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-08-15 20:44:05 -06:00

106 lines
3.5 KiB
Python

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/chatapp
# 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"),
}
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)