fix(documents): bound nested DOCX table parsing

This commit is contained in:
Xubin Ren
2026-07-23 16:53:59 +08:00
parent 60ab580f8b
commit fc9d17eb7b
2 changed files with 125 additions and 13 deletions
+45 -13
View File
@@ -1,6 +1,7 @@
"""Document text extraction utilities for nanobot.""" """Document text extraction utilities for nanobot."""
import mimetypes import mimetypes
from collections.abc import Iterator
from dataclasses import dataclass from dataclasses import dataclass
from pathlib import Path from pathlib import Path
from zipfile import BadZipFile, ZipFile 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_ARCHIVE_MEMBERS = 10_000
_MAX_OFFICE_UNCOMPRESSED_SIZE = 256 * 1024 * 1024 # 256 MB _MAX_OFFICE_UNCOMPRESSED_SIZE = 256 * 1024 * 1024 # 256 MB
_MAX_OFFICE_MEMBER_SIZE = 128 * 1024 * 1024 # 128 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_CONTENT_STREAM_SIZE = 32 * 1024 * 1024 # 32 MB per page
_MAX_PDF_ATTACHMENT_PAGES = 100 _MAX_PDF_ATTACHMENT_PAGES = 100
@@ -87,6 +90,10 @@ class PdfPageRangeError(Exception):
"""Raised when a requested PDF page range is invalid.""" """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) @dataclass(frozen=True, slots=True)
class PdfExtraction: class PdfExtraction:
text: str text: str
@@ -210,7 +217,7 @@ def _extract_docx(path: Path) -> str:
"""Extract text from DOCX using python-docx.""" """Extract text from DOCX using python-docx."""
try: try:
from docx import Document as DocxDocument from docx import Document as DocxDocument
from docx.table import Table from docx.table import Table, _Cell
from docx.text.paragraph import Paragraph from docx.text.paragraph import Paragraph
except ImportError: except ImportError:
return "[error: python-docx not installed]" return "[error: python-docx not installed]"
@@ -219,6 +226,39 @@ def _extract_docx(path: Path) -> str:
return error return error
doc = DocxDocument(path) doc = DocxDocument(path)
collector = _TextCollector(_MAX_TEXT_LENGTH) 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(): for block in doc.iter_inner_content():
if isinstance(block, Paragraph): if isinstance(block, Paragraph):
text = block.text.strip() text = block.text.strip()
@@ -228,22 +268,14 @@ def _extract_docx(path: Path) -> str:
if not isinstance(block, Table): if not isinstance(block, Table):
continue continue
first_row = True first_row = True
for row in block.rows: for row_text in table_rows(block, 1):
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
separator = "\n\n" if first_row else "\n" separator = "\n\n" if first_row else "\n"
first_row = False 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()
return collector.render() return collector.render()
except DocxSafetyError as e:
return f"[error: unsafe DOCX: {e!s}]"
except Exception as e: except Exception as e:
logger.exception("Failed to extract DOCX {}", path) logger.exception("Failed to extract DOCX {}", path)
return f"[error: failed to extract DOCX: {e!s}]" return f"[error: failed to extract DOCX: {e!s}]"
+80
View File
@@ -197,6 +197,86 @@ class TestExtractText:
assert result.index("Applicant details") < result.index("Name\tAda Lovelace") assert result.index("Applicant details") < result.index("Name\tAda Lovelace")
assert result.index("Analytical Engine") < result.index("End of form") 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): def test_extract_text_docx_empty(self, tmp_path: Path):
"""Test extracting text from an empty .docx file.""" """Test extracting text from an empty .docx file."""
from docx import Document from docx import Document