Added server-side image validation and tests

Author: chriscarrollsmith

docs/installation.qmd

@@ -10,7 +10,7 @@ If you use VSCode with Docker to develop in a container, the following VSCode De
 {
 	"name": "Python 3",
 	"image": "mcr.microsoft.com/devcontainers/python:1-3.12-bookworm",
-	"postCreateCommand": "sudo apt update && sudo apt install -y python3-dev libpq-dev graphviz && uv venv && uv sync",
+	"postCreateCommand": "sudo apt update && sudo apt install -y python3-dev libpq-dev graphviz libwebp-dev && uv venv && uv sync",
 	"features": {
 		"ghcr.io/va-h/devcontainers-features/uv:1": {
 			"version": "latest"
@@ -61,7 +61,7 @@ Install Docker Desktop and Docker Compose for your operating system by following
 For Ubuntu/Debian:
 
 ``` bash
-sudo apt update && sudo apt install -y python3-dev libpq-dev
+sudo apt update && sudo apt install -y python3-dev libpq-dev libwebp-dev
 ```
 
 For macOS:

index.qmd

@@ -84,7 +84,7 @@ Install Docker Desktop and Coker Compose for your operating system by following
 For Ubuntu/Debian:
 
 ``` bash
-sudo apt update && sudo apt install -y python3-dev libpq-dev
+sudo apt update && sudo apt install -y python3-dev libpq-dev libwebp-dev
 ```
 
 For macOS:

routers/user.py

@@ -1,41 +1,15 @@
-from fastapi import APIRouter, Depends, Form, UploadFile, File, HTTPException
+from fastapi import APIRouter, Depends, Form, UploadFile, File
 from fastapi.responses import RedirectResponse, Response
 from pydantic import BaseModel, EmailStr
 from sqlmodel import Session
 from typing import Optional
 from utils.models import User, DataIntegrityError
 from utils.auth import get_session, get_authenticated_user, verify_password, PasswordValidationError
-from PIL import Image
-import io
+from utils.images import validate_and_process_image
 
 router = APIRouter(prefix="/user", tags=["user"])
 
 
-# --- Constants ---
-
-
-# 2MB in bytes
-MAX_FILE_SIZE = 2 * 1024 * 1024
-ALLOWED_CONTENT_TYPES = {
-    'image/jpeg',
-    'image/png',
-    'image/webp'
-}
-MIN_DIMENSION = 100
-MAX_DIMENSION = 2000
-TARGET_SIZE = 500
-
-
-# --- Custom Exceptions ---
-
-
-class InvalidImageError(HTTPException):
-    """Raised when an invalid image is uploaded"""
-
-    def __init__(self, message: str = "Invalid image file"):
-        super().__init__(status_code=400, detail=message)
-
-
 # --- Server Request and Response Models ---
 
 
@@ -90,55 +64,12 @@ async def update_profile(
 ):
     # Handle avatar update
     if user_profile.avatar_file:
-        # Check file size
-        if len(user_profile.avatar_file) > MAX_FILE_SIZE:
-            raise InvalidImageError(
-                message="File too large (max 2MB)"
-            )
-
-        # Check file type
-        if user_profile.avatar_content_type not in ALLOWED_CONTENT_TYPES:
-            raise InvalidImageError(
-                message="Invalid file type. Must be JPEG, PNG, or WebP"
-            )
-
-        try:
-            # Open and validate image
-            image = Image.open(io.BytesIO(user_profile.avatar_file))
-            width, height = image.size
-
-            # Check minimum dimensions
-            if width < MIN_DIMENSION or height < MIN_DIMENSION:
-                raise InvalidImageError(
-                    message=f"Image too small. Minimum dimension is {MIN_DIMENSION}px"
-                )
-
-            # Check maximum dimensions
-            if width > MAX_DIMENSION or height > MAX_DIMENSION:
-                raise InvalidImageError(
-                    message=f"Image too large. Maximum dimension is {MAX_DIMENSION}px"
-                )
-
-            # Crop to square and resize
-            min_dim = min(width, height)
-            left = (width - min_dim) // 2
-            top = (height - min_dim) // 2
-            right = left + min_dim
-            bottom = top + min_dim
-            
-            image = image.crop((left, top, right, bottom))
-            image = image.resize((TARGET_SIZE, TARGET_SIZE), Image.Resampling.LANCZOS)
-
-            # Convert back to bytes
-            output = io.BytesIO()
-            image.save(output, format='PNG')
-            user_profile.avatar_file = output.getvalue()
-            user_profile.avatar_content_type = 'image/png'
-
-        except Exception as e:
-            raise InvalidImageError(
-                message="Invalid image file"
-            )
+        processed_image, content_type = validate_and_process_image(
+            user_profile.avatar_file,
+            user_profile.avatar_content_type
+        )
+        user_profile.avatar_file = processed_image
+        user_profile.avatar_content_type = content_type
 
     # Update user details
     user.name = user_profile.name

tests/test_images.py

@@ -0,0 +1,104 @@
+import pytest
+from PIL import Image
+import io
+from utils.images import (
+    validate_and_process_image,
+    InvalidImageError,
+    MAX_FILE_SIZE,
+    MIN_DIMENSION,
+    MAX_DIMENSION,
+    TARGET_SIZE
+)
+
+def create_test_image(width: int, height: int, format: str = 'PNG') -> bytes:
+    """Helper function to create test images"""
+    image = Image.new('RGB', (width, height), color='red')
+    output = io.BytesIO()
+    image.save(output, format=format)
+    return output.getvalue()
+
+def test_webp_dependencies_are_installed():
+    """Test that webp dependencies are installed"""
+    assert '.webp' in Image.registered_extensions(), "WebP dependencies are not installed (e.g., libwebp-dev on Linux)"
+
+def test_valid_square_image():
+    """Test processing a valid square image"""
+    image_data = create_test_image(500, 500)
+    processed_data, content_type = validate_and_process_image(image_data, 'image/png')
+    
+    # Verify the processed image
+    processed_image = Image.open(io.BytesIO(processed_data))
+    assert processed_image.size == (TARGET_SIZE, TARGET_SIZE)
+    assert content_type == 'image/png'
+
+def test_valid_rectangular_image():
+    """Test processing a valid rectangular image"""
+    image_data = create_test_image(800, 600)
+    processed_data, content_type = validate_and_process_image(image_data, 'image/png')
+    
+    # Verify the processed image
+    processed_image = Image.open(io.BytesIO(processed_data))
+    assert processed_image.size == (TARGET_SIZE, TARGET_SIZE)
+    assert content_type == 'image/png'
+
+def test_minimum_size_image():
+    """Test processing an image with minimum allowed dimensions"""
+    image_data = create_test_image(MIN_DIMENSION, MIN_DIMENSION)
+    processed_data, content_type = validate_and_process_image(image_data, 'image/png')
+    
+    processed_image = Image.open(io.BytesIO(processed_data))
+    assert processed_image.size == (TARGET_SIZE, TARGET_SIZE)
+
+def test_too_small_image():
+    """Test that too small images are rejected"""
+    image_data = create_test_image(MIN_DIMENSION - 1, MIN_DIMENSION - 1)
+    with pytest.raises(InvalidImageError) as exc_info:
+        validate_and_process_image(image_data, 'image/png')
+    assert "Image too small" in str(exc_info.value.detail)
+
+def test_too_large_image():
+    """Test that too large images are rejected"""
+    image_data = create_test_image(MAX_DIMENSION + 1, MAX_DIMENSION + 1)
+    with pytest.raises(InvalidImageError) as exc_info:
+        validate_and_process_image(image_data, 'image/png')
+    assert "Image too large" in str(exc_info.value.detail)
+
+def test_invalid_file_type():
+    """Test that invalid file types are rejected"""
+    image_data = create_test_image(500, 500)
+    with pytest.raises(InvalidImageError) as exc_info:
+        validate_and_process_image(image_data, 'image/gif')
+    assert "Invalid file type" in str(exc_info.value.detail)
+
+def test_file_too_large():
+    """Test that files exceeding MAX_FILE_SIZE are rejected"""
+    # Create a large file that exceeds MAX_FILE_SIZE
+    large_image_data = b'0' * (MAX_FILE_SIZE + 1)
+    with pytest.raises(InvalidImageError) as exc_info:
+        validate_and_process_image(large_image_data, 'image/png')
+    assert "File too large" in str(exc_info.value.detail)
+
+def test_corrupt_image_data():
+    """Test that corrupt image data is rejected"""
+    corrupt_data = b'not an image'
+    with pytest.raises(InvalidImageError) as exc_info:
+        validate_and_process_image(corrupt_data, 'image/png')
+    assert "Invalid image file" in str(exc_info.value.detail)
+
+def test_different_image_formats():
+    """Test processing different valid image formats"""
+    formats = [
+        ('JPEG', 'image/jpeg'),
+        ('PNG', 'image/png'),
+        ('WEBP', 'image/webp')
+    ]
+    
+    for format_name, content_type in formats:
+        image_data = create_test_image(500, 500, format_name)
+        processed_data, result_type = validate_and_process_image(image_data, content_type)
+        
+        # Verify the processed image
+        processed_image = Image.open(io.BytesIO(processed_data))
+        assert processed_image.size == (TARGET_SIZE, TARGET_SIZE)
+        # Output should match input format
+        assert result_type == content_type

tests/test_user.py

@@ -1,9 +1,15 @@
 from fastapi.testclient import TestClient
 from httpx import Response
 from sqlmodel import Session
+from unittest.mock import patch
 
 from main import app
 from utils.models import User
+from utils.images import InvalidImageError
+
+# Mock data for consistent testing
+MOCK_IMAGE_DATA = b"processed fake image data"
+MOCK_CONTENT_TYPE = "image/png"
 
 
 def test_update_profile_unauthorized(unauth_client: TestClient):
@@ -23,11 +29,12 @@ def test_update_profile_unauthorized(unauth_client: TestClient):
     assert response.headers["location"] == app.url_path_for("read_login")
 
 
-def test_update_profile_authorized(auth_client: TestClient, test_user: User, session: Session):
+@patch('routers.user.validate_and_process_image')
+def test_update_profile_authorized(mock_validate, auth_client: TestClient, test_user: User, session: Session):
     """Test that authorized users can edit their profile"""
 
-    # Create test image data
-    test_image_data = b"fake image data"
+    # Configure mock to return processed image data
+    mock_validate.return_value = (MOCK_IMAGE_DATA, MOCK_CONTENT_TYPE)
 
     # Update profile
     response: Response = auth_client.post(
@@ -37,7 +44,7 @@ def test_update_profile_authorized(auth_client: TestClient, test_user: User, ses
             "email": "[email protected]",
         },
         files={
-            "avatar_file": ("test_avatar.jpg", test_image_data, "image/jpeg")
+            "avatar_file": ("test_avatar.jpg", b"fake image data", "image/jpeg")
         },
         follow_redirects=False
     )
@@ -48,8 +55,11 @@ def test_update_profile_authorized(auth_client: TestClient, test_user: User, ses
     session.refresh(test_user)
     assert test_user.name == "Updated Name"
     assert test_user.email == "[email protected]"
-    assert test_user.avatar_data == test_image_data
-    assert test_user.avatar_content_type == "image/jpeg"
+    assert test_user.avatar_data == MOCK_IMAGE_DATA
+    assert test_user.avatar_content_type == MOCK_CONTENT_TYPE
+
+    # Verify mock was called correctly
+    mock_validate.assert_called_once()
 
 
 def test_update_profile_without_avatar(auth_client: TestClient, test_user: User, session: Session):
@@ -110,18 +120,21 @@ def test_delete_account_success(auth_client: TestClient, test_user: User, sessio
     assert user is None
 
 
-def test_get_avatar_authorized(auth_client: TestClient, test_user: User):
+@patch('routers.user.validate_and_process_image')
+def test_get_avatar_authorized(mock_validate, auth_client: TestClient, test_user: User):
     """Test getting user avatar"""
+    # Configure mock to return processed image data
+    mock_validate.return_value = (MOCK_IMAGE_DATA, MOCK_CONTENT_TYPE)
+
     # First upload an avatar
-    test_image_data = b"fake image data"
     auth_client.post(
         app.url_path_for("update_profile"),
         data={
             "name": test_user.name,
             "email": test_user.email,
         },
         files={
-            "avatar_file": ("test_avatar.jpg", test_image_data, "image/jpeg")
+            "avatar_file": ("test_avatar.jpg", b"fake image data", "image/jpeg")
         },
     )
 
@@ -130,8 +143,8 @@ def test_get_avatar_authorized(auth_client: TestClient, test_user: User):
         app.url_path_for("get_avatar")
     )
     assert response.status_code == 200
-    assert response.content == test_image_data
-    assert response.headers["content-type"] == "image/jpeg"
+    assert response.content == MOCK_IMAGE_DATA
+    assert response.headers["content-type"] == MOCK_CONTENT_TYPE
 
 
 def test_get_avatar_unauthorized(unauth_client: TestClient):
@@ -142,3 +155,24 @@ def test_get_avatar_unauthorized(unauth_client: TestClient):
     )
     assert response.status_code == 303
     assert response.headers["location"] == app.url_path_for("read_login")
+
+
+# Add new test for invalid image
+@patch('routers.user.validate_and_process_image')
+def test_update_profile_invalid_image(mock_validate, auth_client: TestClient):
+    """Test that invalid images are rejected"""
+    # Configure mock to raise InvalidImageError
+    mock_validate.side_effect = InvalidImageError("Invalid test image")
+    
+    response: Response = auth_client.post(
+        app.url_path_for("update_profile"),
+        data={
+            "name": "Updated Name",
+            "email": "[email protected]",
+        },
+        files={
+            "avatar_file": ("test_avatar.jpg", b"invalid image data", "image/jpeg")
+        },
+    )
+    assert response.status_code == 400
+    assert "Invalid test image" in response.text