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:
@@ -11,6 +11,16 @@ ALLOW_FILE_DELETE=true # Allow deletion of file records
|
|||||||
PROCESSALL_THROTTLE_THRESHOLD=20 # Number of files above which throttling is applied (default: 20)
|
PROCESSALL_THROTTLE_THRESHOLD=20 # Number of files above which throttling is applied (default: 20)
|
||||||
PROCESSALL_THROTTLE_DELAY=3 # Delay in seconds between each task submission when throttling (default: 3)
|
PROCESSALL_THROTTLE_DELAY=3 # Delay in seconds between each task submission when throttling (default: 3)
|
||||||
|
|
||||||
|
# **File Upload Size Limits** (Security - see SECURITY_AUDIT.md)
|
||||||
|
# Maximum file upload size in bytes. Default: 1GB (1073741824 bytes)
|
||||||
|
# Prevents resource exhaustion attacks. Adjust based on your server capacity.
|
||||||
|
MAX_UPLOAD_SIZE=1073741824
|
||||||
|
|
||||||
|
# Maximum size for a single file chunk in bytes (optional)
|
||||||
|
# If set and a file exceeds this size, it will be split into smaller chunks for processing
|
||||||
|
# Default: None (no splitting). Example: 104857600 for 100MB chunks
|
||||||
|
# MAX_SINGLE_FILE_SIZE=104857600
|
||||||
|
|
||||||
# **Authentication**
|
# **Authentication**
|
||||||
AUTH_ENABLED=true
|
AUTH_ENABLED=true
|
||||||
# Generate a secure random string, for example:
|
# Generate a secure random string, for example:
|
||||||
|
|||||||
+30
-2
@@ -142,6 +142,33 @@ This document tracks security vulnerabilities found in DocuElevate and their rem
|
|||||||
- SSH keys
|
- SSH keys
|
||||||
- Explicit exclusion of patterns where needed
|
- Explicit exclusion of patterns where needed
|
||||||
|
|
||||||
|
### 5. File Upload Size Limits
|
||||||
|
**Status:** ✅ FIXED
|
||||||
|
**Severity:** MEDIUM
|
||||||
|
**Description:** No configurable limits on file upload sizes could lead to resource exhaustion attacks and DoS.
|
||||||
|
|
||||||
|
**Fix:** Implemented configurable file upload size limits with the following features:
|
||||||
|
- `MAX_UPLOAD_SIZE`: Maximum file upload size in bytes (default: 1GB)
|
||||||
|
- `MAX_SINGLE_FILE_SIZE`: Optional maximum size for a single file chunk
|
||||||
|
- Automatic file splitting for large PDFs when max_single_file_size is configured
|
||||||
|
- Split files are processed sequentially to prevent overwhelming the system
|
||||||
|
- Clear error messages referencing SECURITY_AUDIT.md for configuration details
|
||||||
|
|
||||||
|
**Configuration:**
|
||||||
|
```bash
|
||||||
|
# Set maximum upload size (default: 1GB)
|
||||||
|
MAX_UPLOAD_SIZE=1073741824
|
||||||
|
|
||||||
|
# Optional: Enable file splitting for large PDFs
|
||||||
|
MAX_SINGLE_FILE_SIZE=104857600 # 100MB chunks
|
||||||
|
```
|
||||||
|
|
||||||
|
**Security Benefits:**
|
||||||
|
- Prevents resource exhaustion from extremely large uploads
|
||||||
|
- Configurable limits allow adaptation to server capacity
|
||||||
|
- File splitting enables processing of large documents without memory issues
|
||||||
|
- Maintains support for large PDF files (up to 1GB by default) as required by use case
|
||||||
|
|
||||||
## Best Practices Implemented
|
## Best Practices Implemented
|
||||||
|
|
||||||
### Security Scanning with Bandit
|
### Security Scanning with Bandit
|
||||||
@@ -203,9 +230,10 @@ ftp = ftplib.FTP() # nosec B321 - Plaintext FTP intentional when configured
|
|||||||
- ✅ Authentication required on all sensitive endpoints (@require_login decorator)
|
- ✅ Authentication required on all sensitive endpoints (@require_login decorator)
|
||||||
- ✅ Path traversal protection in file uploads (basename sanitization)
|
- ✅ Path traversal protection in file uploads (basename sanitization)
|
||||||
- ✅ Unique filenames with UUID to prevent conflicts and overwrites
|
- ✅ Unique filenames with UUID to prevent conflicts and overwrites
|
||||||
|
- ✅ File upload size limits with configurable maximum (default: 1GB)
|
||||||
|
- ✅ Optional file splitting for large PDFs (when max_single_file_size is configured)
|
||||||
- ⏳ **TODO:** Implement rate limiting on API endpoints
|
- ⏳ **TODO:** Implement rate limiting on API endpoints
|
||||||
- ⏳ **TODO:** Add CSRF protection for state-changing operations
|
- ⏳ **TODO:** Add CSRF protection for state-changing operations
|
||||||
- ⏳ **TODO:** Implement request size limits ([#173](https://github.com/christianlouis/DocuElevate/issues/173))
|
|
||||||
- ⏳ **TODO:** Add comprehensive input sanitization for all user inputs ([#172](https://github.com/christianlouis/DocuElevate/issues/172))
|
- ⏳ **TODO:** Add comprehensive input sanitization for all user inputs ([#172](https://github.com/christianlouis/DocuElevate/issues/172))
|
||||||
- ⏳ **TODO:** Implement proper API key rotation mechanisms ([#168](https://github.com/christianlouis/DocuElevate/issues/168))
|
- ⏳ **TODO:** Implement proper API key rotation mechanisms ([#168](https://github.com/christianlouis/DocuElevate/issues/168))
|
||||||
|
|
||||||
@@ -230,7 +258,7 @@ ftp = ftplib.FTP() # nosec B321 - Plaintext FTP intentional when configured
|
|||||||
1. **Add security headers** - Improve browser-side security (HSTS, CSP, X-Frame-Options) ([#174](https://github.com/christianlouis/DocuElevate/issues/174))
|
1. **Add security headers** - Improve browser-side security (HSTS, CSP, X-Frame-Options) ([#174](https://github.com/christianlouis/DocuElevate/issues/174))
|
||||||
2. **Configure CORS properly** - Currently no CORS middleware configured ([#175](https://github.com/christianlouis/DocuElevate/issues/175))
|
2. **Configure CORS properly** - Currently no CORS middleware configured ([#175](https://github.com/christianlouis/DocuElevate/issues/175))
|
||||||
3. **Implement audit logging** - Track security-relevant events ([#170](https://github.com/christianlouis/DocuElevate/issues/170))
|
3. **Implement audit logging** - Track security-relevant events ([#170](https://github.com/christianlouis/DocuElevate/issues/170))
|
||||||
4. **Add file upload size limits** - Prevent resource exhaustion
|
4. ~~**Add file upload size limits**~~ ✅ Implemented - Configurable limits with 1GB default, optional file splitting
|
||||||
5. **Document security architecture** - Security design decisions
|
5. **Document security architecture** - Security design decisions
|
||||||
|
|
||||||
### Low Priority
|
### Low Priority
|
||||||
|
|||||||
+50
-4
@@ -685,13 +685,17 @@ async def ui_upload(request: Request, file: UploadFile = File(...)):
|
|||||||
# Log the mapping between original and safe filename
|
# Log the mapping between original and safe filename
|
||||||
logger.info(f"Saved uploaded file '{safe_filename}' as '{target_filename}'")
|
logger.info(f"Saved uploaded file '{safe_filename}' as '{target_filename}'")
|
||||||
|
|
||||||
# Check file size
|
# Check file size against configured maximum
|
||||||
file_size = os.path.getsize(target_path)
|
file_size = os.path.getsize(target_path)
|
||||||
max_size = 500 * 1024 * 1024 # 500MB
|
max_size = settings.max_upload_size
|
||||||
if file_size > max_size:
|
if file_size > max_size:
|
||||||
# Remove the file if it's too large
|
# Remove the file if it's too large
|
||||||
os.remove(target_path)
|
os.remove(target_path)
|
||||||
raise HTTPException(status_code=413, detail=f"File too large: {file_size} bytes (max {max_size} bytes)")
|
raise HTTPException(
|
||||||
|
status_code=413,
|
||||||
|
detail=f"File too large: {file_size} bytes (max {max_size} bytes). "
|
||||||
|
f"See SECURITY_AUDIT.md for configuration details.",
|
||||||
|
)
|
||||||
|
|
||||||
# Same set of allowed file types as in the IMAP task
|
# Same set of allowed file types as in the IMAP task
|
||||||
ALLOWED_MIME_TYPES = {
|
ALLOWED_MIME_TYPES = {
|
||||||
@@ -727,7 +731,49 @@ async def ui_upload(request: Request, file: UploadFile = File(...)):
|
|||||||
# Check if it's a PDF by extension or MIME type
|
# Check if it's a PDF by extension or MIME type
|
||||||
is_pdf = file_ext == ".pdf" or mime_type == "application/pdf"
|
is_pdf = file_ext == ".pdf" or mime_type == "application/pdf"
|
||||||
|
|
||||||
if is_pdf:
|
# Check if file splitting is needed (only for PDFs)
|
||||||
|
from app.utils.file_splitting import should_split_file
|
||||||
|
|
||||||
|
should_split = is_pdf and should_split_file(target_path, settings.max_single_file_size)
|
||||||
|
|
||||||
|
if should_split:
|
||||||
|
# File needs to be split before processing
|
||||||
|
from app.utils.file_splitting import split_pdf_by_size
|
||||||
|
|
||||||
|
try:
|
||||||
|
logger.info(
|
||||||
|
f"File {target_path} ({file_size} bytes) exceeds max_single_file_size "
|
||||||
|
f"({settings.max_single_file_size} bytes). Splitting..."
|
||||||
|
)
|
||||||
|
split_files = split_pdf_by_size(target_path, settings.max_single_file_size)
|
||||||
|
logger.info(f"Split {target_path} into {len(split_files)} parts")
|
||||||
|
|
||||||
|
# Queue each split file for processing
|
||||||
|
task_ids = []
|
||||||
|
for split_file in split_files:
|
||||||
|
split_filename = os.path.basename(split_file)
|
||||||
|
task = process_document.delay(split_file, original_filename=split_filename)
|
||||||
|
task_ids.append(task.id)
|
||||||
|
logger.info(f"Enqueued split PDF part for processing: {split_file}")
|
||||||
|
|
||||||
|
# Remove the original file after successful splitting
|
||||||
|
os.remove(target_path)
|
||||||
|
|
||||||
|
return {
|
||||||
|
"task_ids": task_ids,
|
||||||
|
"status": "queued",
|
||||||
|
"original_filename": safe_filename,
|
||||||
|
"stored_filename": target_filename,
|
||||||
|
"split_into_parts": len(split_files),
|
||||||
|
"message": f"File split into {len(split_files)} parts for processing",
|
||||||
|
}
|
||||||
|
except Exception as e:
|
||||||
|
logger.exception(f"Failed to split file {target_path}: {str(e)}")
|
||||||
|
# Fall back to processing the whole file
|
||||||
|
logger.warning(f"Falling back to processing whole file due to split error: {str(e)}")
|
||||||
|
should_split = False
|
||||||
|
|
||||||
|
if is_pdf and not should_split:
|
||||||
# If it's a PDF, process directly
|
# If it's a PDF, process directly
|
||||||
task = process_document.delay(target_path, original_filename=safe_filename)
|
task = process_document.delay(target_path, original_filename=safe_filename)
|
||||||
logger.info(f"Enqueued PDF for processing: {target_path}")
|
logger.info(f"Enqueued PDF for processing: {target_path}")
|
||||||
|
|||||||
@@ -164,6 +164,16 @@ class Settings(BaseSettings):
|
|||||||
default=True, description="Send notifications when files are successfully processed"
|
default=True, description="Send notifications when files are successfully processed"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# File upload size limits (for security - see SECURITY_AUDIT.md)
|
||||||
|
max_upload_size: int = Field(
|
||||||
|
default=1073741824, # 1GB in bytes (1024 * 1024 * 1024)
|
||||||
|
description="Maximum file upload size in bytes. Default: 1GB. Prevents resource exhaustion attacks.",
|
||||||
|
)
|
||||||
|
max_single_file_size: Optional[int] = Field(
|
||||||
|
default=None,
|
||||||
|
description="Maximum size for a single file chunk in bytes. If set and file exceeds this, it will be split into smaller chunks for processing. Default: None (no splitting).",
|
||||||
|
)
|
||||||
|
|
||||||
@validator("notification_urls", pre=True)
|
@validator("notification_urls", pre=True)
|
||||||
def parse_notification_urls(cls, v):
|
def parse_notification_urls(cls, v):
|
||||||
"""Parse notification URLs from string or list"""
|
"""Parse notification URLs from string or list"""
|
||||||
|
|||||||
@@ -0,0 +1,163 @@
|
|||||||
|
"""
|
||||||
|
Utility functions for splitting large PDF files into smaller chunks.
|
||||||
|
|
||||||
|
This module provides functionality to split PDF files that exceed a certain size
|
||||||
|
into smaller chunks for processing. Used when MAX_SINGLE_FILE_SIZE is configured.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import logging
|
||||||
|
import os
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import List, Optional
|
||||||
|
|
||||||
|
from PyPDF2 import PdfReader, PdfWriter
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
def split_pdf_by_size(
|
||||||
|
pdf_path: str,
|
||||||
|
max_size_bytes: int,
|
||||||
|
output_dir: Optional[str] = None
|
||||||
|
) -> List[str]:
|
||||||
|
"""
|
||||||
|
Split a PDF file into multiple smaller PDF files based on size constraints.
|
||||||
|
|
||||||
|
The function splits the PDF by distributing pages across multiple output files,
|
||||||
|
ensuring each output file stays under the specified size limit.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
pdf_path: Path to the PDF file to split
|
||||||
|
max_size_bytes: Maximum size for each output file in bytes
|
||||||
|
output_dir: Directory to save split files. If None, uses same directory as input file.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
List of paths to the generated PDF files (in order)
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
FileNotFoundError: If the input PDF file doesn't exist
|
||||||
|
ValueError: If max_size_bytes is too small to fit even one page
|
||||||
|
|
||||||
|
Example:
|
||||||
|
>>> split_files = split_pdf_by_size("large.pdf", 50 * 1024 * 1024) # 50MB max
|
||||||
|
>>> print(f"Split into {len(split_files)} files")
|
||||||
|
"""
|
||||||
|
# Validate input file exists
|
||||||
|
if not os.path.exists(pdf_path):
|
||||||
|
raise FileNotFoundError(f"PDF file not found: {pdf_path}")
|
||||||
|
|
||||||
|
# Determine output directory
|
||||||
|
if output_dir is None:
|
||||||
|
output_dir = os.path.dirname(pdf_path)
|
||||||
|
os.makedirs(output_dir, exist_ok=True)
|
||||||
|
|
||||||
|
# Read the input PDF
|
||||||
|
try:
|
||||||
|
reader = PdfReader(pdf_path)
|
||||||
|
total_pages = len(reader.pages)
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"Failed to read PDF {pdf_path}: {str(e)}")
|
||||||
|
raise ValueError(f"Invalid or corrupted PDF file: {str(e)}")
|
||||||
|
|
||||||
|
if total_pages == 0:
|
||||||
|
logger.warning(f"PDF {pdf_path} has no pages")
|
||||||
|
return []
|
||||||
|
|
||||||
|
# Get base filename without extension
|
||||||
|
base_name = os.path.splitext(os.path.basename(pdf_path))[0]
|
||||||
|
|
||||||
|
output_files = []
|
||||||
|
current_writer = PdfWriter()
|
||||||
|
current_page_count = 0
|
||||||
|
part_number = 1
|
||||||
|
|
||||||
|
logger.info(f"Splitting PDF {pdf_path} ({total_pages} pages) into chunks of max {max_size_bytes} bytes")
|
||||||
|
|
||||||
|
for page_num in range(total_pages):
|
||||||
|
# Add the page to current writer
|
||||||
|
page = reader.pages[page_num]
|
||||||
|
current_writer.add_page(page)
|
||||||
|
current_page_count += 1
|
||||||
|
|
||||||
|
# Write to temporary file to check size
|
||||||
|
temp_output_path = os.path.join(output_dir, f"{base_name}_part{part_number}_temp.pdf")
|
||||||
|
with open(temp_output_path, "wb") as temp_file:
|
||||||
|
current_writer.write(temp_file)
|
||||||
|
|
||||||
|
temp_size = os.path.getsize(temp_output_path)
|
||||||
|
|
||||||
|
# If adding this page exceeds the limit (and we have more than 1 page in current chunk)
|
||||||
|
# save the previous chunk and start a new one
|
||||||
|
if temp_size > max_size_bytes and current_page_count > 1:
|
||||||
|
# Remove the temporary file
|
||||||
|
os.remove(temp_output_path)
|
||||||
|
|
||||||
|
# Create a new writer without the last page
|
||||||
|
previous_writer = PdfWriter()
|
||||||
|
for prev_page_num in range(page_num - current_page_count + 1, page_num):
|
||||||
|
previous_writer.add_page(reader.pages[prev_page_num])
|
||||||
|
|
||||||
|
# Save the previous chunk
|
||||||
|
output_path = os.path.join(output_dir, f"{base_name}_part{part_number}.pdf")
|
||||||
|
with open(output_path, "wb") as output_file:
|
||||||
|
previous_writer.write(output_file)
|
||||||
|
|
||||||
|
output_files.append(output_path)
|
||||||
|
logger.info(f"Created chunk {part_number}: {output_path} ({current_page_count - 1} pages)")
|
||||||
|
|
||||||
|
# Start new chunk with current page
|
||||||
|
part_number += 1
|
||||||
|
current_writer = PdfWriter()
|
||||||
|
current_writer.add_page(page)
|
||||||
|
current_page_count = 1
|
||||||
|
elif temp_size > max_size_bytes and current_page_count == 1:
|
||||||
|
# Single page exceeds limit - this is a problem
|
||||||
|
# We'll keep it anyway but log a warning
|
||||||
|
logger.warning(
|
||||||
|
f"Single page (page {page_num + 1}) exceeds size limit "
|
||||||
|
f"({temp_size} > {max_size_bytes}). Keeping as separate file."
|
||||||
|
)
|
||||||
|
# Rename temp file to final name
|
||||||
|
output_path = os.path.join(output_dir, f"{base_name}_part{part_number}.pdf")
|
||||||
|
os.rename(temp_output_path, output_path)
|
||||||
|
output_files.append(output_path)
|
||||||
|
|
||||||
|
# Start new chunk
|
||||||
|
part_number += 1
|
||||||
|
current_writer = PdfWriter()
|
||||||
|
current_page_count = 0
|
||||||
|
else:
|
||||||
|
# Size is OK, remove temp file and continue
|
||||||
|
os.remove(temp_output_path)
|
||||||
|
|
||||||
|
# Save the last chunk if it has any pages
|
||||||
|
if current_page_count > 0:
|
||||||
|
output_path = os.path.join(output_dir, f"{base_name}_part{part_number}.pdf")
|
||||||
|
with open(output_path, "wb") as output_file:
|
||||||
|
current_writer.write(output_file)
|
||||||
|
output_files.append(output_path)
|
||||||
|
logger.info(f"Created final chunk {part_number}: {output_path} ({current_page_count} pages)")
|
||||||
|
|
||||||
|
logger.info(f"Successfully split PDF into {len(output_files)} files")
|
||||||
|
return output_files
|
||||||
|
|
||||||
|
|
||||||
|
def should_split_file(file_path: str, max_single_file_size: Optional[int]) -> bool:
|
||||||
|
"""
|
||||||
|
Determine if a file should be split based on its size and configuration.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
file_path: Path to the file to check
|
||||||
|
max_single_file_size: Maximum single file size in bytes, or None to disable splitting
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
True if file should be split, False otherwise
|
||||||
|
"""
|
||||||
|
if max_single_file_size is None:
|
||||||
|
return False
|
||||||
|
|
||||||
|
if not os.path.exists(file_path):
|
||||||
|
return False
|
||||||
|
|
||||||
|
file_size = os.path.getsize(file_path)
|
||||||
|
return file_size > max_single_file_size
|
||||||
@@ -31,6 +31,42 @@ Control how the `/processall` endpoint handles large batches of files to prevent
|
|||||||
- Total queue time: (25-1) × 3 = 72 seconds
|
- Total queue time: (25-1) × 3 = 72 seconds
|
||||||
- Prevents API rate limit issues and ensures smooth processing
|
- Prevents API rate limit issues and ensures smooth processing
|
||||||
|
|
||||||
|
### File Upload Size Limits
|
||||||
|
|
||||||
|
**Security Feature**: Control file upload sizes to prevent resource exhaustion attacks. See [SECURITY_AUDIT.md](../SECURITY_AUDIT.md#5-file-upload-size-limits) for security details.
|
||||||
|
|
||||||
|
| **Variable** | **Description** | **Default** |
|
||||||
|
|---------------------------|--------------------------------------------------------------------------------------------------------------|---------------|
|
||||||
|
| `MAX_UPLOAD_SIZE` | Maximum file upload size in bytes. Files exceeding this limit are rejected. | `1073741824` (1GB) |
|
||||||
|
| `MAX_SINGLE_FILE_SIZE` | Optional: Maximum size for a single file chunk in bytes. Files exceeding this are split into smaller parts. | `None` (no splitting) |
|
||||||
|
|
||||||
|
**Configuration Examples:**
|
||||||
|
|
||||||
|
```bash
|
||||||
|
# Default: Allow up to 1GB uploads, no splitting
|
||||||
|
MAX_UPLOAD_SIZE=1073741824
|
||||||
|
|
||||||
|
# Conservative: 100MB max, split files over 50MB
|
||||||
|
MAX_UPLOAD_SIZE=104857600
|
||||||
|
MAX_SINGLE_FILE_SIZE=52428800
|
||||||
|
|
||||||
|
# Large files: 2GB max, split files over 500MB
|
||||||
|
MAX_UPLOAD_SIZE=2147483648
|
||||||
|
MAX_SINGLE_FILE_SIZE=524288000
|
||||||
|
```
|
||||||
|
|
||||||
|
**File Splitting Behavior:**
|
||||||
|
- When `MAX_SINGLE_FILE_SIZE` is configured and a PDF exceeds this size, it is automatically split into smaller chunks
|
||||||
|
- Each chunk is processed sequentially as a separate task
|
||||||
|
- Only works for PDF files (images and office documents are converted to PDF first)
|
||||||
|
- Original file is removed after successful splitting
|
||||||
|
- Useful for very large PDFs to prevent memory issues during processing
|
||||||
|
|
||||||
|
**Use Cases:**
|
||||||
|
- **Default (1GB, no splitting)**: Suitable for most deployments handling typical documents
|
||||||
|
- **With splitting**: Recommended for servers with limited memory or when processing very large scanned documents
|
||||||
|
- **Higher limits**: For environments specifically designed to handle large architectural plans, books, or scanned archives
|
||||||
|
|
||||||
### IMAP Configuration
|
### IMAP Configuration
|
||||||
|
|
||||||
DocuElevate can monitor multiple IMAP mailboxes for document attachments. Each mailbox uses a numbered prefix (e.g., `IMAP1_`, `IMAP2_`).
|
DocuElevate can monitor multiple IMAP mailboxes for document attachments. Each mailbox uses a numbered prefix (e.g., `IMAP1_`, `IMAP2_`).
|
||||||
|
|||||||
@@ -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."""
|
"""Tests for handling invalid or problematic file uploads."""
|
||||||
|
|
||||||
def test_upload_file_too_large(self, client: TestClient):
|
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)
|
# Create a large file content (mock it to avoid memory issues)
|
||||||
large_content = b"x" * 1024 # 1KB for testing
|
large_content = b"x" * 1024 # 1KB for testing
|
||||||
|
|
||||||
with patch("os.path.getsize") as mock_getsize:
|
with patch("os.path.getsize") as mock_getsize:
|
||||||
# Mock the file size to be over 500MB
|
# Mock the file size to be over the configured limit
|
||||||
mock_getsize.return_value = 501 * 1024 * 1024 # 501MB
|
mock_getsize.return_value = settings.max_upload_size + 1
|
||||||
|
|
||||||
response = client.post(
|
response = client.post(
|
||||||
"/api/ui-upload", files={"file": ("huge.pdf", io.BytesIO(large_content), "application/pdf")}
|
"/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 response.status_code == 413 # Request Entity Too Large
|
||||||
assert "too large" in response.json()["detail"].lower()
|
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):
|
def test_upload_executable_file(self, client: TestClient, mock_celery_tasks):
|
||||||
"""Test that executable files are handled (attempted conversion)."""
|
"""Test that executable files are handled (attempted conversion)."""
|
||||||
@@ -411,3 +414,129 @@ class TestUploadMimeTypeDetection:
|
|||||||
assert response.status_code == 200
|
assert response.status_code == 200
|
||||||
# Should route to convert_to_pdf based on .jpg extension
|
# Should route to convert_to_pdf based on .jpg extension
|
||||||
mock_celery_tasks["convert_to_pdf"].assert_called_once()
|
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