Rescue Paperjet v1 implementation

backend/app/api/v1/assets.py

@@ -1,58 +1,101 @@
-from fastapi import APIRouter, Depends, UploadFile, File, HTTPException
-from fastapi.responses import FileResponse
-from sqlalchemy.orm import Session
+"""Authenticated binary assets used by annotations."""
+
 import uuid
-import shutil
 from pathlib import Path
 
-from app.db import get_db
-from app.auth.dependencies import get_current_user
+from fastapi import APIRouter, Depends, File, HTTPException, UploadFile
+from fastapi.responses import FileResponse
+from sqlalchemy import select
+from sqlalchemy.orm import Session
+
+from app.auth.dependencies import get_current_user, verify_csrf
 from app.config import settings
+from app.db import get_db
+from app.models.document import Document
 
 router = APIRouter(prefix="/documents", tags=["Assets"])
+_ASSET_CHUNK_SIZE = 1024 * 1024
+_MAX_ASSET_BYTES = 25 * 1024 * 1024
 
-def get_asset_path(document_id: str, ref: str) -> Path:
-    # Ensure directory exists
-    dir_path = settings.PDF_STORAGE_PATH / "assets" / document_id
-    dir_path.mkdir(parents=True, exist_ok=True)
-    return dir_path / ref
 
-@router.post("/{id}/assets")
+def get_asset_path(document_id: str, ref: str, *, create: bool = False) -> Path:
+    """Return a safe asset path; only writes are allowed to create directories."""
+    try:
+        asset_ref = str(uuid.UUID(ref))
+    except ValueError as err:
+        raise HTTPException(status_code=404, detail="Asset not found") from err
+
+    directory = settings.PDF_STORAGE_PATH / "assets" / document_id
+    if create:
+        directory.mkdir(parents=True, exist_ok=True)
+    return directory / asset_ref
+
+
+def _is_supported_image(filepath: Path) -> bool:
+    with filepath.open("rb") as file:
+        header = file.read(12)
+    return (
+        header.startswith(b"\x89PNG\r\n\x1a\n")
+        or header.startswith(b"\xff\xd8\xff")
+        or header.startswith((b"GIF87a", b"GIF89a"))
+        or (header.startswith(b"RIFF") and header[8:12] == b"WEBP")
+    )
+
+
+@router.post("/{id}/assets", dependencies=[Depends(verify_csrf)])
 async def upload_asset(
     id: str,
     file: UploadFile = File(...),
     user_id: int = Depends(get_current_user),
-    db: Session = Depends(get_db)
-):
-    """Upload a binary asset (like an image or signature) for a document."""
-    # In a real app we'd verify the user owns the document here
-    # For this single-user app, we trust the ID
-
-    if not file.content_type or not file.content_type.startswith("image/"):
-        raise HTTPException(status_code=400, detail="Only image assets are supported")
+    db: Session = Depends(get_db),
+) -> dict[str, str]:
+    """Upload a PNG, JPEG, GIF, or WebP annotation asset."""
+    document = db.scalar(select(Document).where(Document.id == id))
+    if not document:
+        raise HTTPException(status_code=404, detail="Document not found")
 
     ref = str(uuid.uuid4())
-    filepath = get_asset_path(id, ref)
-    
-    try:
-        with open(filepath, "wb") as f:
-            shutil.copyfileobj(file.file, f)
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Failed to save asset: {e}")
+    filepath = get_asset_path(id, ref, create=True)
+    total_bytes = 0
+
+    try:
+        with filepath.open("wb") as destination:
+            while chunk := file.file.read(_ASSET_CHUNK_SIZE):
+                total_bytes += len(chunk)
+                if total_bytes > _MAX_ASSET_BYTES:
+                    raise HTTPException(status_code=413, detail="Image asset is 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 asset") from err
+
+    if not _is_supported_image(filepath):
+        filepath.unlink(missing_ok=True)
+        raise HTTPException(
+            status_code=400,
+            detail="Only PNG, JPEG, GIF, and WebP images are supported",
+        )
+
+    return {"ref": ref, "url": f"/api/v1/documents/{id}/assets/{ref}"}
 
-    return {
-        "ref": ref,
-        "url": f"/api/v1/documents/{id}/assets/{ref}"
-    }
 
 @router.get("/{id}/assets/{ref}")
 async def get_asset(
     id: str,
     ref: str,
-):
-    """Serve a binary asset."""
+    db: Session = Depends(get_db),
+    user_id: int = Depends(get_current_user),
+) -> FileResponse:
+    """Serve an annotation asset to an authenticated session."""
+    document = db.scalar(select(Document).where(Document.id == id))
+    if not document:
+        raise HTTPException(status_code=404, detail="Document not found")
+
     filepath = get_asset_path(id, ref)
     if not filepath.exists() or not filepath.is_file():
         raise HTTPException(status_code=404, detail="Asset not found")
-        
+
     return FileResponse(filepath)