fix: enforce avatar upload size limit before buffering body (#220)

* fix: enforce avatar upload size limit before buffering body

Reject oversized Content-Length headers and stream multipart uploads in
fixed-size chunks so workers never hold an entire large file in memory.

* fix: run async upload-limit tests outside pytest event loop

Use a thread-backed helper so asyncio.run is not called while pytest
already has an active event loop in CI.

Author: chriscarrollsmith

routers/core/user.py

@@ -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,
@@ -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(

tests/utils/test_images.py

@@ -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,
 )
@@ -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(), (
@@ -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"

utils/core/images.py

@@ -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
@@ -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]: