124 lines
5.7 KiB
Python
124 lines
5.7 KiB
Python
"""Хранение аватаров пользователей: валидация загрузки, файлы на диске, URL.
|
||
|
||
Файл лежит на диске `MEDIA_ROOT/avatars/{user_id}.{ext}`; в БД (`users.avatar_path`)
|
||
хранится путь относительно `MEDIA_ROOT` (`avatars/{user_id}.{ext}`) — тот же
|
||
приём, что и у записей аудиотреков (`recordings_dir`, `core/config.py`).
|
||
"""
|
||
|
||
import uuid
|
||
from collections.abc import Callable
|
||
from pathlib import Path
|
||
|
||
from fastapi import UploadFile
|
||
|
||
# Лимит размера загружаемого аватара — 2 МБ.
|
||
MAX_AVATAR_SIZE_BYTES = 2 * 1024 * 1024
|
||
|
||
# Читаем файл чанками, не доверяя заголовку `Content-Length` (клиент может
|
||
# солгать о размере) — реальный размер считается по факту прочитанных байт.
|
||
_CHUNK_SIZE_BYTES = 64 * 1024
|
||
|
||
# Допустимые типы изображений -> расширение файла на диске.
|
||
_ALLOWED_CONTENT_TYPES: dict[str, str] = {
|
||
"image/jpeg": "jpg",
|
||
"image/png": "png",
|
||
"image/webp": "webp",
|
||
}
|
||
|
||
# Магические байты (сигнатуры) форматов — заголовку `Content-Type` от клиента
|
||
# доверять нельзя (легко подделать), реальный формат определяется по
|
||
# содержимому файла.
|
||
_MAGIC_CHECKS: dict[str, Callable[[bytes], bool]] = {
|
||
"image/jpeg": lambda head: head[:3] == b"\xff\xd8\xff",
|
||
"image/png": lambda head: head[:8] == b"\x89PNG\r\n\x1a\n",
|
||
"image/webp": lambda head: head[:4] == b"RIFF" and head[8:12] == b"WEBP",
|
||
}
|
||
|
||
# Достаточно первых 12 байт, чтобы проверить все сигнатуры выше (WebP —
|
||
# самая длинная проверка, требует байты 8..11 включительно).
|
||
_MAGIC_HEAD_SIZE = 12
|
||
|
||
|
||
class AvatarTooLargeError(Exception):
|
||
"""Загружаемый файл превышает `MAX_AVATAR_SIZE_BYTES` (413)."""
|
||
|
||
|
||
class AvatarInvalidTypeError(Exception):
|
||
"""`Content-Type` не входит в список допустимых либо не совпадает с содержимым (415)."""
|
||
|
||
|
||
async def read_and_validate_avatar(file: UploadFile) -> tuple[bytes, str]:
|
||
"""Прочитать содержимое файла аватара чанками и провалидировать тип/размер.
|
||
|
||
Возвращает `(содержимое, расширение)`. Порядок проверок: сначала
|
||
заявленный `Content-Type` (быстрый отсев), затем фактический размер по
|
||
мере чтения, затем магические байты содержимого — заявленный тип должен
|
||
совпасть с реальным (иначе подделка `Content-Type` не даст загрузить,
|
||
например, исполняемый файл под видом `image/png`).
|
||
"""
|
||
declared_type = file.content_type
|
||
if declared_type not in _ALLOWED_CONTENT_TYPES:
|
||
raise AvatarInvalidTypeError(f"unsupported_content_type: {declared_type}")
|
||
|
||
chunks: list[bytes] = []
|
||
total_size = 0
|
||
while True:
|
||
chunk = await file.read(_CHUNK_SIZE_BYTES)
|
||
if not chunk:
|
||
break
|
||
total_size += len(chunk)
|
||
if total_size > MAX_AVATAR_SIZE_BYTES:
|
||
raise AvatarTooLargeError(f"file exceeds {MAX_AVATAR_SIZE_BYTES} bytes")
|
||
chunks.append(chunk)
|
||
content = b"".join(chunks)
|
||
|
||
magic_check = _MAGIC_CHECKS[declared_type]
|
||
if not magic_check(content[:_MAGIC_HEAD_SIZE]):
|
||
raise AvatarInvalidTypeError("content_does_not_match_declared_content_type")
|
||
|
||
return content, _ALLOWED_CONTENT_TYPES[declared_type]
|
||
|
||
|
||
def _avatar_relative_path(user_id: uuid.UUID, ext: str) -> str:
|
||
"""Путь аватара относительно `MEDIA_ROOT`."""
|
||
return f"avatars/{user_id}.{ext}"
|
||
|
||
|
||
def save_avatar(
|
||
media_root: Path, user_id: uuid.UUID, content: bytes, ext: str, *, old_path: str | None
|
||
) -> str:
|
||
"""Сохранить содержимое аватара на диск, удалить предыдущий файл (если был другого формата).
|
||
|
||
Возвращает новый относительный путь (`users.avatar_path`).
|
||
"""
|
||
avatars_dir = media_root / "avatars"
|
||
avatars_dir.mkdir(parents=True, exist_ok=True)
|
||
delete_avatar(media_root, old_path)
|
||
relative_path = _avatar_relative_path(user_id, ext)
|
||
(media_root / relative_path).write_bytes(content)
|
||
return relative_path
|
||
|
||
|
||
def delete_avatar(media_root: Path, avatar_path: str | None) -> None:
|
||
"""Удалить файл аватара с диска, если он существует; `None`/отсутствие файла — no-op."""
|
||
if not avatar_path:
|
||
return
|
||
file_path = media_root / avatar_path
|
||
file_path.unlink(missing_ok=True)
|
||
|
||
|
||
def avatar_url(media_root: Path, avatar_path: str | None) -> str | None:
|
||
"""URL аватара с cache-busting параметром `v={mtime файла}`; `None`, если аватара нет.
|
||
|
||
`mtime` — не хранящееся в БД значение (файл может быть перезалит в обход
|
||
ORM, например, вручную на проде), поэтому считывается со диска на лету.
|
||
"""
|
||
if not avatar_path:
|
||
return None
|
||
file_path = media_root / avatar_path
|
||
try:
|
||
mtime = int(file_path.stat().st_mtime)
|
||
except FileNotFoundError:
|
||
return None
|
||
return f"/media/{avatar_path}?v={mtime}"
|