fix(cv): isolate document parsing
Run untrusted document decoders in a secret-free, resource-bounded child process and terminate its process tree on deadline. Harden the production container and enforce parser and lint gates in CI.
This commit is contained in:
cesnimda
+158
-141
@@ -3,18 +3,20 @@ from fastapi.responses import JSONResponse
|
||||
from pydantic import BaseModel, Field
|
||||
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
|
||||
from cachetools import TTLCache
|
||||
from PIL import Image
|
||||
from pypdf import PdfReader
|
||||
from docx import Document
|
||||
import fitz
|
||||
import asyncio
|
||||
import hashlib
|
||||
import hmac
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import re
|
||||
import shutil
|
||||
import signal
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import torch
|
||||
import pytesseract
|
||||
import threading
|
||||
import time
|
||||
import zipfile
|
||||
@@ -94,17 +96,22 @@ MAX_INPUT_CHARS = 20000
|
||||
MAX_CONTEXT_CHARS = 2200
|
||||
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"}
|
||||
PARSER_CHILD_PATH = Path(__file__).with_name("parser_child.py")
|
||||
PARSER_TIMEOUT_SECONDS = max(5, min(int(os.getenv("PARSER_TIMEOUT_SECONDS", "25")), 120))
|
||||
PARSER_CPU_SECONDS = max(2, min(int(os.getenv("PARSER_CPU_SECONDS", "20")), PARSER_TIMEOUT_SECONDS))
|
||||
PARSER_MEMORY_BYTES = max(256, min(int(os.getenv("PARSER_MEMORY_MIB", "512")), 2048)) * 1024 * 1024
|
||||
PARSER_MAX_RESULT_BYTES = 1024 * 1024
|
||||
_parser_capacity = threading.BoundedSemaphore(value=1)
|
||||
PARSER_TEMP_ROOT = Path(os.getenv("PARSER_TEMP_ROOT", tempfile.gettempdir()), "jobtracker-cv-parser").resolve()
|
||||
PARSER_STALE_WORK_SECONDS = max(300, min(int(os.getenv("PARSER_STALE_WORK_SECONDS", "3600")), 86_400))
|
||||
_parser_temp_lock = threading.Lock()
|
||||
_parser_temp_ready = False
|
||||
OLLAMA_BASE_URL = os.getenv("OLLAMA_BASE_URL", "http://127.0.0.1:11434").rstrip("/")
|
||||
OLLAMA_MODEL = os.getenv("OLLAMA_MODEL", "")
|
||||
|
||||
@@ -280,6 +287,11 @@ async def health():
|
||||
"gpu_name": GPU_NAME,
|
||||
"ocr_available": True,
|
||||
"ocr_languages": OCR_LANGUAGES,
|
||||
"parser_isolated": True,
|
||||
"parser_concurrency": 1,
|
||||
"parser_timeout_seconds": PARSER_TIMEOUT_SECONDS,
|
||||
"parser_cpu_seconds": PARSER_CPU_SECONDS,
|
||||
"parser_memory_mib": PARSER_MEMORY_BYTES // (1024 * 1024),
|
||||
"model_loaded": MODEL_LOADED,
|
||||
"model_disabled": MODEL_DISABLED,
|
||||
"summarize_available": MODEL_LOADED and not MODEL_DISABLED,
|
||||
@@ -1055,47 +1067,6 @@ async def purge_cache():
|
||||
return {"cleared": cleared}
|
||||
|
||||
|
||||
def _normalize_text(value: str) -> str:
|
||||
"""Normalize extraction noise without destroying section, bullet, or table boundaries."""
|
||||
value = value.replace("\x00", " ").replace("\r\n", "\n").replace("\r", "\n")
|
||||
lines = []
|
||||
blank = False
|
||||
for raw_line in value.split("\n"):
|
||||
line = re.sub(r"[\t\f\v]+", " ", raw_line).strip()
|
||||
if not line:
|
||||
if lines and not blank:
|
||||
lines.append("")
|
||||
blank = True
|
||||
continue
|
||||
lines.append(line)
|
||||
blank = False
|
||||
return "\n".join(lines).strip()
|
||||
|
||||
|
||||
def _ocr_image(image: Image.Image) -> str:
|
||||
if image.mode not in ("RGB", "L"):
|
||||
image = image.convert("RGB")
|
||||
text = pytesseract.image_to_string(image, lang=OCR_LANGUAGES)
|
||||
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:
|
||||
@@ -1137,76 +1108,144 @@ def _validate_file_signature(extension: str, data: bytes) -> None:
|
||||
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:
|
||||
_PARSER_ERRORS = {
|
||||
"unsupported_type": (400, "This file type is not supported for AI extraction."),
|
||||
"text_encoding": (400, "The uploaded text file is not valid UTF-8."),
|
||||
"docx_signature": (400, "The uploaded file does not match its DOCX extension."),
|
||||
"docx_entry_limit": (422, "The DOCX contains too many files."),
|
||||
"docx_entry_size": (422, "The DOCX contains an oversized file."),
|
||||
"docx_ratio": (422, "The DOCX compression ratio is not supported."),
|
||||
"docx_expanded_size": (422, "The DOCX expands beyond the supported size."),
|
||||
"pdf_page_limit": (422, "The PDF contains too many pages."),
|
||||
"pdf_ocr_limit": (422, "The PDF requires too much OCR processing."),
|
||||
"image_dimensions": (422, "The document image dimensions are not supported."),
|
||||
"image_pixel_limit": (422, "The document image contains too many pixels."),
|
||||
"image_frames": (422, "Multi-frame document images are not supported."),
|
||||
"extracted_text_limit": (422, "The document contains too much extracted text."),
|
||||
"no_text": (422, "AI extraction did not find readable text in the uploaded file."),
|
||||
"parser_failed": (422, "Document extraction failed."),
|
||||
}
|
||||
|
||||
|
||||
def _parser_environment(temp_root: str) -> dict[str, str]:
|
||||
allowed = ("PATH", "LANG", "LC_ALL", "TESSDATA_PREFIX", "SYSTEMROOT", "WINDIR")
|
||||
environment = {key: os.environ[key] for key in allowed if os.environ.get(key)}
|
||||
environment.update({
|
||||
"PYTHONUTF8": "1",
|
||||
"PYTHONDONTWRITEBYTECODE": "1",
|
||||
"OMP_NUM_THREADS": "1",
|
||||
"OPENBLAS_NUM_THREADS": "1",
|
||||
"TMPDIR": temp_root,
|
||||
"TEMP": temp_root,
|
||||
"TMP": temp_root,
|
||||
})
|
||||
return environment
|
||||
|
||||
|
||||
def _prepare_parser_temp_root() -> Path:
|
||||
global _parser_temp_ready
|
||||
with _parser_temp_lock:
|
||||
PARSER_TEMP_ROOT.mkdir(mode=0o700, parents=True, exist_ok=True)
|
||||
if _parser_temp_ready:
|
||||
return PARSER_TEMP_ROOT
|
||||
cutoff = time.time() - PARSER_STALE_WORK_SECONDS
|
||||
for candidate in PARSER_TEMP_ROOT.glob("work-*"):
|
||||
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 = _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 _check_extracted_size(_normalize_text("\n".join(ocr_pages))), True, page_count
|
||||
if candidate.is_symlink() or not candidate.is_dir():
|
||||
continue
|
||||
resolved = candidate.resolve()
|
||||
if resolved.parent != PARSER_TEMP_ROOT or candidate.stat().st_mtime > cutoff:
|
||||
continue
|
||||
shutil.rmtree(resolved)
|
||||
except OSError:
|
||||
continue
|
||||
_parser_temp_ready = True
|
||||
return PARSER_TEMP_ROOT
|
||||
|
||||
|
||||
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
|
||||
for block in blocks:
|
||||
if hasattr(block, "rows"):
|
||||
for row in block.rows:
|
||||
cells = [cell.text.strip() for cell in row.cells if cell.text and cell.text.strip()]
|
||||
if cells:
|
||||
parts.append(" | ".join(cells))
|
||||
elif getattr(block, "text", "").strip():
|
||||
text = block.text.strip()
|
||||
style = (getattr(getattr(block, "style", None), "name", "") or "").lower()
|
||||
if "role" in style and parts:
|
||||
parts.append("")
|
||||
parts.append(f"- {text}" if "bullet" in style else text)
|
||||
return _check_extracted_size(_normalize_text("\n".join(parts)))
|
||||
def _terminate_parser_process(process: subprocess.Popen) -> None:
|
||||
if process.poll() is not None:
|
||||
return
|
||||
if os.name == "posix":
|
||||
try:
|
||||
os.killpg(process.pid, signal.SIGKILL)
|
||||
except ProcessLookupError:
|
||||
pass
|
||||
else:
|
||||
try:
|
||||
subprocess.run(
|
||||
["taskkill", "/PID", str(process.pid), "/T", "/F"],
|
||||
stdin=subprocess.DEVNULL,
|
||||
stdout=subprocess.DEVNULL,
|
||||
stderr=subprocess.DEVNULL,
|
||||
timeout=5,
|
||||
check=False,
|
||||
)
|
||||
except (OSError, subprocess.SubprocessError):
|
||||
process.kill()
|
||||
|
||||
|
||||
def _extract_plain_text(data: bytes) -> str:
|
||||
def _raise_parser_error(code: str) -> None:
|
||||
fallback = _PARSER_ERRORS["parser_failed"]
|
||||
status_code, detail = _PARSER_ERRORS.get(code, fallback) if isinstance(code, str) else fallback
|
||||
raise HTTPException(status_code=status_code, detail=detail)
|
||||
|
||||
|
||||
def _run_parser_child(extension: str, data: bytes) -> dict[str, object]:
|
||||
if not _parser_capacity.acquire(blocking=False):
|
||||
raise HTTPException(status_code=503, detail="Document extraction is busy. Try again shortly.")
|
||||
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))
|
||||
with tempfile.TemporaryDirectory(prefix="work-", dir=_prepare_parser_temp_root()) as temp_root:
|
||||
input_path = Path(temp_root, "input.bin")
|
||||
output_path = Path(temp_root, "result.json")
|
||||
input_path.write_bytes(data)
|
||||
command = [
|
||||
sys.executable,
|
||||
str(PARSER_CHILD_PATH),
|
||||
"--input", str(input_path),
|
||||
"--output", str(output_path),
|
||||
"--extension", extension,
|
||||
"--cpu-seconds", str(PARSER_CPU_SECONDS),
|
||||
"--memory-bytes", str(PARSER_MEMORY_BYTES),
|
||||
]
|
||||
creation_flags = subprocess.CREATE_NEW_PROCESS_GROUP if os.name == "nt" else 0
|
||||
try:
|
||||
process = subprocess.Popen(
|
||||
command,
|
||||
cwd=temp_root,
|
||||
env=_parser_environment(temp_root),
|
||||
stdin=subprocess.DEVNULL,
|
||||
stdout=subprocess.DEVNULL,
|
||||
stderr=subprocess.DEVNULL,
|
||||
start_new_session=os.name == "posix",
|
||||
creationflags=creation_flags,
|
||||
)
|
||||
except OSError:
|
||||
_raise_parser_error("parser_failed")
|
||||
try:
|
||||
process.wait(timeout=PARSER_TIMEOUT_SECONDS)
|
||||
except subprocess.TimeoutExpired as exc:
|
||||
_terminate_parser_process(process)
|
||||
try:
|
||||
process.wait(timeout=5)
|
||||
except subprocess.TimeoutExpired:
|
||||
process.kill()
|
||||
raise HTTPException(status_code=504, detail="Document extraction timed out.") from exc
|
||||
|
||||
if not output_path.exists() or output_path.stat().st_size > PARSER_MAX_RESULT_BYTES:
|
||||
_raise_parser_error("parser_failed")
|
||||
try:
|
||||
result = json.loads(output_path.read_text(encoding="utf-8"))
|
||||
except (OSError, UnicodeError, json.JSONDecodeError):
|
||||
_raise_parser_error("parser_failed")
|
||||
if not isinstance(result, dict) or result.get("ok") is not True:
|
||||
_raise_parser_error(result.get("code") if isinstance(result, dict) else "parser_failed")
|
||||
text = result.get("text")
|
||||
if not isinstance(text, str) or not text or len(text) > MAX_EXTRACT_CHARS:
|
||||
_raise_parser_error("parser_failed")
|
||||
return result
|
||||
finally:
|
||||
_parser_capacity.release()
|
||||
|
||||
|
||||
@app.post("/extract-text")
|
||||
@@ -1228,36 +1267,14 @@ async def extract_text(file: UploadFile = File(...)):
|
||||
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"}:
|
||||
text = _extract_plain_text(data)
|
||||
ocr_used = False
|
||||
page_count = None
|
||||
elif extension == ".docx":
|
||||
text = _extract_docx_text(data)
|
||||
ocr_used = False
|
||||
page_count = None
|
||||
elif extension == ".pdf":
|
||||
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
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as 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.")
|
||||
result = await asyncio.to_thread(_run_parser_child, extension, data)
|
||||
text = result["text"]
|
||||
|
||||
return {
|
||||
"text": text,
|
||||
"ocr_used": ocr_used,
|
||||
"ocr_used": result.get("ocr_used", False),
|
||||
"content_type": file.content_type,
|
||||
"page_count": page_count,
|
||||
"page_count": result.get("page_count"),
|
||||
"characters": len(text),
|
||||
"file_name": filename,
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user