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
@@ -2,8 +2,11 @@ FROM python:3.12.10-slim-bookworm
|
||||
ENV PIP_NO_CACHE_DIR=1 \
|
||||
PIP_DISABLE_PIP_VERSION_CHECK=1 \
|
||||
PYTHONUNBUFFERED=1 \
|
||||
PYTHONDONTWRITEBYTECODE=1 \
|
||||
TRANSFORMERS_NO_TF=1 \
|
||||
HF_HUB_DISABLE_TELEMETRY=1
|
||||
HF_HUB_DISABLE_TELEMETRY=1 \
|
||||
HF_HOME=/home/app/.cache/huggingface \
|
||||
PARSER_TEMP_ROOT=/tmp
|
||||
WORKDIR /app
|
||||
RUN apt-get update \
|
||||
&& apt-get install -y --no-install-recommends tesseract-ocr tesseract-ocr-eng \
|
||||
@@ -12,5 +15,10 @@ COPY requirements-linux.lock ./
|
||||
RUN python -m pip install --upgrade pip setuptools wheel \
|
||||
&& python -m pip install --require-hashes --extra-index-url https://download.pytorch.org/whl/cpu -r requirements-linux.lock
|
||||
COPY . .
|
||||
RUN groupadd --system app \
|
||||
&& useradd --system --gid app --home-dir /home/app --create-home app \
|
||||
&& mkdir -p /home/app/.cache/huggingface \
|
||||
&& chown -R app:app /app /home/app
|
||||
USER app
|
||||
EXPOSE 8001
|
||||
CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "8001"]
|
||||
|
||||
@@ -39,6 +39,14 @@ If the host is missing `python3-venv` or `pip`, use the bootstrap script instead
|
||||
## Docker
|
||||
The Dockerfile installs Tesseract OCR so scanned PDFs and supported images can be processed inside the container.
|
||||
|
||||
## Document parser isolation
|
||||
|
||||
`/extract-text` performs only bounded upload and signature checks in the web process. Actual TXT, Markdown, PDF, DOCX and image decoding runs in a single-capacity child process with a minimal environment. Linux children receive CPU, address-space, output-size and file-descriptor limits; every platform enforces a parent deadline and terminates the parser process tree on timeout.
|
||||
|
||||
Production Compose runs the service as a non-root user with a read-only root filesystem, no Linux capabilities, `no-new-privileges`, PID/CPU/memory limits and a bounded `/tmp` tmpfs. The Hugging Face cache is the only named writable volume. Parser work directories are removed after each request and stale `work-*` directories are reconciled on the next extraction.
|
||||
|
||||
Parser limits can be made more conservative with `PARSER_TIMEOUT_SECONDS`, `PARSER_CPU_SECONDS`, `PARSER_MEMORY_MIB` and `PARSER_STALE_WORK_SECONDS`. Increasing them should follow measured benign-CV evidence; disabling import is the safe fallback when the budget is insufficient.
|
||||
|
||||
## Tests
|
||||
|
||||
Run the summarizer unit tests with:
|
||||
|
||||
+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,
|
||||
}
|
||||
|
||||
@@ -0,0 +1,253 @@
|
||||
"""Resource-bounded CV document decoder invoked only as a child process."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import re
|
||||
import sys
|
||||
import zipfile
|
||||
|
||||
|
||||
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"}
|
||||
|
||||
|
||||
class ParserError(Exception):
|
||||
def __init__(self, code: str):
|
||||
super().__init__(code)
|
||||
self.code = code
|
||||
|
||||
|
||||
def _apply_resource_limits(cpu_seconds: int, memory_bytes: int) -> None:
|
||||
if os.name != "posix":
|
||||
return
|
||||
import resource
|
||||
|
||||
resource.setrlimit(resource.RLIMIT_CPU, (cpu_seconds, cpu_seconds + 1))
|
||||
resource.setrlimit(resource.RLIMIT_AS, (memory_bytes, memory_bytes))
|
||||
resource.setrlimit(resource.RLIMIT_FSIZE, (1024 * 1024, 1024 * 1024))
|
||||
resource.setrlimit(resource.RLIMIT_NOFILE, (64, 64))
|
||||
resource.setrlimit(resource.RLIMIT_CORE, (0, 0))
|
||||
|
||||
|
||||
def normalize_text(value: str) -> str:
|
||||
value = value.replace("\x00", " ").replace("\r\n", "\n").replace("\r", "\n")
|
||||
lines: list[str] = []
|
||||
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 _check_extracted_size(text: str, limit: int = MAX_EXTRACT_CHARS) -> str:
|
||||
if len(text) > limit:
|
||||
raise ParserError("extracted_text_limit")
|
||||
return text
|
||||
|
||||
|
||||
def _check_image_size(image, *, pixel_limit: int = MAX_IMAGE_PIXELS) -> None:
|
||||
width, height = image.size
|
||||
if width <= 0 or height <= 0 or width > MAX_IMAGE_DIMENSION or height > MAX_IMAGE_DIMENSION:
|
||||
raise ParserError("image_dimensions")
|
||||
if width * height > pixel_limit:
|
||||
raise ParserError("image_pixel_limit")
|
||||
if getattr(image, "n_frames", 1) != 1:
|
||||
raise ParserError("image_frames")
|
||||
|
||||
|
||||
def validate_docx_container(
|
||||
data: bytes,
|
||||
*,
|
||||
entry_limit: int = MAX_DOCX_ENTRIES,
|
||||
total_limit: int = MAX_DOCX_UNCOMPRESSED_BYTES,
|
||||
entry_size_limit: int = MAX_DOCX_ENTRY_BYTES,
|
||||
ratio_limit: int = MAX_DOCX_COMPRESSION_RATIO,
|
||||
) -> 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 ParserError("docx_signature")
|
||||
if len(entries) > entry_limit:
|
||||
raise ParserError("docx_entry_limit")
|
||||
total_size = 0
|
||||
for entry in entries:
|
||||
total_size += entry.file_size
|
||||
if entry.file_size > entry_size_limit:
|
||||
raise ParserError("docx_entry_size")
|
||||
if entry.file_size > 0 and entry.compress_size == 0:
|
||||
raise ParserError("docx_ratio")
|
||||
if entry.compress_size > 0 and entry.file_size / entry.compress_size > ratio_limit:
|
||||
raise ParserError("docx_ratio")
|
||||
if total_size > total_limit:
|
||||
raise ParserError("docx_expanded_size")
|
||||
except ParserError:
|
||||
raise
|
||||
except (zipfile.BadZipFile, OSError) as exc:
|
||||
raise ParserError("docx_signature") from exc
|
||||
|
||||
|
||||
def extract_plain_text(data: bytes) -> str:
|
||||
try:
|
||||
decoded = data.decode("utf-8-sig")
|
||||
except UnicodeDecodeError as exc:
|
||||
raise ParserError("text_encoding") from exc
|
||||
return _check_extracted_size(normalize_text(decoded))
|
||||
|
||||
|
||||
def extract_docx_text(data: bytes) -> str:
|
||||
validate_docx_container(data)
|
||||
from docx import Document
|
||||
|
||||
document = Document(io.BytesIO(data))
|
||||
parts: list[str] = []
|
||||
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 _ocr_image(image) -> str:
|
||||
import pytesseract
|
||||
|
||||
if image.mode not in ("RGB", "L"):
|
||||
image = image.convert("RGB")
|
||||
return normalize_text(pytesseract.image_to_string(image, lang=OCR_LANGUAGES))
|
||||
|
||||
|
||||
def extract_image_text(data: bytes, *, pixel_limit: int = MAX_IMAGE_PIXELS) -> str:
|
||||
from PIL import Image
|
||||
|
||||
with Image.open(io.BytesIO(data)) as image:
|
||||
_check_image_size(image, pixel_limit=pixel_limit)
|
||||
return _check_extracted_size(_ocr_image(image))
|
||||
|
||||
|
||||
def extract_pdf_text(data: bytes, *, page_limit: int = MAX_PDF_PAGES) -> tuple[str, bool, int]:
|
||||
from pypdf import PdfReader
|
||||
|
||||
page_count = 0
|
||||
extracted_pages: list[str] = []
|
||||
try:
|
||||
reader = PdfReader(io.BytesIO(data))
|
||||
page_count = len(reader.pages)
|
||||
if page_count > page_limit:
|
||||
raise ParserError("pdf_page_limit")
|
||||
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 ParserError:
|
||||
raise
|
||||
except Exception:
|
||||
extracted_pages = []
|
||||
|
||||
text = _check_extracted_size(normalize_text("\n".join(extracted_pages)))
|
||||
if len(text) >= 80:
|
||||
return text, False, page_count
|
||||
|
||||
import fitz
|
||||
from PIL import Image
|
||||
|
||||
document = fitz.open(stream=data, filetype="pdf")
|
||||
try:
|
||||
page_count = document.page_count
|
||||
if page_count > page_limit:
|
||||
raise ParserError("pdf_page_limit")
|
||||
ocr_pages: list[str] = []
|
||||
total_pixels = 0
|
||||
for page in document:
|
||||
pixmap = page.get_pixmap(matrix=fitz.Matrix(2, 2), alpha=False)
|
||||
page_pixels = pixmap.width * pixmap.height
|
||||
total_pixels += page_pixels
|
||||
if page_pixels > MAX_PDF_PAGE_PIXELS or total_pixels > MAX_PDF_OCR_PIXELS:
|
||||
raise ParserError("pdf_ocr_limit")
|
||||
with Image.open(io.BytesIO(pixmap.tobytes("png"))) as image:
|
||||
_check_image_size(image, pixel_limit=MAX_PDF_PAGE_PIXELS)
|
||||
ocr_pages.append(_ocr_image(image))
|
||||
_check_extracted_size("\n".join(ocr_pages))
|
||||
return _check_extracted_size(normalize_text("\n".join(ocr_pages))), True, page_count
|
||||
finally:
|
||||
document.close()
|
||||
|
||||
|
||||
def extract(extension: str, data: bytes) -> dict[str, object]:
|
||||
if extension in {".txt", ".md"}:
|
||||
text = extract_plain_text(data)
|
||||
return {"text": text, "ocr_used": False, "page_count": None}
|
||||
if extension == ".docx":
|
||||
text = extract_docx_text(data)
|
||||
return {"text": text, "ocr_used": False, "page_count": None}
|
||||
if extension == ".pdf":
|
||||
text, ocr_used, page_count = extract_pdf_text(data)
|
||||
return {"text": text, "ocr_used": ocr_used, "page_count": page_count}
|
||||
if extension in IMAGE_EXTENSIONS:
|
||||
text = extract_image_text(data)
|
||||
return {"text": text, "ocr_used": True, "page_count": 1}
|
||||
raise ParserError("unsupported_type")
|
||||
|
||||
|
||||
def _write_result(path: Path, result: dict[str, object]) -> None:
|
||||
path.write_text(json.dumps(result, ensure_ascii=False, separators=(",", ":")), encoding="utf-8")
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser(add_help=False)
|
||||
parser.add_argument("--input", required=True)
|
||||
parser.add_argument("--output", required=True)
|
||||
parser.add_argument("--extension", required=True)
|
||||
parser.add_argument("--cpu-seconds", type=int, default=20)
|
||||
parser.add_argument("--memory-bytes", type=int, default=512 * 1024 * 1024)
|
||||
args = parser.parse_args()
|
||||
output_path = Path(args.output)
|
||||
try:
|
||||
_apply_resource_limits(args.cpu_seconds, args.memory_bytes)
|
||||
data = Path(args.input).read_bytes()
|
||||
result = extract(args.extension.lower(), data)
|
||||
if not result.get("text"):
|
||||
raise ParserError("no_text")
|
||||
_write_result(output_path, {"ok": True, **result})
|
||||
return 0
|
||||
except ParserError as exc:
|
||||
_write_result(output_path, {"ok": False, "code": exc.code})
|
||||
return 2
|
||||
except BaseException:
|
||||
_write_result(output_path, {"ok": False, "code": "parser_failed"})
|
||||
return 3
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -1,3 +1,3 @@
|
||||
-r requirements.txt
|
||||
pytest==8.3.5
|
||||
httpx==0.28.1
|
||||
httpx2==2.12.0
|
||||
|
||||
@@ -1,10 +1,14 @@
|
||||
import importlib
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import textwrap
|
||||
import zipfile
|
||||
from pathlib import Path
|
||||
|
||||
from docx import Document
|
||||
from fastapi.testclient import TestClient
|
||||
from PIL import Image
|
||||
from pypdf import PdfWriter
|
||||
@@ -67,6 +71,10 @@ def test_health_reports_runtime_without_ollama_and_without_forcing_model_load(mo
|
||||
assert payload["model_loaded"] is False
|
||||
assert payload["model_disabled"] is True
|
||||
assert payload["summarize_available"] is False
|
||||
assert payload["parser_isolated"] is True
|
||||
assert payload["parser_concurrency"] == 1
|
||||
assert payload["parser_timeout_seconds"] == 25
|
||||
assert payload["parser_memory_mib"] == 512
|
||||
assert "disabled" in payload["model_load_error"].lower()
|
||||
assert payload["ollama_configured"] is False
|
||||
assert payload["ollama_model"] is None
|
||||
@@ -517,7 +525,7 @@ def test_health_stays_open_so_probes_and_healthchecks_work(monkeypatch):
|
||||
def test_extract_text_rejects_oversized_upload_before_parsing(monkeypatch):
|
||||
module = load_app_module(monkeypatch)
|
||||
monkeypatch.setattr(module, "MAX_EXTRACT_FILE_BYTES", 32)
|
||||
monkeypatch.setattr(module, "_extract_plain_text", lambda data: (_ for _ in ()).throw(AssertionError("parser must not run")))
|
||||
monkeypatch.setattr(module, "_run_parser_child", lambda extension, data: (_ for _ in ()).throw(AssertionError("parser must not run")))
|
||||
client = TestClient(module.app)
|
||||
|
||||
response = client.post("/extract-text", files={"file": ("large.txt", b"x" * 64, "text/plain")})
|
||||
@@ -547,33 +555,36 @@ def test_extract_text_rejects_binary_plain_text(monkeypatch):
|
||||
|
||||
|
||||
def test_extract_text_rejects_pdf_over_page_limit(monkeypatch):
|
||||
module = load_app_module(monkeypatch)
|
||||
monkeypatch.setattr(module, "MAX_PDF_PAGES", 1)
|
||||
load_app_module(monkeypatch)
|
||||
import parser_child
|
||||
|
||||
writer = PdfWriter()
|
||||
writer.add_blank_page(width=100, height=100)
|
||||
writer.add_blank_page(width=100, height=100)
|
||||
payload = io.BytesIO()
|
||||
writer.write(payload)
|
||||
client = TestClient(module.app)
|
||||
|
||||
response = client.post("/extract-text", files={"file": ("resume.pdf", payload.getvalue(), "application/pdf")})
|
||||
|
||||
assert response.status_code == 422
|
||||
assert response.json()["detail"] == "The PDF contains too many pages."
|
||||
try:
|
||||
parser_child.extract_pdf_text(payload.getvalue(), page_limit=1)
|
||||
except parser_child.ParserError as exc:
|
||||
assert exc.code == "pdf_page_limit"
|
||||
else:
|
||||
raise AssertionError("Expected the PDF page limit to reject the document")
|
||||
|
||||
|
||||
def test_extract_text_rejects_image_pixel_limit_before_ocr(monkeypatch):
|
||||
module = load_app_module(monkeypatch)
|
||||
monkeypatch.setattr(module, "MAX_IMAGE_PIXELS", 50)
|
||||
monkeypatch.setattr(module, "_ocr_image", lambda image: (_ for _ in ()).throw(AssertionError("OCR must not run")))
|
||||
load_app_module(monkeypatch)
|
||||
import parser_child
|
||||
|
||||
monkeypatch.setattr(parser_child, "_ocr_image", lambda image: (_ for _ in ()).throw(AssertionError("OCR must not run")))
|
||||
payload = io.BytesIO()
|
||||
Image.new("RGB", (10, 10), "white").save(payload, format="PNG")
|
||||
client = TestClient(module.app)
|
||||
|
||||
response = client.post("/extract-text", files={"file": ("resume.png", payload.getvalue(), "image/png")})
|
||||
|
||||
assert response.status_code == 422
|
||||
assert response.json()["detail"] == "The document image contains too many pixels."
|
||||
try:
|
||||
parser_child.extract_image_text(payload.getvalue(), pixel_limit=50)
|
||||
except parser_child.ParserError as exc:
|
||||
assert exc.code == "image_pixel_limit"
|
||||
else:
|
||||
raise AssertionError("Expected the image pixel limit to reject the document")
|
||||
|
||||
|
||||
def test_docx_container_rejects_excessive_entries(monkeypatch):
|
||||
@@ -596,7 +607,7 @@ def test_docx_container_rejects_excessive_entries(monkeypatch):
|
||||
|
||||
def test_extract_text_does_not_leak_parser_exception(monkeypatch):
|
||||
module = load_app_module(monkeypatch)
|
||||
monkeypatch.setattr(module, "_extract_pdf_text", lambda data: (_ for _ in ()).throw(RuntimeError("private path C:/secret")))
|
||||
monkeypatch.setattr(module, "_run_parser_child", lambda extension, data: module._raise_parser_error("parser_failed"))
|
||||
client = TestClient(module.app)
|
||||
|
||||
response = client.post("/extract-text", files={"file": ("resume.pdf", b"%PDF-invalid", "application/pdf")})
|
||||
@@ -607,16 +618,19 @@ def test_extract_text_does_not_leak_parser_exception(monkeypatch):
|
||||
|
||||
|
||||
def test_extraction_normalization_preserves_sections_and_bullets(monkeypatch):
|
||||
module = load_app_module(monkeypatch)
|
||||
load_app_module(monkeypatch)
|
||||
import parser_child
|
||||
|
||||
normalized = module._normalize_text("PROFILE\n\nEngineer\n• Built APIs\n• Shipped services")
|
||||
normalized = parser_child.normalize_text("PROFILE\n\nEngineer\n• Built APIs\n• Shipped services")
|
||||
|
||||
assert normalized.splitlines() == ["PROFILE", "", "Engineer", "• Built APIs", "• Shipped services"]
|
||||
|
||||
|
||||
def test_docx_extraction_preserves_paragraph_and_table_boundaries(monkeypatch):
|
||||
module = load_app_module(monkeypatch)
|
||||
document = module.Document()
|
||||
load_app_module(monkeypatch)
|
||||
import parser_child
|
||||
|
||||
document = Document()
|
||||
document.add_heading("Technical Skills", level=1)
|
||||
table = document.add_table(rows=2, cols=2)
|
||||
table.cell(0, 0).text = "Backend"
|
||||
@@ -626,13 +640,129 @@ def test_docx_extraction_preserves_paragraph_and_table_boundaries(monkeypatch):
|
||||
payload = io.BytesIO()
|
||||
document.save(payload)
|
||||
|
||||
extracted = module._extract_docx_text(payload.getvalue())
|
||||
extracted = parser_child.extract_docx_text(payload.getvalue())
|
||||
|
||||
assert "Technical Skills" in extracted
|
||||
assert "Backend | C#, .NET" in extracted
|
||||
assert "DevOps | Docker, Linux" in extracted
|
||||
|
||||
|
||||
def test_extract_text_runs_in_isolated_child(monkeypatch):
|
||||
module = load_app_module(monkeypatch)
|
||||
client = TestClient(module.app)
|
||||
|
||||
response = client.post(
|
||||
"/extract-text",
|
||||
files={"file": ("resume.md", b"# Ada Lovelace\n\n## Skills\nC#", "text/markdown")},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json()["text"] == "# Ada Lovelace\n\n## Skills\nC#"
|
||||
|
||||
|
||||
def test_parser_child_receives_no_provider_secrets(monkeypatch, tmp_path):
|
||||
module = load_app_module(monkeypatch)
|
||||
monkeypatch.setenv("GEMINI_API_KEY", "must-not-leak")
|
||||
monkeypatch.setenv("GROQ_API_KEY", "must-not-leak")
|
||||
child = tmp_path / "environment_probe.py"
|
||||
child.write_text(textwrap.dedent("""
|
||||
import argparse, json, os
|
||||
from pathlib import Path
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--input')
|
||||
parser.add_argument('--output')
|
||||
parser.add_argument('--extension')
|
||||
parser.add_argument('--cpu-seconds')
|
||||
parser.add_argument('--memory-bytes')
|
||||
args = parser.parse_args()
|
||||
leaked = any(os.environ.get(key) for key in ('GEMINI_API_KEY', 'GROQ_API_KEY', 'AI_SERVICE_TOKEN'))
|
||||
Path(args.output).write_text(json.dumps({'ok': not leaked, 'text': 'safe', 'ocr_used': False, 'page_count': None} if not leaked else {'ok': False, 'code': 'parser_failed'}), encoding='utf-8')
|
||||
"""), encoding="utf-8")
|
||||
monkeypatch.setattr(module, "PARSER_CHILD_PATH", child)
|
||||
|
||||
assert module._run_parser_child(".txt", b"safe")["text"] == "safe"
|
||||
|
||||
|
||||
def test_parser_child_timeout_is_stable_and_releases_capacity(monkeypatch, tmp_path):
|
||||
module = load_app_module(monkeypatch)
|
||||
child = tmp_path / "sleeping_parser.py"
|
||||
child.write_text("import time\ntime.sleep(30)\n", encoding="utf-8")
|
||||
monkeypatch.setattr(module, "PARSER_CHILD_PATH", child)
|
||||
monkeypatch.setattr(module, "PARSER_TIMEOUT_SECONDS", 0.1)
|
||||
|
||||
try:
|
||||
module._run_parser_child(".txt", b"safe")
|
||||
except module.HTTPException as exc:
|
||||
assert exc.status_code == 504
|
||||
assert exc.detail == "Document extraction timed out."
|
||||
else:
|
||||
raise AssertionError("Expected the parser deadline to terminate the child")
|
||||
|
||||
assert module._parser_capacity.acquire(blocking=False)
|
||||
module._parser_capacity.release()
|
||||
|
||||
|
||||
def test_parser_timeout_terminates_descendants(monkeypatch, tmp_path):
|
||||
module = load_app_module(monkeypatch)
|
||||
pid_path = tmp_path / "descendant.pid"
|
||||
child = tmp_path / "parser_with_descendant.py"
|
||||
child.write_text(textwrap.dedent(f"""
|
||||
import subprocess, sys, time
|
||||
from pathlib import Path
|
||||
descendant = subprocess.Popen([sys.executable, '-c', 'import time; time.sleep(30)'])
|
||||
Path({str(pid_path)!r}).write_text(str(descendant.pid), encoding='ascii')
|
||||
time.sleep(30)
|
||||
"""), encoding="utf-8")
|
||||
monkeypatch.setattr(module, "PARSER_CHILD_PATH", child)
|
||||
monkeypatch.setattr(module, "PARSER_TIMEOUT_SECONDS", 0.2)
|
||||
|
||||
try:
|
||||
module._run_parser_child(".txt", b"safe")
|
||||
except module.HTTPException as exc:
|
||||
assert exc.status_code == 504
|
||||
else:
|
||||
raise AssertionError("Expected the parser deadline to terminate the process tree")
|
||||
|
||||
descendant_pid = int(pid_path.read_text(encoding="ascii"))
|
||||
if os.name == "nt":
|
||||
listing = subprocess.run(
|
||||
["tasklist", "/FI", f"PID eq {descendant_pid}", "/FO", "CSV", "/NH"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=False,
|
||||
).stdout
|
||||
assert f'"{descendant_pid}"' not in listing
|
||||
else:
|
||||
try:
|
||||
os.kill(descendant_pid, 0)
|
||||
except ProcessLookupError:
|
||||
pass
|
||||
else:
|
||||
os.kill(descendant_pid, 9)
|
||||
raise AssertionError("Parser descendant survived the parent deadline")
|
||||
|
||||
|
||||
def test_parser_temp_cleanup_removes_only_stale_work_directories(monkeypatch, tmp_path):
|
||||
module = load_app_module(monkeypatch)
|
||||
parser_root = tmp_path / "parser-root"
|
||||
stale = parser_root / "work-stale"
|
||||
recent = parser_root / "work-recent"
|
||||
unrelated = parser_root / "keep-me"
|
||||
stale.mkdir(parents=True)
|
||||
recent.mkdir()
|
||||
unrelated.mkdir()
|
||||
old = 1_000_000
|
||||
os.utime(stale, (old, old))
|
||||
monkeypatch.setattr(module, "PARSER_TEMP_ROOT", parser_root.resolve())
|
||||
monkeypatch.setattr(module, "PARSER_STALE_WORK_SECONDS", 300)
|
||||
monkeypatch.setattr(module, "_parser_temp_ready", False)
|
||||
|
||||
assert module._prepare_parser_temp_root() == parser_root.resolve()
|
||||
assert not stale.exists()
|
||||
assert recent.exists()
|
||||
assert unrelated.exists()
|
||||
|
||||
|
||||
def test_cache_purge_requires_service_token_and_clears_content(monkeypatch):
|
||||
module = load_app_module(monkeypatch, service_token="s3cret")
|
||||
module.cache["synthetic-key"] = "synthetic-summary"
|
||||
|
||||
Reference in New Issue
Block a user