Rescue Paperjet v1 implementation

backend/app/services/storage.py

@@ -1,53 +1,71 @@
-"""Document storage service."""
+"""Document and binary storage helpers."""
 
-import shutil
 import uuid
-import pymupdf
 from pathlib import Path
-from fastapi import UploadFile, HTTPException
+
+import pymupdf
+from fastapi import HTTPException, UploadFile
 
 from app.config import settings
 
-def validate_pdf(filepath: Path) -> None:
-    """Validate PDF magic bytes and PyMuPDF openability."""
-    with open(filepath, "rb") as f:
-        header = f.read(5)
-        if header != b"%PDF-":
+_COPY_CHUNK_SIZE = 1024 * 1024
+
+
+def validate_pdf(filepath: Path) -> int:
+    """Validate a PDF's magic bytes and return its page count."""
+    with filepath.open("rb") as file:
+        if file.read(5) != b"%PDF-":
             raise ValueError("Not a valid PDF file (missing magic bytes)")
 
     try:
-        doc = pymupdf.open(filepath)
-        doc.close()
-    except Exception as e:
-        raise ValueError(f"Failed to open PDF with PyMuPDF: {e}")
+        with pymupdf.open(filepath) as document:
+            page_count = len(document)
+    except Exception as err:
+        raise ValueError(f"Failed to open PDF with PyMuPDF: {err}") from err
+
+    if page_count == 0:
+        raise ValueError("PDF does not contain any pages")
+    return page_count
+
+
+def safe_filename(filename: str | None) -> str:
+    """Return a display-safe basename without allowing path components."""
+    name = Path(filename or "Untitled.pdf").name.strip()
+    return name or "Untitled.pdf"
+
 
 async def save_upload_file(upload_file: UploadFile) -> tuple[str, int, int]:
-    """Save an uploaded file to disk and validate it."""
+    """Stream an uploaded PDF to storage, enforcing the configured size limit."""
     file_id = str(uuid.uuid4())
+    settings.PDF_STORAGE_PATH.mkdir(parents=True, exist_ok=True)
     filepath = settings.PDF_STORAGE_PATH / f"{file_id}.pdf"
+    max_bytes = settings.MAX_UPLOAD_MB * 1024 * 1024
+    total_bytes = 0
 
     try:
-        with open(filepath, "wb") as f:
-            shutil.copyfileobj(upload_file.file, f)
-    except Exception as e:
-        if filepath.exists():
-            filepath.unlink()
-        raise HTTPException(status_code=500, detail="Failed to save file")
+        with filepath.open("wb") as destination:
+            while chunk := upload_file.file.read(_COPY_CHUNK_SIZE):
+                total_bytes += len(chunk)
+                if total_bytes > max_bytes:
+                    raise HTTPException(status_code=413, detail="File too large")
+                destination.write(chunk)
+    except HTTPException:
+        filepath.unlink(missing_ok=True)
+        raise
+    except Exception as err:
+        filepath.unlink(missing_ok=True)
+        raise HTTPException(status_code=500, detail="Failed to save file") from err
 
-    page_count = 0
     try:
-        doc = pymupdf.open(filepath)
-        page_count = len(doc)
-        doc.close()
-    except Exception as e:
-        filepath.unlink()
-        raise HTTPException(status_code=400, detail=f"Failed to open PDF with PyMuPDF: {e}")
-    
-    size = filepath.stat().st_size
-    return file_id, size, page_count
+        page_count = validate_pdf(filepath)
+    except ValueError as err:
+        filepath.unlink(missing_ok=True)
+        raise HTTPException(status_code=400, detail=str(err)) from err
+
+    return file_id, total_bytes, page_count
+
 
 def delete_pdf_file(file_id: str) -> None:
     """Delete a PDF file from storage."""
     filepath = settings.PDF_STORAGE_PATH / f"{file_id}.pdf"
-    if filepath.exists():
-        filepath.unlink()
+    filepath.unlink(missing_ok=True)