From a1ad802e2283c3bb5c367d5f7e982e432f4259e7 Mon Sep 17 00:00:00 2001 From: Alex Date: Tue, 4 Aug 2026 14:04:44 +0100 Subject: [PATCH 1/2] feat: limit docling attachment file sizes explicitly --- application/api/user/tasks.py | 10 +- application/core/settings.py | 3 + application/parser/file/docling_parser.py | 144 +++++++++++++ application/parser/file/tabular_parser.py | 2 +- application/worker.py | 138 +++++++++++- tests/parser/file/test_docling_parser.py | 247 ++++++++++++++++++++++ tests/parser/file/test_tabular_parser.py | 22 +- tests/worker/test_attachment_worker.py | 165 +++++++++++++++ 8 files changed, 714 insertions(+), 17 deletions(-) diff --git a/application/api/user/tasks.py b/application/api/user/tasks.py index fd577e24..d1fc60d2 100644 --- a/application/api/user/tasks.py +++ b/application/api/user/tasks.py @@ -6,6 +6,7 @@ from sqlalchemy.exc import DataError from application.api.user.idempotency import with_idempotency from application.celery_init import celery from application.worker import ( + AttachmentRejectedError, agent_webhook_worker, attachment_worker, ingest_worker, @@ -235,10 +236,11 @@ def _emit_attachment_poison_event(task_name, bound): ) -# ``dont_autoretry_for``: a DataError is deterministic (poison payload, -# e.g. NUL bytes or an over-long value) — retrying re-fails identically -# and multiplies log noise, so it goes straight to the failure path. -@celery.task(**DURABLE_TASK, dont_autoretry_for=(DataError,)) +# ``dont_autoretry_for``: a DataError (poison payload, e.g. NUL bytes or an +# over-long value) or an AttachmentRejectedError (zip bomb) is deterministic — +# retrying re-fails identically and multiplies log noise, so it goes straight +# to the failure path. +@celery.task(**DURABLE_TASK, dont_autoretry_for=(DataError, AttachmentRejectedError)) @with_idempotency( task_name="store_attachment", on_poison=_emit_attachment_poison_event, ) diff --git a/application/core/settings.py b/application/core/settings.py index 8db6dc24..62a6c2ce 100644 --- a/application/core/settings.py +++ b/application/core/settings.py @@ -95,6 +95,9 @@ class Settings(BaseSettings): # Pages docling's threaded pipeline buffers in flight; the library # default (100) drives worker RSS to ~3 GB on a mid-size PDF. DOCLING_PIPELINE_QUEUE_MAX_SIZE: int = 2 + DOCLING_TABULAR_MAX_BYTES: int = 2_000_000 + DOCLING_MARKUP_MAX_BYTES: int = 8_000_000 + ATTACHMENT_TEXT_MAX_BYTES: int = 5_000_000 VECTOR_STORE: str = "faiss" # "faiss" or "elasticsearch" or "qdrant" or "milvus" or "lancedb" or "pgvector" # Allow-list of retriever keys an agent may use. Values must match the # ``RetrieverCreator.retrievers`` registry keys (``classic`` / ``default``), diff --git a/application/parser/file/docling_parser.py b/application/parser/file/docling_parser.py index ff0d3656..d52d2b2c 100644 --- a/application/parser/file/docling_parser.py +++ b/application/parser/file/docling_parser.py @@ -8,6 +8,9 @@ images (PNG, JPEG, TIFF, BMP, WEBP), WebVTT, and specialized XML formats. """ import importlib.util import logging +import os +import tempfile +import zipfile from pathlib import Path from typing import Dict, List, Optional, Union @@ -39,6 +42,115 @@ def _apply_pipeline_caps(pipeline_options) -> None: setattr(pipeline_options, name, value) +def _tabular_content_size(file: Path) -> int: + """Effective content size of a tabular file, in bytes. + + Docling's memory scales with cell count, and for XLSX the cell count is + hidden by zip compression — a 2 MB xlsx can hold ~850k cells. So for the + zip-based ``.xlsx`` the measure is the inner-uncompressed size (read from + the central directory without decompressing); for plain-text CSV the + on-disk size already reflects the content. + + Args: + file: Path to the tabular file. + + Returns: + Content size in bytes, or -1 when it can't be determined. + """ + path = Path(file) + try: + if path.suffix.lower() == ".xlsx": + with zipfile.ZipFile(path) as zf: + return sum(info.file_size for info in zf.infolist()) + return path.stat().st_size + except (OSError, zipfile.BadZipFile): + try: + return path.stat().st_size + except OSError: + return -1 + + +def _exceeds_tabular_gate(file: Path) -> bool: + """Whether a tabular file is too large (by content) to hand to docling. + + Docling materializes a ``TableCell`` model per cell, so tabular memory + scales with cell count, not on-disk bytes (~11 KB of RSS per 4-cell CSV + row — an 88 MB CSV needs ~26 GB). Oversized CSV/XLSX files are routed to + the lightweight parsers in ``tabular_parser`` instead. + + Args: + file: Path to the tabular file about to be parsed. + + Returns: + True when the file's content exceeds ``DOCLING_TABULAR_MAX_BYTES`` + (and the gate is enabled), False otherwise or when size is unknown. + """ + from application.core.settings import settings + + max_bytes = settings.DOCLING_TABULAR_MAX_BYTES + if max_bytes <= 0: + return False + return _tabular_content_size(file) > max_bytes + + +def _capped_markup_copy(file: Path) -> Optional[str]: + """Head-truncate an oversized markup file to a temp copy for docling. + + HTML/VTT have no lightweight full-content parser to fall back to (unlike + CSV/XLSX), so element-dense markup is bounded by parsing only the first + ``DOCLING_MARKUP_MAX_BYTES`` — enough context for retrieval, and docling's + lenient HTML/VTT backends handle a truncated tail. Cuts on a line boundary + when one is reasonably close to the limit. + + Args: + file: Path to the markup file about to be parsed. + + Returns: + Path to a temp copy the caller must delete, or None when the file is + within the limit / the gate is disabled / the size can't be read. + """ + from application.core.settings import settings + + max_bytes = settings.DOCLING_MARKUP_MAX_BYTES + if max_bytes <= 0: + return None + try: + if Path(file).stat().st_size <= max_bytes: + return None + with open(file, "rb") as src: + head = src.read(max_bytes) + except OSError: + return None + cut = head.rfind(b"\n") + if cut > max_bytes // 2: + head = head[: cut + 1] + with tempfile.NamedTemporaryFile( + delete=False, suffix=Path(file).suffix + ) as tmp: + tmp.write(head) + return tmp.name + + +def _parse_markup_bounded( + parser: "DoclingParser", file: Path, errors: str +) -> Union[str, List[str]]: + """Run ``parser`` on ``file``, first truncating it if it is oversized markup.""" + capped = _capped_markup_copy(file) + if capped is None: + return DoclingParser.parse_file(parser, file, errors) + logger.warning( + f"Markup {Path(file).name} exceeds DOCLING_MARKUP_MAX_BYTES; " + "parsing a head-truncated copy to bound memory" + ) + try: + return DoclingParser.parse_file(parser, Path(capped), errors) + finally: + try: + os.unlink(capped) + except OSError: + pass + + class DoclingParser(BaseParser): """Parser using docling for advanced document processing. @@ -288,6 +400,18 @@ class DoclingXLSXParser(DoclingParser): def __init__(self): super().__init__(table_structure=True, export_format="markdown") + def parse_file(self, file: Path, errors: str = "ignore") -> Union[str, List[str]]: + """Parse an XLSX file, delegating oversized files to ``ExcelParser``.""" + if _exceeds_tabular_gate(file): + logger.warning( + f"XLSX {file.name} exceeds DOCLING_TABULAR_MAX_BYTES; " + "using lightweight Excel parser instead of docling" + ) + from application.parser.file.tabular_parser import ExcelParser + + return ExcelParser().parse_file(file, errors) + return super().parse_file(file, errors) + class DoclingHTMLParser(DoclingParser): """Docling-based HTML parser.""" @@ -295,6 +419,10 @@ class DoclingHTMLParser(DoclingParser): def __init__(self): super().__init__(export_format="markdown") + def parse_file(self, file: Path, errors: str = "ignore") -> Union[str, List[str]]: + """Parse HTML, truncating oversized element-dense markup first.""" + return _parse_markup_bounded(self, file, errors) + class DoclingImageParser(DoclingParser): """Docling-based image parser with OCR and RapidOCR support. @@ -325,6 +453,18 @@ class DoclingCSVParser(DoclingParser): def __init__(self): super().__init__(table_structure=True, export_format="markdown") + def parse_file(self, file: Path, errors: str = "ignore") -> Union[str, List[str]]: + """Parse a CSV file, delegating oversized files to the plain ``CSVParser``.""" + if _exceeds_tabular_gate(file): + logger.warning( + f"CSV {file.name} exceeds DOCLING_TABULAR_MAX_BYTES; " + "using plain CSV parser instead of docling" + ) + from application.parser.file.tabular_parser import CSVParser + + return CSVParser().parse_file(file, errors) + return super().parse_file(file, errors) + class DoclingMarkdownParser(DoclingParser): """Docling-based Markdown parser.""" @@ -346,6 +486,10 @@ class DoclingVTTParser(DoclingParser): def __init__(self): super().__init__(export_format="markdown") + def parse_file(self, file: Path, errors: str = "ignore") -> Union[str, List[str]]: + """Parse WebVTT, truncating oversized cue-dense tracks first.""" + return _parse_markup_bounded(self, file, errors) + class DoclingXMLParser(DoclingParser): """Docling-based XML parser (USPTO, JATS).""" diff --git a/application/parser/file/tabular_parser.py b/application/parser/file/tabular_parser.py index 40971b3c..cc9c821e 100644 --- a/application/parser/file/tabular_parser.py +++ b/application/parser/file/tabular_parser.py @@ -200,7 +200,7 @@ class ExcelParser(BaseParser): raise ValueError("pandas module is required to read Excel files.") df = pd.read_excel(file, **self._pandas_config) - headers = df.columns.tolist() + headers = [str(h) for h in df.columns.tolist()] header_row = f"{self._header_prefix}{self._col_joiner.join(headers)}" if not self._concat_rows: diff --git a/application/worker.py b/application/worker.py index 5bd8abb7..9ea99b35 100755 --- a/application/worker.py +++ b/application/worker.py @@ -1469,6 +1469,114 @@ def sync_worker(self, frequency): } +# Line-oriented text formats that stay parseable after a head-truncation. +# Structured/binary formats (pdf, docx, xlsx, json, html, ...) are excluded — +# cutting them mid-stream would break their parsers entirely. +_TRUNCATABLE_ATTACHMENT_SUFFIXES = { + ".csv", + ".tsv", + ".txt", + ".log", + ".md", + ".mdx", + ".rst", + ".jsonl", +} + + +# Zip-container document formats. A malicious archive can declare a tiny +# on-disk size but a huge inner-uncompressed payload (or a huge entry count), +# so these are size-checked before a parser touches them. +_ZIP_CONTAINER_ATTACHMENT_EXTS = frozenset({".docx", ".xlsx", ".pptx", ".epub"}) + + +class AttachmentRejectedError(Exception): + """A deterministic reason an attachment must not be parsed (e.g. zip bomb). + + Raised before parsing and marked non-retryable on the Celery task so a + poison upload fails once instead of retrying identically. + """ + + +def _reject_attachment_zip_bomb(local_path: str) -> None: + """Reject a zip-container attachment that decompresses to too much. + + Mirrors the ingest-path guard (``document_reader._reject_zip_bomb``) on the + attachment path, which previously had no such check. Reads only the zip + central directory (declared sizes), never decompressing. + + Args: + local_path: Filesystem path of the attachment about to be parsed. + + Raises: + AttachmentRejectedError: If the archive exceeds the entry-count or + inner-uncompressed-size caps. + """ + suffix = os.path.splitext(local_path)[1].lower() + if suffix not in _ZIP_CONTAINER_ATTACHMENT_EXTS: + return + max_entries = int(getattr(settings, "DOCUMENT_MAX_ARCHIVE_ENTRIES", 10000)) + cap = int(getattr(settings, "DOCUMENT_MAX_DECOMPRESSED_BYTES", 300 * 1024 * 1024)) + try: + with zipfile.ZipFile(local_path) as zf: + infos = zf.infolist() + except zipfile.BadZipFile: + # Not a readable zip; the format parser will surface a clean error. + return + if len(infos) > max_entries: + raise AttachmentRejectedError( + f"archive has too many entries ({len(infos)} > {max_entries})" + ) + total = 0 + for info in infos: + total += info.file_size + if total > cap: + raise AttachmentRejectedError( + f"archive decompresses to too much data (exceeds {cap} bytes)" + ) + + +def _bounded_attachment_copy(local_path: str) -> tuple[str, bool]: + """Bound how much of a text attachment reaches the parser. + + Attachment content is capped at ~250k chars after parsing, so bytes past + ``ATTACHMENT_TEXT_MAX_BYTES`` only cost parse time and memory. Oversized + line-oriented text files are head-truncated on a line boundary into a + temp copy; the stored original is never modified (local storage hands the + canonical file path to ``process_file`` callbacks). + + Args: + local_path: Filesystem path handed to the parse callback. + + Returns: + Tuple of (path to parse, whether it is a temp copy the caller must + delete). + """ + max_bytes = settings.ATTACHMENT_TEXT_MAX_BYTES + if max_bytes <= 0: + return local_path, False + suffix = os.path.splitext(local_path)[1].lower() + if suffix not in _TRUNCATABLE_ATTACHMENT_SUFFIXES: + return local_path, False + try: + if os.path.getsize(local_path) <= max_bytes: + return local_path, False + with open(local_path, "rb") as src: + head = src.read(max_bytes) + except OSError: + return local_path, False + cut = head.rfind(b"\n") + if cut > 0: + head = head[: cut + 1] + with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp: + tmp.write(head) + logging.warning( + f"Attachment {os.path.basename(local_path)} exceeds " + f"ATTACHMENT_TEXT_MAX_BYTES ({max_bytes}); parsing first {len(head)} bytes" + ) + return tmp.name, True + + def attachment_worker(self, file_info, user): """ Process and store a single attachment without vectorization. @@ -1508,17 +1616,25 @@ def attachment_worker(self, file_info, user): file_extractor = get_default_file_extractor( ocr_enabled=settings.DOCLING_OCR_ATTACHMENTS_ENABLED ) - attachment_document = storage.process_file( - relative_path, - lambda local_path, **kwargs: SimpleDirectoryReader( - input_files=[local_path], - exclude_hidden=True, - errors="ignore", - file_extractor=file_extractor, - file_metadata=metadata_from_filename, - ) - .load_data()[0], - ) + def _parse_local_file(local_path: str, **kwargs) -> Document: + _reject_attachment_zip_bomb(local_path) + parse_path, is_temp_copy = _bounded_attachment_copy(local_path) + try: + return SimpleDirectoryReader( + input_files=[parse_path], + exclude_hidden=True, + errors="ignore", + file_extractor=file_extractor, + file_metadata=metadata_from_filename, + ).load_data()[0] + finally: + if is_temp_copy: + try: + os.unlink(parse_path) + except OSError: + pass + + attachment_document = storage.process_file(relative_path, _parse_local_file) content = attachment_document.text parser_metadata = { key: value diff --git a/tests/parser/file/test_docling_parser.py b/tests/parser/file/test_docling_parser.py index 97f96368..13b3a8b2 100644 --- a/tests/parser/file/test_docling_parser.py +++ b/tests/parser/file/test_docling_parser.py @@ -4,6 +4,7 @@ Covers: DoclingParser (init, _init_parser, _get_ocr_options, _export_content, parse_file), subclass initialization, error handling. """ +import os from pathlib import Path from unittest.mock import MagicMock, patch @@ -503,3 +504,249 @@ class TestApplyPipelineCaps: assert not hasattr(opts, "queue_max_size") assert not hasattr(opts, "layout_batch_size") + + +# ===================================================================== +# Tabular size gate (CSV / XLSX) +# ===================================================================== + + +@pytest.mark.unit +class TestDoclingTabularSizeGate: + """Oversized tabular files must bypass docling. + + Docling materializes a ``TableCell`` object per cell (measured ~11 KB of + RSS per 4-cell CSV row), so a multi-MB CSV balloons the worker by tens of + GB. Above ``DOCLING_TABULAR_MAX_BYTES`` the docling tabular parsers must + delegate to the lightweight parsers in ``tabular_parser``. + """ + + def _write_csv(self, tmp_path: Path, rows: int = 50) -> Path: + path = tmp_path / "data.csv" + path.write_text("\n".join(f"{i},{i * 2}" for i in range(rows)) + "\n") + return path + + def test_oversized_csv_delegates_to_plain_csv_parser(self, tmp_path, monkeypatch): + from application.core.settings import settings + from application.parser.file.docling_parser import DoclingCSVParser, DoclingParser + + monkeypatch.setattr(settings, "DOCLING_TABULAR_MAX_BYTES", 64) + docling_parse = MagicMock(name="docling_parse") + monkeypatch.setattr(DoclingParser, "parse_file", docling_parse) + + path = self._write_csv(tmp_path) + assert path.stat().st_size > 64 + + out = DoclingCSVParser().parse_file(path) + + docling_parse.assert_not_called() + # Plain ``CSVParser`` output: rows joined with ", ", newline-separated. + assert out.startswith("0, 0\n1, 2\n") + + def test_small_csv_still_uses_docling(self, tmp_path, monkeypatch): + from application.core.settings import settings + from application.parser.file.docling_parser import DoclingCSVParser, DoclingParser + + monkeypatch.setattr(settings, "DOCLING_TABULAR_MAX_BYTES", 10_000_000) + docling_parse = MagicMock(name="docling_parse", return_value="DOCLING") + monkeypatch.setattr(DoclingParser, "parse_file", docling_parse) + + out = DoclingCSVParser().parse_file(self._write_csv(tmp_path)) + + assert out == "DOCLING" + docling_parse.assert_called_once() + + def test_gate_disabled_when_max_bytes_is_zero(self, tmp_path, monkeypatch): + from application.core.settings import settings + from application.parser.file.docling_parser import DoclingCSVParser, DoclingParser + + monkeypatch.setattr(settings, "DOCLING_TABULAR_MAX_BYTES", 0) + docling_parse = MagicMock(name="docling_parse", return_value="DOCLING") + monkeypatch.setattr(DoclingParser, "parse_file", docling_parse) + + out = DoclingCSVParser().parse_file(self._write_csv(tmp_path, rows=5000)) + + assert out == "DOCLING" + docling_parse.assert_called_once() + + def test_oversized_xlsx_delegates_to_excel_parser(self, tmp_path, monkeypatch): + from application.core.settings import settings + from application.parser.file.docling_parser import DoclingParser, DoclingXLSXParser + from application.parser.file.tabular_parser import ExcelParser + + monkeypatch.setattr(settings, "DOCLING_TABULAR_MAX_BYTES", 64) + docling_parse = MagicMock(name="docling_parse") + monkeypatch.setattr(DoclingParser, "parse_file", docling_parse) + excel_parse = MagicMock(name="excel_parse", return_value="PLAIN-XLSX") + monkeypatch.setattr(ExcelParser, "parse_file", excel_parse) + + path = tmp_path / "big.xlsx" + path.write_bytes(b"x" * 200) + + out = DoclingXLSXParser().parse_file(path) + + assert out == "PLAIN-XLSX" + docling_parse.assert_not_called() + excel_parse.assert_called_once() + + def test_small_xlsx_still_uses_docling(self, tmp_path, monkeypatch): + from application.core.settings import settings + from application.parser.file.docling_parser import DoclingParser, DoclingXLSXParser + + monkeypatch.setattr(settings, "DOCLING_TABULAR_MAX_BYTES", 10_000_000) + docling_parse = MagicMock(name="docling_parse", return_value="DOCLING") + monkeypatch.setattr(DoclingParser, "parse_file", docling_parse) + + path = tmp_path / "small.xlsx" + path.write_bytes(b"x" * 200) + + assert DoclingXLSXParser().parse_file(path) == "DOCLING" + docling_parse.assert_called_once() + + +# ===================================================================== +# Tabular content-size gate — XLSX compression must not defeat it +# ===================================================================== + + +def _make_xlsx(path: Path, rows: int) -> None: + from openpyxl import Workbook + + wb = Workbook(write_only=True) + ws = wb.create_sheet() + for i in range(rows): + ws.append([i, i * 2, i % 7, i % 13, i * 3]) + wb.save(str(path)) + + +@pytest.mark.unit +class TestTabularContentSize: + """XLSX is zip-compressed, so the gate must measure inner-uncompressed + size, not on-disk bytes — otherwise a small-on-disk / many-cell xlsx + (2.44 GB in docling) slips under a byte gate.""" + + def test_xlsx_content_size_is_inner_not_ondisk(self, tmp_path): + from application.parser.file.docling_parser import _tabular_content_size + + path = tmp_path / "data.xlsx" + _make_xlsx(path, rows=5000) + on_disk = path.stat().st_size + inner = _tabular_content_size(path) + # Repetitive numeric data compresses hard: inner XML >> zip on disk. + assert inner > on_disk + + def test_csv_content_size_is_ondisk(self, tmp_path): + from application.parser.file.docling_parser import _tabular_content_size + + path = tmp_path / "data.csv" + path.write_text("a,b\n1,2\n3,4\n") + assert _tabular_content_size(path) == path.stat().st_size + + def test_compressed_xlsx_over_inner_gate_delegates(self, tmp_path, monkeypatch): + """The regression: on-disk < threshold < inner-uncompressed must gate.""" + from application.core.settings import settings + from application.parser.file.docling_parser import ( + DoclingParser, + DoclingXLSXParser, + _tabular_content_size, + ) + from application.parser.file.tabular_parser import ExcelParser + + path = tmp_path / "wide.xlsx" + _make_xlsx(path, rows=5000) + on_disk = path.stat().st_size + inner = _tabular_content_size(path) + assert inner > on_disk, "premise: compression hides cell count" + + # A byte gate between the two would have sent this to docling; the + # inner-size gate must catch it. + monkeypatch.setattr( + settings, "DOCLING_TABULAR_MAX_BYTES", (on_disk + inner) // 2 + ) + docling_parse = MagicMock(name="docling_parse") + monkeypatch.setattr(DoclingParser, "parse_file", docling_parse) + excel_parse = MagicMock(name="excel_parse", return_value="PLAIN") + monkeypatch.setattr(ExcelParser, "parse_file", excel_parse) + + out = DoclingXLSXParser().parse_file(path) + + assert out == "PLAIN" + docling_parse.assert_not_called() + excel_parse.assert_called_once() + + +# ===================================================================== +# Markup gate (HTML / VTT) — truncate oversized element-dense markup +# ===================================================================== + + +@pytest.mark.unit +class TestDoclingMarkupGate: + def _write(self, tmp_path: Path, name: str, nbytes: int) -> Path: + path = tmp_path / name + # newline-terminated lines so the line-boundary trim has something to cut + path.write_text(("x" * 63 + "\n") * (nbytes // 64 + 1)) + return path + + @pytest.mark.parametrize("name", ["big.html", "big.vtt"]) + def test_oversized_markup_parses_truncated_copy(self, tmp_path, monkeypatch, name): + from application.core.settings import settings + from application.parser.file import docling_parser as dp + + monkeypatch.setattr(settings, "DOCLING_MARKUP_MAX_BYTES", 512) + path = self._write(tmp_path, name, 4096) + assert path.stat().st_size > 512 + + seen = {} + + def fake_parse(self_parser, file, errors="ignore"): + p = str(file) + seen["path"] = p + seen["size"] = os.path.getsize(p) + return "PARSED" + + monkeypatch.setattr(dp.DoclingParser, "parse_file", fake_parse) + + cls = dp.DoclingHTMLParser if name.endswith(".html") else dp.DoclingVTTParser + out = cls().parse_file(path) + + assert out == "PARSED" + assert seen["path"] != str(path), "must parse a temp copy, not the original" + assert seen["size"] <= 512, "temp copy must be truncated to the cap" + assert not os.path.exists(seen["path"]), "temp copy must be cleaned up" + assert path.stat().st_size > 512, "original must be untouched" + + def test_small_markup_parses_original(self, tmp_path, monkeypatch): + from application.core.settings import settings + from application.parser.file import docling_parser as dp + + monkeypatch.setattr(settings, "DOCLING_MARKUP_MAX_BYTES", 10_000_000) + path = self._write(tmp_path, "small.html", 1024) + seen = {} + + def fake_parse(self_parser, file, errors="ignore"): + seen["path"] = str(file) + return "PARSED" + + monkeypatch.setattr(dp.DoclingParser, "parse_file", fake_parse) + out = dp.DoclingHTMLParser().parse_file(path) + + assert out == "PARSED" + assert seen["path"] == str(path) + + def test_markup_gate_disabled_when_zero(self, tmp_path, monkeypatch): + from application.core.settings import settings + from application.parser.file import docling_parser as dp + + monkeypatch.setattr(settings, "DOCLING_MARKUP_MAX_BYTES", 0) + path = self._write(tmp_path, "big.vtt", 8192) + seen = {} + + def fake_parse(self_parser, file, errors="ignore"): + seen["path"] = str(file) + return "PARSED" + + monkeypatch.setattr(dp.DoclingParser, "parse_file", fake_parse) + dp.DoclingVTTParser().parse_file(path) + + assert seen["path"] == str(path), "disabled gate must parse the original" diff --git a/tests/parser/file/test_tabular_parser.py b/tests/parser/file/test_tabular_parser.py index 3dbab672..30adc505 100644 --- a/tests/parser/file/test_tabular_parser.py +++ b/tests/parser/file/test_tabular_parser.py @@ -212,4 +212,24 @@ def test_excel_parser_import_error(excel_parser): import sys with patch.dict(sys.modules, {"pandas": None}): with pytest.raises(ValueError, match="pandas module is required to read Excel files"): - excel_parser.parse_file(Path("test.xlsx")) \ No newline at end of file + excel_parser.parse_file(Path("test.xlsx")) + +def test_excel_numeric_headers_do_not_crash(tmp_path): + """Regression: a headerless/numeric xlsx gives integer column labels, and + ``", ".join(headers)`` used to raise TypeError. The XLSX size-gate delegates + to this parser, so it must handle numeric headers.""" + from openpyxl import Workbook + + from application.parser.file.tabular_parser import ExcelParser + + wb = Workbook() + ws = wb.active + for i in range(5): + ws.append([i, i * 2, i * 3]) + path = tmp_path / "numeric.xlsx" + wb.save(str(path)) + + out = ExcelParser().parse_file(path) + + assert isinstance(out, str) + assert len(out) > 0 diff --git a/tests/worker/test_attachment_worker.py b/tests/worker/test_attachment_worker.py index ecca13fe..cbac1489 100644 --- a/tests/worker/test_attachment_worker.py +++ b/tests/worker/test_attachment_worker.py @@ -9,6 +9,7 @@ assert one concrete row is visible after the task returns. from __future__ import annotations +from pathlib import Path from unittest.mock import MagicMock import pytest @@ -65,3 +66,167 @@ class TestAttachmentWorker: assert row["upload_path"] == "uploads/user1/notes.txt" assert row["content"] == "hello world" assert row["user_id"] == "user1" + + +@pytest.mark.unit +class TestBoundedAttachmentCopy: + """``_bounded_attachment_copy`` head-truncates oversized line-oriented + text attachments to a temp copy before parsing. + + The parsed content is capped at ~250k chars downstream anyway, so bytes + beyond the cap only cost parse time and memory. Local storage hands the + canonical stored file to the processor, so truncation must never happen + in place. + """ + + def _write(self, tmp_path, name: str, data: bytes): + path = tmp_path / name + path.write_bytes(data) + return path + + def test_oversized_csv_is_copied_and_truncated(self, tmp_path, monkeypatch): + from application import worker + from application.core.settings import settings + + monkeypatch.setattr(settings, "ATTACHMENT_TEXT_MAX_BYTES", 1024) + original = self._write( + tmp_path, "big.csv", b"".join(b"%d,%d\n" % (i, i) for i in range(1000)) + ) + original_size = original.stat().st_size + + parse_path, is_temp = worker._bounded_attachment_copy(str(original)) + + assert is_temp is True + assert parse_path != str(original) + assert parse_path.endswith(".csv") + copied = Path(parse_path).read_bytes() + assert 0 < len(copied) <= 1024 + assert copied.endswith(b"\n"), "must cut on a line boundary" + # The stored original must be untouched. + assert original.stat().st_size == original_size + Path(parse_path).unlink() + + def test_small_file_returned_as_is(self, tmp_path, monkeypatch): + from application import worker + from application.core.settings import settings + + monkeypatch.setattr(settings, "ATTACHMENT_TEXT_MAX_BYTES", 1024) + original = self._write(tmp_path, "small.csv", b"a,b\n1,2\n") + + parse_path, is_temp = worker._bounded_attachment_copy(str(original)) + + assert parse_path == str(original) + assert is_temp is False + + def test_non_text_suffix_is_never_truncated(self, tmp_path, monkeypatch): + from application import worker + from application.core.settings import settings + + monkeypatch.setattr(settings, "ATTACHMENT_TEXT_MAX_BYTES", 64) + original = self._write(tmp_path, "doc.pdf", b"%PDF-1.7 " + b"x" * 500) + + parse_path, is_temp = worker._bounded_attachment_copy(str(original)) + + assert parse_path == str(original) + assert is_temp is False + + def test_cap_zero_disables_truncation(self, tmp_path, monkeypatch): + from application import worker + from application.core.settings import settings + + monkeypatch.setattr(settings, "ATTACHMENT_TEXT_MAX_BYTES", 0) + original = self._write(tmp_path, "big.csv", b"1,2\n" * 1000) + + parse_path, is_temp = worker._bounded_attachment_copy(str(original)) + + assert parse_path == str(original) + assert is_temp is False + + def test_single_line_without_newline_falls_back_to_hard_cut( + self, tmp_path, monkeypatch + ): + from application import worker + from application.core.settings import settings + + monkeypatch.setattr(settings, "ATTACHMENT_TEXT_MAX_BYTES", 256) + original = self._write(tmp_path, "oneline.txt", b"x" * 5000) + + parse_path, is_temp = worker._bounded_attachment_copy(str(original)) + + assert is_temp is True + data = Path(parse_path).read_bytes() + assert len(data) == 256 + Path(parse_path).unlink() + + +@pytest.mark.unit +class TestAttachmentZipBombGuard: + """``_reject_attachment_zip_bomb`` brings the ingest-path zip-bomb guard to + the attachment path (which previously had none): a zip-container attachment + that declares too many entries / too much inner data is rejected before any + parser touches it, with a non-retryable error.""" + + def _make_xlsx(self, path: Path, rows: int = 20): + from openpyxl import Workbook + + wb = Workbook(write_only=True) + ws = wb.create_sheet() + for i in range(rows): + ws.append([i, i * 2, i * 3]) + wb.save(str(path)) + + def test_rejects_when_inner_size_exceeds_cap(self, tmp_path, monkeypatch): + from application import worker + from application.core.settings import settings + + path = tmp_path / "book.xlsx" + self._make_xlsx(path) + monkeypatch.setattr(settings, "DOCUMENT_MAX_DECOMPRESSED_BYTES", 100) + + with pytest.raises(worker.AttachmentRejectedError): + worker._reject_attachment_zip_bomb(str(path)) + + def test_rejects_when_too_many_entries(self, tmp_path, monkeypatch): + from application import worker + from application.core.settings import settings + + path = tmp_path / "book.xlsx" + self._make_xlsx(path) + monkeypatch.setattr(settings, "DOCUMENT_MAX_ARCHIVE_ENTRIES", 1) + + with pytest.raises(worker.AttachmentRejectedError): + worker._reject_attachment_zip_bomb(str(path)) + + def test_allows_reasonable_archive(self, tmp_path, monkeypatch): + from application import worker + from application.core.settings import settings + + path = tmp_path / "book.xlsx" + self._make_xlsx(path) + monkeypatch.setattr(settings, "DOCUMENT_MAX_DECOMPRESSED_BYTES", 300 * 1024 * 1024) + monkeypatch.setattr(settings, "DOCUMENT_MAX_ARCHIVE_ENTRIES", 10000) + + # Must not raise. + worker._reject_attachment_zip_bomb(str(path)) + + def test_non_container_suffix_is_ignored(self, tmp_path, monkeypatch): + from application import worker + from application.core.settings import settings + + path = tmp_path / "notes.txt" + path.write_bytes(b"x" * 5000) + monkeypatch.setattr(settings, "DOCUMENT_MAX_DECOMPRESSED_BYTES", 1) + + # Text files are not zip containers — never inspected, never rejected. + worker._reject_attachment_zip_bomb(str(path)) + + def test_corrupt_zip_is_left_to_the_parser(self, tmp_path, monkeypatch): + from application import worker + from application.core.settings import settings + + path = tmp_path / "broken.xlsx" + path.write_bytes(b"not a real zip") + monkeypatch.setattr(settings, "DOCUMENT_MAX_DECOMPRESSED_BYTES", 1) + + # BadZipFile → return quietly; the format parser surfaces a clean error. + worker._reject_attachment_zip_bomb(str(path)) From 15b6b03b47d69553854171b9a528f93693ac132c Mon Sep 17 00:00:00 2001 From: Alex Date: Tue, 4 Aug 2026 15:07:54 +0100 Subject: [PATCH 2/2] fix: minor size cleanup --- application/parser/document_reader.py | 37 +++++++++++++++-- application/parser/file/docling_parser.py | 5 +-- application/utils.py | 22 ++++++++++ application/worker.py | 49 +++++++---------------- tests/parser/test_document_reader.py | 38 ++++++++++++++++++ tests/test_utils.py | 48 ++++++++++++++++++++++ tests/worker/test_attachment_worker.py | 20 +++++++++ 7 files changed, 177 insertions(+), 42 deletions(-) diff --git a/application/parser/document_reader.py b/application/parser/document_reader.py index 3edbfb6a..752f7c33 100644 --- a/application/parser/document_reader.py +++ b/application/parser/document_reader.py @@ -15,7 +15,7 @@ import os import tempfile import zipfile from pathlib import Path -from typing import Any, Dict, Iterator, List, Optional +from typing import Any, Dict, Iterator, List, Optional, Union from application.core.settings import settings from application.parser.file.bulk import get_default_file_extractor @@ -105,14 +105,26 @@ def _max_input_bytes() -> int: _ZIP_CONTAINER_EXTENSIONS = frozenset({".docx", ".xlsx", ".pptx", ".epub"}) -def _reject_zip_bomb(data: bytes, suffix: str) -> Optional[str]: - """Return an error string if a zip-based document declares an implausible expansion, else None.""" +def _zip_bomb_reason(source: Union[bytes, str, Path], suffix: str) -> Optional[str]: + """Return an error string if a zip container declares an implausible expansion, else None. + + Reads only the central directory (declared sizes), never decompressing. + + Args: + source: Document bytes, or a filesystem path to the document. + suffix: Lowercased file extension; non-zip formats are not gated here. + + Returns: + The rejection reason, or None when the archive is acceptable, is not a + zip container, or is not a readable zip. + """ if suffix not in _ZIP_CONTAINER_EXTENSIONS: return None max_entries = int(getattr(settings, "DOCUMENT_MAX_ARCHIVE_ENTRIES", 10000)) cap = int(getattr(settings, "DOCUMENT_MAX_DECOMPRESSED_BYTES", 300 * 1024 * 1024)) + opened = io.BytesIO(source) if isinstance(source, bytes) else source try: - with zipfile.ZipFile(io.BytesIO(data)) as zf: + with zipfile.ZipFile(opened) as zf: infos = zf.infolist() if len(infos) > max_entries: return f"document archive has too many entries ({len(infos)} > {max_entries})." @@ -127,6 +139,23 @@ def _reject_zip_bomb(data: bytes, suffix: str) -> Optional[str]: return None +def _reject_zip_bomb(data: bytes, suffix: str) -> Optional[str]: + """Return an error string if a zip-based document declares an implausible expansion, else None.""" + return _zip_bomb_reason(data, suffix) + + +def reject_zip_bomb_path(path: Union[str, Path]) -> Optional[str]: + """Path-taking sibling of ``_reject_zip_bomb`` for callers that have a file, not bytes. + + Args: + path: Filesystem path of the document about to be parsed. + + Returns: + The rejection reason, or None when the file is acceptable / not a zip container. + """ + return _zip_bomb_reason(path, os.path.splitext(str(path))[1].lower()) + + def _resolve_ocr_enabled(ocr: str) -> bool: """Resolve the OCR flag from the ``ocr`` arg and the deployment setting.""" if ocr == "on": diff --git a/application/parser/file/docling_parser.py b/application/parser/file/docling_parser.py index d52d2b2c..50f5ef98 100644 --- a/application/parser/file/docling_parser.py +++ b/application/parser/file/docling_parser.py @@ -15,6 +15,7 @@ from pathlib import Path from typing import Dict, List, Optional, Union from application.parser.file.base_parser import BaseParser +from application.utils import truncate_to_line_boundary logger = logging.getLogger(__name__) @@ -121,9 +122,7 @@ def _capped_markup_copy(file: Path) -> Optional[str]: head = src.read(max_bytes) except OSError: return None - cut = head.rfind(b"\n") - if cut > max_bytes // 2: - head = head[: cut + 1] + head = truncate_to_line_boundary(head) with tempfile.NamedTemporaryFile( delete=False, suffix=Path(file).suffix ) as tmp: diff --git a/application/utils.py b/application/utils.py index eea3569c..60032893 100644 --- a/application/utils.py +++ b/application/utils.py @@ -79,6 +79,28 @@ def strip_null_bytes(value): return value +def truncate_to_line_boundary(data: bytes) -> bytes: + """Trim a head-truncated byte window back to its last line boundary. + + The trim is skipped when it would discard more than half the window — a + file whose only newline sits near the start (or at byte 0) would otherwise + collapse to a few bytes, which is far worse than a partial final line. + Callers pass a window already read at their size cap, so the cut never + grows the result. + + Args: + data: The head window read from an oversized file. + + Returns: + ``data`` up to and including its last newline, or ``data`` unchanged + when no newline is far enough in to be worth cutting at. + """ + cut = data.rfind(b"\n") + if cut > len(data) // 2: + return data[: cut + 1] + return data + + def num_tokens_from_string(string: str) -> int: encoding = get_encoding() if isinstance(string, str): diff --git a/application/worker.py b/application/worker.py index 9ea99b35..1b91c231 100755 --- a/application/worker.py +++ b/application/worker.py @@ -48,7 +48,12 @@ from application.storage.db.repositories.wiki_pages import ( from application.storage.db.session import db_readonly, db_session from application.storage.db.source_config import SourceConfig from application.storage.storage_creator import StorageCreator -from application.utils import count_tokens_docs, num_tokens_from_string, safe_filename +from application.utils import ( + count_tokens_docs, + num_tokens_from_string, + safe_filename, + truncate_to_line_boundary, +) # Constants @@ -1484,12 +1489,6 @@ _TRUNCATABLE_ATTACHMENT_SUFFIXES = { } -# Zip-container document formats. A malicious archive can declare a tiny -# on-disk size but a huge inner-uncompressed payload (or a huge entry count), -# so these are size-checked before a parser touches them. -_ZIP_CONTAINER_ATTACHMENT_EXTS = frozenset({".docx", ".xlsx", ".pptx", ".epub"}) - - class AttachmentRejectedError(Exception): """A deterministic reason an attachment must not be parsed (e.g. zip bomb). @@ -1501,9 +1500,8 @@ class AttachmentRejectedError(Exception): def _reject_attachment_zip_bomb(local_path: str) -> None: """Reject a zip-container attachment that decompresses to too much. - Mirrors the ingest-path guard (``document_reader._reject_zip_bomb``) on the - attachment path, which previously had no such check. Reads only the zip - central directory (declared sizes), never decompressing. + Applies the shared guard (``document_reader.reject_zip_bomb_path``) on the + attachment path, which previously had no such check. Args: local_path: Filesystem path of the attachment about to be parsed. @@ -1512,28 +1510,11 @@ def _reject_attachment_zip_bomb(local_path: str) -> None: AttachmentRejectedError: If the archive exceeds the entry-count or inner-uncompressed-size caps. """ - suffix = os.path.splitext(local_path)[1].lower() - if suffix not in _ZIP_CONTAINER_ATTACHMENT_EXTS: - return - max_entries = int(getattr(settings, "DOCUMENT_MAX_ARCHIVE_ENTRIES", 10000)) - cap = int(getattr(settings, "DOCUMENT_MAX_DECOMPRESSED_BYTES", 300 * 1024 * 1024)) - try: - with zipfile.ZipFile(local_path) as zf: - infos = zf.infolist() - except zipfile.BadZipFile: - # Not a readable zip; the format parser will surface a clean error. - return - if len(infos) > max_entries: - raise AttachmentRejectedError( - f"archive has too many entries ({len(infos)} > {max_entries})" - ) - total = 0 - for info in infos: - total += info.file_size - if total > cap: - raise AttachmentRejectedError( - f"archive decompresses to too much data (exceeds {cap} bytes)" - ) + from application.parser.document_reader import reject_zip_bomb_path + + reason = reject_zip_bomb_path(local_path) + if reason is not None: + raise AttachmentRejectedError(reason) def _bounded_attachment_copy(local_path: str) -> tuple[str, bool]: @@ -1565,9 +1546,7 @@ def _bounded_attachment_copy(local_path: str) -> tuple[str, bool]: head = src.read(max_bytes) except OSError: return local_path, False - cut = head.rfind(b"\n") - if cut > 0: - head = head[: cut + 1] + head = truncate_to_line_boundary(head) with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp: tmp.write(head) logging.warning( diff --git a/tests/parser/test_document_reader.py b/tests/parser/test_document_reader.py index e43e7547..7d594eb0 100644 --- a/tests/parser/test_document_reader.py +++ b/tests/parser/test_document_reader.py @@ -66,6 +66,44 @@ def test_reject_zip_bomb_ignores_non_zip_formats(): assert dr._reject_zip_bomb(b"not a zip", ".docx") is None +@pytest.mark.unit +def test_reject_zip_bomb_path_matches_the_bytes_variant(tmp_path, monkeypatch): + # The path-taking sibling (used by the attachment worker, which has a file + # rather than bytes) must reach the same verdict as the bytes variant. + monkeypatch.setattr(dr.settings, "DOCUMENT_MAX_DECOMPRESSED_BYTES", 1000, raising=False) + data = _make_zip({"word/document.xml": b"A" * 50_000}) + path = tmp_path / "bomb.docx" + path.write_bytes(data) + + reason = dr.reject_zip_bomb_path(str(path)) + + assert reason is not None and "too much data" in reason + assert reason == dr._reject_zip_bomb(data, ".docx") + + +@pytest.mark.unit +def test_reject_zip_bomb_path_allows_reasonable_archive(tmp_path, monkeypatch): + monkeypatch.setattr(dr.settings, "DOCUMENT_MAX_DECOMPRESSED_BYTES", 300 * 1024 * 1024, raising=False) + monkeypatch.setattr(dr.settings, "DOCUMENT_MAX_ARCHIVE_ENTRIES", 10000, raising=False) + path = tmp_path / "fine.docx" + path.write_bytes(_make_zip({"word/document.xml": b"hello"})) + + assert dr.reject_zip_bomb_path(path) is None + + +@pytest.mark.unit +def test_reject_zip_bomb_path_ignores_non_container_and_corrupt_files(tmp_path, monkeypatch): + monkeypatch.setattr(dr.settings, "DOCUMENT_MAX_DECOMPRESSED_BYTES", 1, raising=False) + text = tmp_path / "notes.txt" + text.write_bytes(b"x" * 5000) + broken = tmp_path / "broken.xlsx" + broken.write_bytes(b"not a real zip") + + # Non-container suffix is never inspected; a corrupt zip is left to the parser. + assert dr.reject_zip_bomb_path(str(text)) is None + assert dr.reject_zip_bomb_path(str(broken)) is None + + @pytest.mark.unit def test_size_cap_rejects_oversize(monkeypatch): monkeypatch.setattr(dr.settings, "DOCUMENT_PARSE_MAX_BYTES", 8, raising=False) diff --git a/tests/test_utils.py b/tests/test_utils.py index 40567289..b8e15048 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -20,6 +20,7 @@ from application.utils import ( num_tokens_from_object_or_list, num_tokens_from_string, safe_filename, + truncate_to_line_boundary, validate_function_name, validate_required_fields, ) @@ -666,6 +667,53 @@ class TestGetHashEdgeCases: assert len(h) == 32 +class TestTruncateToLineBoundary: + """Trimming a head window back to its last line boundary. + + Shared by the attachment size gate and the docling markup gate. The trim + is skipped whenever it would cost more than half the window: a partial + final line is always better than losing the content. + """ + + @pytest.mark.unit + def test_trims_to_last_newline(self): + data = b"aaaa\nbbbb\ncccc" + assert truncate_to_line_boundary(data) == b"aaaa\nbbbb\n" + + @pytest.mark.unit + def test_already_line_terminated_is_unchanged(self): + data = b"aaaa\nbbbb\n" + assert truncate_to_line_boundary(data) == data + + @pytest.mark.unit + def test_no_newline_falls_back_to_the_hard_cut(self): + data = b"x" * 500 + assert truncate_to_line_boundary(data) == data + + @pytest.mark.unit + def test_leading_newline_only_keeps_the_content(self): + # rfind returns 0 here; trimming would collapse the window to a single + # byte, so the partial line is kept instead. + data = b"\n" + b"x" * 499 + assert truncate_to_line_boundary(data) == data + + @pytest.mark.unit + def test_early_newline_keeps_the_content(self): + # A short first line followed by one huge line: cutting at byte 10 + # would discard ~98% of the window. + data = b"header\n" + b"x" * 493 + assert truncate_to_line_boundary(data) == data + + @pytest.mark.unit + def test_newline_past_the_halfway_mark_is_used(self): + data = b"x" * 300 + b"\n" + b"y" * 199 + assert truncate_to_line_boundary(data) == b"x" * 300 + b"\n" + + @pytest.mark.unit + def test_empty_input(self): + assert truncate_to_line_boundary(b"") == b"" + + class TestValidateFunctionNameEdgeCases: @pytest.mark.unit diff --git a/tests/worker/test_attachment_worker.py b/tests/worker/test_attachment_worker.py index cbac1489..c6a63aca 100644 --- a/tests/worker/test_attachment_worker.py +++ b/tests/worker/test_attachment_worker.py @@ -158,6 +158,26 @@ class TestBoundedAttachmentCopy: assert len(data) == 256 Path(parse_path).unlink() + def test_leading_newline_does_not_collapse_the_copy(self, tmp_path, monkeypatch): + """A window whose only newline sits at byte 0 must keep its content. + + ``rfind`` returns 0 for this shape, so cutting at that boundary would + write a one-byte copy and throw the attachment away — the partial + final line is the better trade. + """ + from application import worker + from application.core.settings import settings + + monkeypatch.setattr(settings, "ATTACHMENT_TEXT_MAX_BYTES", 256) + original = self._write(tmp_path, "leading.log", b"\n" + b"x" * 5000) + + parse_path, is_temp = worker._bounded_attachment_copy(str(original)) + + assert is_temp is True + data = Path(parse_path).read_bytes() + assert len(data) == 256 + Path(parse_path).unlink() + @pytest.mark.unit class TestAttachmentZipBombGuard: