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.

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,15 @@
 import pytest
 from PIL import Image
 import io
+import asyncio
+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,
 )
@@ -89,6 +94,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:
+        asyncio.run(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 = asyncio.run(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]: