mirror of
https://github.com/escalante29/WealthySmart.git
synced 2026-07-17 11:28:47 +02:00
Harden PDF uploads: size cap, magic-byte check, non-blocking parse
read_pdf_upload() bounds reads at 10 MB and requires the %PDF magic before anything reaches pdftotext (SEC-06); rejections land in the per-file errors list with the filename. Parsing now runs via asyncio.to_thread so a slow PDF no longer blocks the event loop for up to 30s (BE-14). Review notes: temp files were already context-managed (BE-09 overstated) and the pension batch already commits atomically (BE-10 overstated). 69 tests green. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
@@ -2,11 +2,14 @@ from datetime import datetime
|
|||||||
from dateutil.relativedelta import relativedelta
|
from dateutil.relativedelta import relativedelta
|
||||||
from typing import Optional
|
from typing import Optional
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
|
||||||
from fastapi import APIRouter, Depends, Query, UploadFile
|
from fastapi import APIRouter, Depends, Query, UploadFile
|
||||||
from pydantic import BaseModel
|
from pydantic import BaseModel
|
||||||
from sqlmodel import Session, select
|
from sqlmodel import Session, select
|
||||||
|
|
||||||
from app.auth import get_current_user
|
from app.auth import get_current_user
|
||||||
|
from app.uploads import read_pdf_upload
|
||||||
from app.db import get_session
|
from app.db import get_session
|
||||||
from app.models.models import (
|
from app.models.models import (
|
||||||
Category,
|
Category,
|
||||||
@@ -204,8 +207,9 @@ async def upload_municipal_receipt(
|
|||||||
errors: list[str] = []
|
errors: list[str] = []
|
||||||
|
|
||||||
try:
|
try:
|
||||||
pdf_bytes = await file.read()
|
pdf_bytes = await read_pdf_upload(file)
|
||||||
data = extract_municipal_receipt(pdf_bytes, filename)
|
# pdftotext runs as a subprocess; keep the event loop free
|
||||||
|
data = await asyncio.to_thread(extract_municipal_receipt, pdf_bytes, filename)
|
||||||
except ValueError as e:
|
except ValueError as e:
|
||||||
return MunicipalReceiptUploadResult(imported=0, updated=0, errors=[str(e)])
|
return MunicipalReceiptUploadResult(imported=0, updated=0, errors=[str(e)])
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
|
|||||||
@@ -1,10 +1,13 @@
|
|||||||
from datetime import date
|
from datetime import date
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
|
||||||
from fastapi import APIRouter, Depends, UploadFile
|
from fastapi import APIRouter, Depends, UploadFile
|
||||||
from pydantic import BaseModel
|
from pydantic import BaseModel
|
||||||
from sqlmodel import Session, select
|
from sqlmodel import Session, select
|
||||||
|
|
||||||
from app.auth import get_current_user
|
from app.auth import get_current_user
|
||||||
|
from app.uploads import read_pdf_upload
|
||||||
from app.db import get_session
|
from app.db import get_session
|
||||||
from app.models.models import Bank, PensionSnapshot, PensionSnapshotRead
|
from app.models.models import Bank, PensionSnapshot, PensionSnapshotRead
|
||||||
from app.services.pension_pdf import parse_pension_pdf
|
from app.services.pension_pdf import parse_pension_pdf
|
||||||
@@ -115,8 +118,11 @@ async def upload_pension_pdfs(
|
|||||||
for file in files:
|
for file in files:
|
||||||
filename = file.filename or "unknown.pdf"
|
filename = file.filename or "unknown.pdf"
|
||||||
try:
|
try:
|
||||||
pdf_bytes = await file.read()
|
pdf_bytes = await read_pdf_upload(file)
|
||||||
fund_snapshots = parse_pension_pdf(pdf_bytes, filename)
|
# pdftotext runs as a subprocess (up to 30s); keep the event loop free
|
||||||
|
fund_snapshots = await asyncio.to_thread(
|
||||||
|
parse_pension_pdf, pdf_bytes, filename
|
||||||
|
)
|
||||||
except ValueError as e:
|
except ValueError as e:
|
||||||
errors.append(str(e))
|
errors.append(str(e))
|
||||||
continue
|
continue
|
||||||
|
|||||||
22
backend/app/uploads.py
Normal file
22
backend/app/uploads.py
Normal file
@@ -0,0 +1,22 @@
|
|||||||
|
from fastapi import UploadFile
|
||||||
|
|
||||||
|
MAX_UPLOAD_BYTES = 10 * 1024 * 1024 # 10 MB — statements are a few hundred KB
|
||||||
|
|
||||||
|
|
||||||
|
async def read_pdf_upload(
|
||||||
|
file: UploadFile, max_bytes: int = MAX_UPLOAD_BYTES
|
||||||
|
) -> bytes:
|
||||||
|
"""Read an uploaded PDF with a size cap and magic-byte check.
|
||||||
|
|
||||||
|
Raises ValueError with a user-readable, filename-prefixed message —
|
||||||
|
upload endpoints surface these in their per-file `errors` list.
|
||||||
|
"""
|
||||||
|
filename = file.filename or "unknown.pdf"
|
||||||
|
data = await file.read(max_bytes + 1)
|
||||||
|
if len(data) > max_bytes:
|
||||||
|
raise ValueError(
|
||||||
|
f"{filename}: file exceeds the {max_bytes // (1024 * 1024)} MB limit"
|
||||||
|
)
|
||||||
|
if not data.startswith(b"%PDF"):
|
||||||
|
raise ValueError(f"{filename}: not a PDF file")
|
||||||
|
return data
|
||||||
41
backend/tests/test_uploads.py
Normal file
41
backend/tests/test_uploads.py
Normal file
@@ -0,0 +1,41 @@
|
|||||||
|
"""Upload validation: size cap and PDF magic-byte check (SEC-06)."""
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
from io import BytesIO
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from fastapi import UploadFile
|
||||||
|
|
||||||
|
from app.uploads import read_pdf_upload
|
||||||
|
|
||||||
|
|
||||||
|
def make_upload(data: bytes, filename: str = "test.pdf") -> UploadFile:
|
||||||
|
return UploadFile(file=BytesIO(data), filename=filename)
|
||||||
|
|
||||||
|
|
||||||
|
def test_valid_pdf_accepted():
|
||||||
|
data = b"%PDF-1.7 fake content"
|
||||||
|
out = asyncio.run(read_pdf_upload(make_upload(data)))
|
||||||
|
assert out == data
|
||||||
|
|
||||||
|
|
||||||
|
def test_non_pdf_rejected_with_filename():
|
||||||
|
with pytest.raises(ValueError, match=r"evil\.pdf: not a PDF"):
|
||||||
|
asyncio.run(read_pdf_upload(make_upload(b"MZ\x90\x00binary", "evil.pdf")))
|
||||||
|
|
||||||
|
|
||||||
|
def test_oversized_rejected():
|
||||||
|
big = b"%PDF" + b"x" * 100
|
||||||
|
with pytest.raises(ValueError, match="exceeds the 0 MB limit"):
|
||||||
|
asyncio.run(read_pdf_upload(make_upload(big), max_bytes=50))
|
||||||
|
|
||||||
|
|
||||||
|
def test_exact_limit_accepted():
|
||||||
|
data = b"%PDF" + b"x" * 46 # exactly 50 bytes
|
||||||
|
out = asyncio.run(read_pdf_upload(make_upload(data), max_bytes=50))
|
||||||
|
assert len(out) == 50
|
||||||
|
|
||||||
|
|
||||||
|
def test_empty_file_rejected():
|
||||||
|
with pytest.raises(ValueError, match="not a PDF"):
|
||||||
|
asyncio.run(read_pdf_upload(make_upload(b"")))
|
||||||
Reference in New Issue
Block a user