Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 7 additions & 1 deletion routers/core/user.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,10 @@
)
from utils.core.images import (
validate_and_process_image,
read_upload_with_size_limit,
reject_oversized_content_length,
MAX_FILE_SIZE,
MAX_AVATAR_UPLOAD_BYTES,
MIN_DIMENSION,
MAX_DIMENSION,
ALLOWED_CONTENT_TYPES,
Expand Down Expand Up @@ -131,7 +134,10 @@ async def update_profile(
# Handle avatar update
if avatar_changed:
assert avatar_file is not None
avatar_data = await avatar_file.read()
reject_oversized_content_length(
request.headers.get("content-length"), MAX_AVATAR_UPLOAD_BYTES
)
avatar_data = await read_upload_with_size_limit(avatar_file, MAX_FILE_SIZE)
avatar_content_type = avatar_file.content_type

processed_image, content_type = validate_and_process_image(
Expand Down
47 changes: 47 additions & 0 deletions tests/utils/test_images.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,16 @@
import pytest
from PIL import Image
import io
import asyncio
from concurrent.futures import ThreadPoolExecutor
from unittest.mock import AsyncMock, MagicMock
from utils.core.images import (
validate_and_process_image,
read_upload_with_size_limit,
reject_oversized_content_length,
InvalidImageError,
MAX_FILE_SIZE,
MAX_AVATAR_UPLOAD_BYTES,
MIN_DIMENSION,
MAX_DIMENSION,
)
Expand All @@ -18,6 +24,12 @@ def create_test_image(width: int, height: int, format: str = "PNG") -> bytes:
return output.getvalue()


def _run_async(coro):
"""Run a coroutine when pytest may already have an event loop active."""
with ThreadPoolExecutor(max_workers=1) as executor:
return executor.submit(asyncio.run, coro).result()


def test_webp_dependencies_are_installed():
"""Test that webp dependencies are installed"""
assert ".webp" in Image.registered_extensions(), (
Expand Down Expand Up @@ -89,6 +101,41 @@ def test_file_too_large():
assert "File too large" in str(exc_info.value.detail)


def test_read_upload_with_size_limit_rejects_oversized_stream():
upload = MagicMock()
chunk_size = 64 * 1024
upload.read = AsyncMock(
side_effect=[b"x" * chunk_size, b"x" * (MAX_FILE_SIZE - chunk_size + 1), b""]
)

with pytest.raises(InvalidImageError) as exc_info:
_run_async(read_upload_with_size_limit(upload, MAX_FILE_SIZE))

assert "File too large" in str(exc_info.value.detail)
assert upload.read.await_count == 2


def test_read_upload_with_size_limit_accepts_valid_stream():
upload = MagicMock()
upload.read = AsyncMock(side_effect=[b"abc", b"def", b""])

data = _run_async(read_upload_with_size_limit(upload, MAX_FILE_SIZE))

assert data == b"abcdef"


def test_reject_oversized_content_length():
with pytest.raises(InvalidImageError):
reject_oversized_content_length(
str(MAX_AVATAR_UPLOAD_BYTES + 1), MAX_AVATAR_UPLOAD_BYTES
)


def test_reject_oversized_content_length_allows_missing_or_valid():
reject_oversized_content_length(None, MAX_AVATAR_UPLOAD_BYTES)
reject_oversized_content_length(str(MAX_FILE_SIZE), MAX_AVATAR_UPLOAD_BYTES)


def test_corrupt_image_data():
"""Test that corrupt image data is rejected"""
corrupt_data = b"not an image"
Expand Down
36 changes: 36 additions & 0 deletions utils/core/images.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,13 +2,17 @@
from PIL import Image
import io
from typing import Tuple
from fastapi import UploadFile
from exceptions.http_exceptions import InvalidImageError


# --- Constants ---


MAX_FILE_SIZE = 2 * 1024 * 1024 # 2MB in bytes
# Multipart overhead for the profile form (fields + boundaries).
MAX_AVATAR_UPLOAD_BYTES = MAX_FILE_SIZE + 64 * 1024
READ_CHUNK_SIZE = 64 * 1024
ALLOWED_CONTENT_TYPES = {"image/jpeg": "JPEG", "image/png": "PNG", "image/webp": "WEBP"}
MIN_DIMENSION = 100
MAX_DIMENSION = 2000
Expand All @@ -17,6 +21,38 @@
# --- Functions ---


def reject_oversized_content_length(content_length: str | None, max_bytes: int) -> None:
"""Reject requests whose Content-Length exceeds the byte budget."""
if not content_length:
return
try:
declared_size = int(content_length)
except ValueError:
return
if declared_size > max_bytes:
raise InvalidImageError(message="File too large (max 2MB)")


async def read_upload_with_size_limit(
upload_file: UploadFile, max_bytes: int = MAX_FILE_SIZE
) -> bytes:
"""
Read an upload in fixed-size chunks, aborting before buffering more than
max_bytes into memory.
"""
chunks: list[bytes] = []
total = 0
while True:
chunk = await upload_file.read(READ_CHUNK_SIZE)
if not chunk:
break
total += len(chunk)
if total > max_bytes:
raise InvalidImageError(message="File too large (max 2MB)")
chunks.append(chunk)
return b"".join(chunks)


def validate_and_process_image(
image_data: bytes, content_type: str | None
) -> Tuple[bytes, str]:
Expand Down
Loading