From fc9d17eb7b6b8f8e8fa0fe2a7057e141e952ae5f Mon Sep 17 00:00:00 2001 From: Xubin Ren <52506698+Re-bin@users.noreply.github.com> Date: Thu, 23 Jul 2026 16:23:29 +0800 Subject: [PATCH] fix(documents): bound nested DOCX table parsing --- nanobot/utils/document.py | 58 ++++++++++++++++++------ tests/test_document_parsing.py | 80 ++++++++++++++++++++++++++++++++++ 2 files changed, 125 insertions(+), 13 deletions(-) diff --git a/nanobot/utils/document.py b/nanobot/utils/document.py index 5570a1ee..a2a26654 100644 --- a/nanobot/utils/document.py +++ b/nanobot/utils/document.py @@ -1,6 +1,7 @@ """Document text extraction utilities for nanobot.""" import mimetypes +from collections.abc import Iterator from dataclasses import dataclass from pathlib import Path from zipfile import BadZipFile, ZipFile @@ -43,6 +44,8 @@ _MAX_EXTRACT_FILE_SIZE = 50 * 1024 * 1024 # 50 MB _MAX_OFFICE_ARCHIVE_MEMBERS = 10_000 _MAX_OFFICE_UNCOMPRESSED_SIZE = 256 * 1024 * 1024 # 256 MB _MAX_OFFICE_MEMBER_SIZE = 128 * 1024 * 1024 # 128 MB +_MAX_DOCX_TABLE_CELLS = 100_000 +_MAX_DOCX_TABLE_DEPTH = 8 _MAX_PDF_CONTENT_STREAM_SIZE = 32 * 1024 * 1024 # 32 MB per page _MAX_PDF_ATTACHMENT_PAGES = 100 @@ -87,6 +90,10 @@ class PdfPageRangeError(Exception): """Raised when a requested PDF page range is invalid.""" +class DocxSafetyError(Exception): + """Raised when a DOCX table exceeds a parser safety boundary.""" + + @dataclass(frozen=True, slots=True) class PdfExtraction: text: str @@ -210,7 +217,7 @@ def _extract_docx(path: Path) -> str: """Extract text from DOCX using python-docx.""" try: from docx import Document as DocxDocument - from docx.table import Table + from docx.table import Table, _Cell from docx.text.paragraph import Paragraph except ImportError: return "[error: python-docx not installed]" @@ -219,6 +226,39 @@ def _extract_docx(path: Path) -> str: return error doc = DocxDocument(path) collector = _TextCollector(_MAX_TEXT_LENGTH) + table_cell_count = 0 + + def cell_text(cell: _Cell, depth: int) -> str: + parts: list[str] = [] + for block in cell.iter_inner_content(): + if isinstance(block, Paragraph): + text = " ".join(block.text.split()) + if text: + parts.append(text) + elif isinstance(block, Table): + parts.extend(row.replace("\t", " | ") for row in table_rows(block, depth + 1)) + return " ".join(parts) + + def table_rows(table: Table, depth: int) -> Iterator[str]: + nonlocal table_cell_count + if depth > _MAX_DOCX_TABLE_DEPTH: + raise DocxSafetyError( + f"table nesting exceeds {_MAX_DOCX_TABLE_DEPTH} levels" + ) + for row in table.rows: + cells: list[str] = [] + # row.cells expands w:gridSpan before callers can apply a bound. + # Physical w:tc elements keep malformed documents proportional to XML size. + for tc in row._tr.tc_lst: + table_cell_count += 1 + if table_cell_count > _MAX_DOCX_TABLE_CELLS: + raise DocxSafetyError( + f"document contains more than {_MAX_DOCX_TABLE_CELLS} table cells" + ) + cells.append(cell_text(_Cell(tc, table), depth)) + if any(cells): + yield "\t".join(cells) + for block in doc.iter_inner_content(): if isinstance(block, Paragraph): text = block.text.strip() @@ -228,22 +268,14 @@ def _extract_docx(path: Path) -> str: if not isinstance(block, Table): continue first_row = True - for row in block.rows: - cells: list[str] = [] - seen_cells: set[int] = set() - for cell in row.cells: - cell_id = id(cell._tc) - if cell_id in seen_cells: - continue - seen_cells.add(cell_id) - cells.append(" ".join(cell.text.split())) - if not any(cells): - continue + for row_text in table_rows(block, 1): separator = "\n\n" if first_row else "\n" first_row = False - if not collector.add("\t".join(cells), separator=separator): + if not collector.add(row_text, separator=separator): return collector.render() return collector.render() + except DocxSafetyError as e: + return f"[error: unsafe DOCX: {e!s}]" except Exception as e: logger.exception("Failed to extract DOCX {}", path) return f"[error: failed to extract DOCX: {e!s}]" diff --git a/tests/test_document_parsing.py b/tests/test_document_parsing.py index 5657e784..ae552cca 100644 --- a/tests/test_document_parsing.py +++ b/tests/test_document_parsing.py @@ -197,6 +197,86 @@ class TestExtractText: assert result.index("Applicant details") < result.index("Name\tAda Lovelace") assert result.index("Analytical Engine") < result.index("End of form") + def test_extract_text_docx_preserves_nested_table_text(self, tmp_path: Path): + """Nested layout tables must not silently drop form fields.""" + from docx import Document + + docx_file = tmp_path / "nested-form.docx" + doc = Document() + outer_cell = doc.add_table(rows=1, cols=1).cell(0, 0) + outer_cell.add_paragraph("Contact") + nested = outer_cell.add_table(rows=1, cols=2) + nested.cell(0, 0).text = "Email" + nested.cell(0, 1).text = "ada@example.com" + doc.save(docx_file) + + result = extract_text(docx_file) + + assert result is not None + assert "Contact" in result + assert "Email" in result + assert "ada@example.com" in result + assert result.index("Contact") < result.index("Email") < result.index("ada@example.com") + + def test_extract_text_docx_does_not_expand_grid_spans( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ): + """Physical cells avoid python-docx's eager gridSpan expansion.""" + from docx import Document + from docx.table import _Row + + docx_file = tmp_path / "merged.docx" + doc = Document() + table = doc.add_table(rows=1, cols=2) + table.cell(0, 0).merge(table.cell(0, 1)).text = "Only once" + doc.save(docx_file) + + def fail_on_expansion(_row: _Row): + pytest.fail("row.cells expands gridSpan before extraction can apply a bound") + + monkeypatch.setattr(_Row, "cells", property(fail_on_expansion)) + + assert extract_text(docx_file) == "Only once" + + def test_extract_text_docx_bounds_physical_table_cells( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ): + """Large tables fail safely even when their text output would be empty.""" + from docx import Document + + from nanobot.utils import document as document_utils + + docx_file = tmp_path / "too-many-cells.docx" + doc = Document() + doc.add_table(rows=1, cols=2) + doc.save(docx_file) + monkeypatch.setattr(document_utils, "_MAX_DOCX_TABLE_CELLS", 1) + + result = extract_text(docx_file) + + assert result is not None + assert result.startswith("[error: unsafe DOCX:") + + def test_extract_text_docx_bounds_table_nesting( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ): + """Deeply nested tables fail safely instead of recursing without a bound.""" + from docx import Document + + from nanobot.utils import document as document_utils + + docx_file = tmp_path / "nested-too-deep.docx" + doc = Document() + outer_cell = doc.add_table(rows=1, cols=1).cell(0, 0) + outer_cell.add_table(rows=1, cols=1).cell(0, 0).text = "Nested" + doc.save(docx_file) + monkeypatch.setattr(document_utils, "_MAX_DOCX_TABLE_DEPTH", 1) + + result = extract_text(docx_file) + + assert result is not None + assert result.startswith("[error: unsafe DOCX:") + def test_extract_text_docx_empty(self, tmp_path: Path): """Test extracting text from an empty .docx file.""" from docx import Document