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: