feat(security): add configurable file upload size limits with optional splitting
- Add MAX_UPLOAD_SIZE config (default 1GB) to prevent resource exhaustion - Add MAX_SINGLE_FILE_SIZE config for optional PDF file splitting - Implement automatic PDF splitting when files exceed single file limit - Update upload endpoint to use configured limits instead of hardcoded 500MB - Add comprehensive tests for upload limits and file splitting - Document configuration in ConfigurationGuide.md and SECURITY_AUDIT.md - Reference SECURITY_AUDIT.md in error messages for user guidance Co-authored-by: christianlouis <361235+christianlouis@users.noreply.github.com>
This commit is contained in:
copilot-swe-agent[bot]
@@ -0,0 +1,210 @@
|
||||
"""
|
||||
Tests for the file splitting utility module.
|
||||
|
||||
Tests cover:
|
||||
- PDF splitting by size
|
||||
- Handling of edge cases (empty PDFs, single-page PDFs, etc.)
|
||||
- Error handling
|
||||
- should_split_file function
|
||||
"""
|
||||
|
||||
import os
|
||||
import pytest
|
||||
import tempfile
|
||||
from unittest.mock import patch, Mock
|
||||
from PyPDF2 import PdfWriter, PdfReader
|
||||
|
||||
from app.utils.file_splitting import split_pdf_by_size, should_split_file
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_multipage_pdf():
|
||||
"""Create a sample multi-page PDF for testing."""
|
||||
with tempfile.NamedTemporaryFile(mode="wb", suffix=".pdf", delete=False) as f:
|
||||
writer = PdfWriter()
|
||||
|
||||
# Add 5 pages to the PDF
|
||||
for i in range(5):
|
||||
writer.add_blank_page(width=200, height=200)
|
||||
|
||||
writer.write(f)
|
||||
pdf_path = f.name
|
||||
|
||||
yield pdf_path
|
||||
|
||||
# Cleanup
|
||||
if os.path.exists(pdf_path):
|
||||
os.remove(pdf_path)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_single_page_pdf():
|
||||
"""Create a sample single-page PDF for testing."""
|
||||
with tempfile.NamedTemporaryFile(mode="wb", suffix=".pdf", delete=False) as f:
|
||||
writer = PdfWriter()
|
||||
writer.add_blank_page(width=200, height=200)
|
||||
writer.write(f)
|
||||
pdf_path = f.name
|
||||
|
||||
yield pdf_path
|
||||
|
||||
# Cleanup
|
||||
if os.path.exists(pdf_path):
|
||||
os.remove(pdf_path)
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestSplitPdfBySize:
|
||||
"""Tests for the split_pdf_by_size function."""
|
||||
|
||||
def test_split_pdf_basic(self, sample_multipage_pdf):
|
||||
"""Test basic PDF splitting functionality."""
|
||||
# Use a small size limit to force splitting
|
||||
max_size = 5000 # 5KB - should split the 5-page PDF
|
||||
|
||||
split_files = split_pdf_by_size(sample_multipage_pdf, max_size)
|
||||
|
||||
# Verify files were created
|
||||
assert len(split_files) > 1, "PDF should be split into multiple files"
|
||||
|
||||
# Verify all split files exist
|
||||
for split_file in split_files:
|
||||
assert os.path.exists(split_file), f"Split file {split_file} should exist"
|
||||
assert os.path.getsize(split_file) > 0, f"Split file {split_file} should not be empty"
|
||||
|
||||
# Verify total pages match original
|
||||
original_reader = PdfReader(sample_multipage_pdf)
|
||||
total_split_pages = sum(len(PdfReader(f).pages) for f in split_files)
|
||||
assert total_split_pages == len(original_reader.pages), "Total pages should match original"
|
||||
|
||||
# Cleanup split files
|
||||
for split_file in split_files:
|
||||
if os.path.exists(split_file):
|
||||
os.remove(split_file)
|
||||
|
||||
def test_split_pdf_with_output_dir(self, sample_multipage_pdf):
|
||||
"""Test PDF splitting with custom output directory."""
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
max_size = 5000
|
||||
|
||||
split_files = split_pdf_by_size(sample_multipage_pdf, max_size, output_dir=temp_dir)
|
||||
|
||||
# Verify files are in the specified directory
|
||||
for split_file in split_files:
|
||||
assert os.path.dirname(split_file) == temp_dir, "Split files should be in output_dir"
|
||||
assert os.path.exists(split_file), f"Split file {split_file} should exist"
|
||||
|
||||
# Files will be cleaned up with temp_dir
|
||||
|
||||
def test_split_pdf_single_page(self, sample_single_page_pdf):
|
||||
"""Test splitting a single-page PDF."""
|
||||
max_size = 1000 # Very small limit
|
||||
|
||||
split_files = split_pdf_by_size(sample_single_page_pdf, max_size)
|
||||
|
||||
# Should create at least one file (might be just the single page)
|
||||
assert len(split_files) >= 1, "Should create at least one output file"
|
||||
|
||||
# Verify the split file exists
|
||||
for split_file in split_files:
|
||||
assert os.path.exists(split_file), f"Split file {split_file} should exist"
|
||||
|
||||
# Cleanup
|
||||
for split_file in split_files:
|
||||
if os.path.exists(split_file):
|
||||
os.remove(split_file)
|
||||
|
||||
def test_split_pdf_large_limit(self, sample_multipage_pdf):
|
||||
"""Test that PDF is not split when limit is very large."""
|
||||
max_size = 10 * 1024 * 1024 # 10MB - much larger than test PDF
|
||||
|
||||
split_files = split_pdf_by_size(sample_multipage_pdf, max_size)
|
||||
|
||||
# Should create only one file (no splitting needed)
|
||||
assert len(split_files) == 1, "PDF should not be split with large limit"
|
||||
|
||||
# Verify total pages match
|
||||
original_reader = PdfReader(sample_multipage_pdf)
|
||||
split_reader = PdfReader(split_files[0])
|
||||
assert len(split_reader.pages) == len(original_reader.pages), "All pages should be in single file"
|
||||
|
||||
# Cleanup
|
||||
for split_file in split_files:
|
||||
if os.path.exists(split_file):
|
||||
os.remove(split_file)
|
||||
|
||||
def test_split_pdf_file_not_found(self):
|
||||
"""Test error handling when PDF file doesn't exist."""
|
||||
with pytest.raises(FileNotFoundError):
|
||||
split_pdf_by_size("/nonexistent/file.pdf", 1000)
|
||||
|
||||
def test_split_pdf_invalid_pdf(self):
|
||||
"""Test error handling with invalid/corrupted PDF."""
|
||||
# Create a file that's not a valid PDF
|
||||
with tempfile.NamedTemporaryFile(mode="w", suffix=".pdf", delete=False) as f:
|
||||
f.write("This is not a PDF file")
|
||||
invalid_pdf = f.name
|
||||
|
||||
try:
|
||||
with pytest.raises(ValueError, match="Invalid or corrupted PDF"):
|
||||
split_pdf_by_size(invalid_pdf, 1000)
|
||||
finally:
|
||||
if os.path.exists(invalid_pdf):
|
||||
os.remove(invalid_pdf)
|
||||
|
||||
def test_split_pdf_naming_convention(self, sample_multipage_pdf):
|
||||
"""Test that split files follow expected naming convention."""
|
||||
max_size = 5000
|
||||
|
||||
split_files = split_pdf_by_size(sample_multipage_pdf, max_size)
|
||||
|
||||
# Verify naming pattern: basename_partN.pdf
|
||||
base_name = os.path.splitext(os.path.basename(sample_multipage_pdf))[0]
|
||||
|
||||
for i, split_file in enumerate(split_files, start=1):
|
||||
filename = os.path.basename(split_file)
|
||||
assert filename.startswith(base_name), f"Filename should start with {base_name}"
|
||||
assert f"_part{i}.pdf" in filename, f"Filename should contain _part{i}.pdf"
|
||||
|
||||
# Cleanup
|
||||
for split_file in split_files:
|
||||
if os.path.exists(split_file):
|
||||
os.remove(split_file)
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestShouldSplitFile:
|
||||
"""Tests for the should_split_file function."""
|
||||
|
||||
def test_should_split_when_file_exceeds_limit(self, sample_multipage_pdf):
|
||||
"""Test that function returns True when file exceeds limit."""
|
||||
file_size = os.path.getsize(sample_multipage_pdf)
|
||||
max_size = file_size - 1 # Set limit just below file size
|
||||
|
||||
result = should_split_file(sample_multipage_pdf, max_size)
|
||||
assert result is True, "Should return True when file exceeds limit"
|
||||
|
||||
def test_should_not_split_when_file_under_limit(self, sample_multipage_pdf):
|
||||
"""Test that function returns False when file is under limit."""
|
||||
file_size = os.path.getsize(sample_multipage_pdf)
|
||||
max_size = file_size + 1000 # Set limit above file size
|
||||
|
||||
result = should_split_file(sample_multipage_pdf, max_size)
|
||||
assert result is False, "Should return False when file is under limit"
|
||||
|
||||
def test_should_not_split_when_limit_is_none(self, sample_multipage_pdf):
|
||||
"""Test that function returns False when max_single_file_size is None."""
|
||||
result = should_split_file(sample_multipage_pdf, None)
|
||||
assert result is False, "Should return False when limit is None (splitting disabled)"
|
||||
|
||||
def test_should_not_split_when_file_not_exists(self):
|
||||
"""Test that function returns False when file doesn't exist."""
|
||||
result = should_split_file("/nonexistent/file.pdf", 1000)
|
||||
assert result is False, "Should return False when file doesn't exist"
|
||||
|
||||
def test_should_not_split_exact_size(self, sample_multipage_pdf):
|
||||
"""Test behavior when file size exactly matches limit."""
|
||||
file_size = os.path.getsize(sample_multipage_pdf)
|
||||
|
||||
result = should_split_file(sample_multipage_pdf, file_size)
|
||||
assert result is False, "Should return False when file size equals limit"
|
||||
+132
-3
@@ -157,13 +157,15 @@ class TestInvalidFileUploads:
|
||||
"""Tests for handling invalid or problematic file uploads."""
|
||||
|
||||
def test_upload_file_too_large(self, client: TestClient):
|
||||
"""Test that files over 500MB are rejected."""
|
||||
"""Test that files exceeding MAX_UPLOAD_SIZE are rejected."""
|
||||
from app.config import settings
|
||||
|
||||
# Create a large file content (mock it to avoid memory issues)
|
||||
large_content = b"x" * 1024 # 1KB for testing
|
||||
|
||||
with patch("os.path.getsize") as mock_getsize:
|
||||
# Mock the file size to be over 500MB
|
||||
mock_getsize.return_value = 501 * 1024 * 1024 # 501MB
|
||||
# Mock the file size to be over the configured limit
|
||||
mock_getsize.return_value = settings.max_upload_size + 1
|
||||
|
||||
response = client.post(
|
||||
"/api/ui-upload", files={"file": ("huge.pdf", io.BytesIO(large_content), "application/pdf")}
|
||||
@@ -171,6 +173,7 @@ class TestInvalidFileUploads:
|
||||
|
||||
assert response.status_code == 413 # Request Entity Too Large
|
||||
assert "too large" in response.json()["detail"].lower()
|
||||
assert "SECURITY_AUDIT.md" in response.json()["detail"]
|
||||
|
||||
def test_upload_executable_file(self, client: TestClient, mock_celery_tasks):
|
||||
"""Test that executable files are handled (attempted conversion)."""
|
||||
@@ -411,3 +414,129 @@ class TestUploadMimeTypeDetection:
|
||||
assert response.status_code == 200
|
||||
# Should route to convert_to_pdf based on .jpg extension
|
||||
mock_celery_tasks["convert_to_pdf"].assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
class TestFileSplitting:
|
||||
"""Tests for file splitting functionality when MAX_SINGLE_FILE_SIZE is configured."""
|
||||
|
||||
def test_pdf_splitting_when_configured(self, client: TestClient, sample_pdf_path: str, mock_celery_tasks):
|
||||
"""Test that PDFs are split when they exceed MAX_SINGLE_FILE_SIZE."""
|
||||
from app.config import settings
|
||||
|
||||
# Mock settings to enable file splitting with a very small limit
|
||||
with patch.object(settings, "max_single_file_size", 100): # 100 bytes limit
|
||||
# Mock the split_pdf_by_size function to return fake split files
|
||||
with patch("app.api.files.split_pdf_by_size") as mock_split:
|
||||
mock_split.return_value = [
|
||||
"/workdir/test_part1.pdf",
|
||||
"/workdir/test_part2.pdf",
|
||||
"/workdir/test_part3.pdf",
|
||||
]
|
||||
|
||||
# Mock should_split_file to return True
|
||||
with patch("app.api.files.should_split_file", return_value=True):
|
||||
with open(sample_pdf_path, "rb") as f:
|
||||
response = client.post("/api/ui-upload", files={"file": ("large.pdf", f, "application/pdf")})
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
|
||||
# Verify response indicates splitting occurred
|
||||
assert "split_into_parts" in data
|
||||
assert data["split_into_parts"] == 3
|
||||
assert "task_ids" in data
|
||||
assert len(data["task_ids"]) == 3
|
||||
assert "message" in data
|
||||
assert "split" in data["message"].lower()
|
||||
|
||||
# Verify each split file was queued for processing
|
||||
assert mock_celery_tasks["process_document"].call_count == 3
|
||||
|
||||
def test_no_splitting_when_not_configured(self, client: TestClient, sample_pdf_path: str, mock_celery_tasks):
|
||||
"""Test that PDFs are not split when MAX_SINGLE_FILE_SIZE is None."""
|
||||
from app.config import settings
|
||||
|
||||
# Ensure max_single_file_size is None (default)
|
||||
with patch.object(settings, "max_single_file_size", None):
|
||||
with open(sample_pdf_path, "rb") as f:
|
||||
response = client.post("/api/ui-upload", files={"file": ("document.pdf", f, "application/pdf")})
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
|
||||
# Verify no splitting occurred
|
||||
assert "split_into_parts" not in data
|
||||
assert "task_id" in data # Single task ID, not task_ids array
|
||||
assert data["status"] == "queued"
|
||||
|
||||
# Verify file was processed directly without splitting
|
||||
mock_celery_tasks["process_document"].assert_called_once()
|
||||
|
||||
def test_no_splitting_for_small_files(self, client: TestClient, sample_pdf_path: str, mock_celery_tasks):
|
||||
"""Test that small PDFs are not split even when MAX_SINGLE_FILE_SIZE is configured."""
|
||||
from app.config import settings
|
||||
|
||||
# Configure a very large limit
|
||||
with patch.object(settings, "max_single_file_size", 1000000000): # 1GB limit
|
||||
with open(sample_pdf_path, "rb") as f:
|
||||
response = client.post("/api/ui-upload", files={"file": ("small.pdf", f, "application/pdf")})
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
|
||||
# Verify no splitting occurred for small file
|
||||
assert "split_into_parts" not in data
|
||||
assert "task_id" in data
|
||||
|
||||
# Verify file was processed directly
|
||||
mock_celery_tasks["process_document"].assert_called_once()
|
||||
|
||||
def test_splitting_fallback_on_error(self, client: TestClient, sample_pdf_path: str, mock_celery_tasks):
|
||||
"""Test that if splitting fails, the file is processed as a whole."""
|
||||
from app.config import settings
|
||||
|
||||
with patch.object(settings, "max_single_file_size", 100): # Small limit
|
||||
with patch("app.api.files.should_split_file", return_value=True):
|
||||
# Mock split_pdf_by_size to raise an exception
|
||||
with patch("app.api.files.split_pdf_by_size", side_effect=Exception("Split failed")):
|
||||
with open(sample_pdf_path, "rb") as f:
|
||||
response = client.post("/api/ui-upload", files={"file": ("document.pdf", f, "application/pdf")})
|
||||
|
||||
# Should still succeed, falling back to processing the whole file
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
|
||||
# Verify no splitting data in response
|
||||
assert "split_into_parts" not in data
|
||||
assert "task_id" in data
|
||||
|
||||
# File should be processed as a whole
|
||||
mock_celery_tasks["process_document"].assert_called_once()
|
||||
|
||||
def test_non_pdf_not_split(self, client: TestClient, mock_celery_tasks):
|
||||
"""Test that non-PDF files are never split, even with MAX_SINGLE_FILE_SIZE configured."""
|
||||
from app.config import settings
|
||||
|
||||
with patch.object(settings, "max_single_file_size", 100): # Small limit
|
||||
# Upload an image file
|
||||
image_content = (
|
||||
b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01"
|
||||
b"\x08\x06\x00\x00\x00\x1f\x15\xc4\x89\x00\x00\x00\nIDATx\x9cc\x00\x01"
|
||||
b"\x00\x00\x05\x00\x01\r\n-\xb4\x00\x00\x00\x00IEND\xaeB`\x82"
|
||||
)
|
||||
|
||||
response = client.post(
|
||||
"/api/ui-upload", files={"file": ("image.png", io.BytesIO(image_content), "image/png")}
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
|
||||
# Images should not be split (they're converted to PDF first)
|
||||
assert "split_into_parts" not in data
|
||||
assert "task_id" in data
|
||||
|
||||
# Should be queued for conversion
|
||||
mock_celery_tasks["convert_to_pdf"].assert_called_once()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user