fix(cv): harden document parsing
Upgrade and hash-lock upload-facing parser dependencies, reject resource-heavy or mismatched inputs, remove unsafe backend binary fallbacks, and prevent internal parser failures from leaking to users.
This commit is contained in:
cesnimda
+99
-8
@@ -17,6 +17,7 @@ import torch
|
||||
import pytesseract
|
||||
import threading
|
||||
import time
|
||||
import zipfile
|
||||
from urllib import request as urllib_request
|
||||
from urllib.error import URLError, HTTPError
|
||||
from contextvars import ContextVar
|
||||
@@ -91,7 +92,17 @@ async def require_service_token(request: Request, call_next):
|
||||
MODEL_NAME = "sshleifer/distilbart-cnn-12-6"
|
||||
MAX_INPUT_CHARS = 20000
|
||||
MAX_CONTEXT_CHARS = 2200
|
||||
MAX_EXTRACT_FILE_BYTES = 8 * 1024 * 1024
|
||||
MAX_EXTRACT_FILE_BYTES = 5 * 1024 * 1024
|
||||
MAX_EXTRACT_CHARS = 200_000
|
||||
MAX_PDF_PAGES = 40
|
||||
MAX_DOCX_ENTRIES = 256
|
||||
MAX_DOCX_UNCOMPRESSED_BYTES = 32 * 1024 * 1024
|
||||
MAX_DOCX_ENTRY_BYTES = 8 * 1024 * 1024
|
||||
MAX_DOCX_COMPRESSION_RATIO = 100
|
||||
MAX_IMAGE_DIMENSION = 12_000
|
||||
MAX_IMAGE_PIXELS = 40_000_000
|
||||
MAX_PDF_PAGE_PIXELS = 12_000_000
|
||||
MAX_PDF_OCR_PIXELS = 120_000_000
|
||||
OCR_LANGUAGES = "eng"
|
||||
IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp"}
|
||||
OLLAMA_BASE_URL = os.getenv("OLLAMA_BASE_URL", "http://127.0.0.1:11434").rstrip("/")
|
||||
@@ -1068,36 +1079,110 @@ def _ocr_image(image: Image.Image) -> str:
|
||||
return _normalize_text(text)
|
||||
|
||||
|
||||
def _check_extracted_size(text: str) -> str:
|
||||
if len(text) > MAX_EXTRACT_CHARS:
|
||||
raise HTTPException(status_code=422, detail="The document contains too much extracted text.")
|
||||
return text
|
||||
|
||||
|
||||
def _check_image_size(image: Image.Image, *, pixel_limit: int | None = None) -> None:
|
||||
pixel_limit = MAX_IMAGE_PIXELS if pixel_limit is None else pixel_limit
|
||||
width, height = image.size
|
||||
if width <= 0 or height <= 0 or width > MAX_IMAGE_DIMENSION or height > MAX_IMAGE_DIMENSION:
|
||||
raise HTTPException(status_code=422, detail="The document image dimensions are not supported.")
|
||||
if width * height > pixel_limit:
|
||||
raise HTTPException(status_code=422, detail="The document image contains too many pixels.")
|
||||
if getattr(image, "n_frames", 1) != 1:
|
||||
raise HTTPException(status_code=422, detail="Multi-frame document images are not supported.")
|
||||
|
||||
|
||||
def _validate_docx_container(data: bytes) -> None:
|
||||
try:
|
||||
with zipfile.ZipFile(io.BytesIO(data)) as archive:
|
||||
entries = archive.infolist()
|
||||
names = {entry.filename for entry in entries}
|
||||
if "[Content_Types].xml" not in names or "word/document.xml" not in names:
|
||||
raise HTTPException(status_code=400, detail="The uploaded file does not match its DOCX extension.")
|
||||
if len(entries) > MAX_DOCX_ENTRIES:
|
||||
raise HTTPException(status_code=422, detail="The DOCX contains too many files.")
|
||||
total_size = 0
|
||||
for entry in entries:
|
||||
total_size += entry.file_size
|
||||
if entry.file_size > MAX_DOCX_ENTRY_BYTES:
|
||||
raise HTTPException(status_code=422, detail="The DOCX contains an oversized file.")
|
||||
if entry.file_size > 0 and entry.compress_size == 0:
|
||||
raise HTTPException(status_code=422, detail="The DOCX compression ratio is not supported.")
|
||||
if entry.compress_size > 0 and entry.file_size / entry.compress_size > MAX_DOCX_COMPRESSION_RATIO:
|
||||
raise HTTPException(status_code=422, detail="The DOCX compression ratio is not supported.")
|
||||
if total_size > MAX_DOCX_UNCOMPRESSED_BYTES:
|
||||
raise HTTPException(status_code=422, detail="The DOCX expands beyond the supported size.")
|
||||
except HTTPException:
|
||||
raise
|
||||
except (zipfile.BadZipFile, OSError) as exc:
|
||||
raise HTTPException(status_code=400, detail="The uploaded file does not match its DOCX extension.") from exc
|
||||
|
||||
|
||||
def _validate_file_signature(extension: str, data: bytes) -> None:
|
||||
matches = {
|
||||
".pdf": data.startswith(b"%PDF-"),
|
||||
".docx": data.startswith(b"PK"),
|
||||
".png": data.startswith(b"\x89PNG\r\n\x1a\n"),
|
||||
".jpg": data.startswith(b"\xff\xd8\xff"),
|
||||
".jpeg": data.startswith(b"\xff\xd8\xff"),
|
||||
".webp": len(data) >= 12 and data.startswith(b"RIFF") and data[8:12] == b"WEBP",
|
||||
}
|
||||
if extension in matches and not matches[extension]:
|
||||
raise HTTPException(status_code=400, detail=f"The uploaded file does not match its {extension[1:].upper()} extension.")
|
||||
if extension in {".txt", ".md"} and b"\x00" in data:
|
||||
raise HTTPException(status_code=400, detail="The uploaded text file contains binary data.")
|
||||
|
||||
|
||||
def _extract_pdf_text(data: bytes) -> tuple[str, bool, int]:
|
||||
page_count = 0
|
||||
extracted_pages = []
|
||||
try:
|
||||
reader = PdfReader(io.BytesIO(data))
|
||||
page_count = len(reader.pages)
|
||||
if page_count > MAX_PDF_PAGES:
|
||||
raise HTTPException(status_code=422, detail="The PDF contains too many pages.")
|
||||
for page in reader.pages:
|
||||
try:
|
||||
extracted_pages.append(page.extract_text(extraction_mode="layout") or "")
|
||||
except TypeError:
|
||||
extracted_pages.append(page.extract_text() or "")
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception:
|
||||
extracted_pages = []
|
||||
|
||||
text = _normalize_text("\n".join(extracted_pages))
|
||||
text = _check_extracted_size(_normalize_text("\n".join(extracted_pages)))
|
||||
if len(text) >= 80:
|
||||
return text, False, page_count
|
||||
|
||||
doc = fitz.open(stream=data, filetype="pdf")
|
||||
page_count = doc.page_count
|
||||
if page_count > MAX_PDF_PAGES:
|
||||
doc.close()
|
||||
raise HTTPException(status_code=422, detail="The PDF contains too many pages.")
|
||||
ocr_pages = []
|
||||
total_pixels = 0
|
||||
for page in doc:
|
||||
pix = page.get_pixmap(matrix=fitz.Matrix(2, 2), alpha=False)
|
||||
page_pixels = pix.width * pix.height
|
||||
total_pixels += page_pixels
|
||||
if page_pixels > MAX_PDF_PAGE_PIXELS or total_pixels > MAX_PDF_OCR_PIXELS:
|
||||
doc.close()
|
||||
raise HTTPException(status_code=422, detail="The PDF requires too much OCR processing.")
|
||||
image = Image.open(io.BytesIO(pix.tobytes("png")))
|
||||
_check_image_size(image, pixel_limit=MAX_PDF_PAGE_PIXELS)
|
||||
ocr_pages.append(_ocr_image(image))
|
||||
_check_extracted_size("\n".join(ocr_pages))
|
||||
doc.close()
|
||||
return _normalize_text("\n".join(ocr_pages)), True, page_count
|
||||
return _check_extracted_size(_normalize_text("\n".join(ocr_pages))), True, page_count
|
||||
|
||||
|
||||
def _extract_docx_text(data: bytes) -> str:
|
||||
_validate_docx_container(data)
|
||||
document = Document(io.BytesIO(data))
|
||||
parts = []
|
||||
blocks = document.iter_inner_content() if hasattr(document, "iter_inner_content") else document.paragraphs
|
||||
@@ -1113,11 +1198,15 @@ def _extract_docx_text(data: bytes) -> str:
|
||||
if "role" in style and parts:
|
||||
parts.append("")
|
||||
parts.append(f"- {text}" if "bullet" in style else text)
|
||||
return _normalize_text("\n".join(parts))
|
||||
return _check_extracted_size(_normalize_text("\n".join(parts)))
|
||||
|
||||
|
||||
def _extract_plain_text(data: bytes) -> str:
|
||||
return _normalize_text(data.decode("utf-8", errors="ignore"))
|
||||
try:
|
||||
decoded = data.decode("utf-8-sig")
|
||||
except UnicodeDecodeError as exc:
|
||||
raise HTTPException(status_code=400, detail="The uploaded text file is not valid UTF-8.") from exc
|
||||
return _check_extracted_size(_normalize_text(decoded))
|
||||
|
||||
|
||||
@app.post("/extract-text")
|
||||
@@ -1135,6 +1224,9 @@ async def extract_text(file: UploadFile = File(...)):
|
||||
raise HTTPException(status_code=400, detail="The uploaded file was empty.")
|
||||
if len(data) > MAX_EXTRACT_FILE_BYTES:
|
||||
raise HTTPException(status_code=400, detail="The uploaded file is too large for AI extraction.")
|
||||
if extension not in {".txt", ".md", ".docx", ".pdf", *IMAGE_EXTENSIONS}:
|
||||
raise HTTPException(status_code=400, detail="This file type is not supported for AI extraction.")
|
||||
_validate_file_signature(extension, data)
|
||||
|
||||
try:
|
||||
if extension in {".txt", ".md"}:
|
||||
@@ -1149,15 +1241,14 @@ async def extract_text(file: UploadFile = File(...)):
|
||||
text, ocr_used, page_count = _extract_pdf_text(data)
|
||||
elif extension in IMAGE_EXTENSIONS:
|
||||
image = Image.open(io.BytesIO(data))
|
||||
_check_image_size(image)
|
||||
text = _ocr_image(image)
|
||||
ocr_used = True
|
||||
page_count = 1
|
||||
else:
|
||||
raise HTTPException(status_code=400, detail="This file type is not supported for AI extraction.")
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as exc:
|
||||
raise HTTPException(status_code=500, detail=f"AI extraction failed: {exc}") from exc
|
||||
raise HTTPException(status_code=422, detail="Document extraction failed.") from exc
|
||||
|
||||
if not text:
|
||||
raise HTTPException(status_code=422, detail="AI extraction did not find readable text in the uploaded file.")
|
||||
|
||||
Reference in New Issue
Block a user