generated from kgod/ai-review-template
自内部仓库剥离深度优化与 RAG 知识库后的交付版本: - Builder 对话式简历生成(FSM + 意图路由 LLM 兜底增强) - 条目级轻度优化:事实覆盖门禁 + STAR/bullet 修复链,功能/简介/成果与技术栈同级保护 - 简历导入:DOCX/PDF 解析、结构归一、手机号脱敏 - PostgreSQL 运行时 + Alembic 迁移链 Co-Authored-By: Claude <noreply@anthropic.com>
50 lines
2.2 KiB
Python
50 lines
2.2 KiB
Python
"""Import upload guards: decompression-bomb docx must be rejected before parsing (H2 DoS)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import io
|
|
import zipfile
|
|
|
|
import pytest
|
|
from docx import Document
|
|
|
|
from app.document_extractors import ImportExtractionError, extract_text
|
|
|
|
CT = ('<?xml version="1.0"?><Types xmlns="http://schemas.openxmlformats.org/package/2006/content-types">'
|
|
'<Default Extension="xml" ContentType="application/xml"/>'
|
|
'<Default Extension="rels" ContentType="application/vnd.openxmlformats-package.relationships+xml"/>'
|
|
'<Override PartName="/word/document.xml" ContentType="application/vnd.openxmlformats-officedocument.wordprocessingml.document.main+xml"/></Types>')
|
|
RELS = ('<?xml version="1.0"?><Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">'
|
|
'<Relationship Id="rId1" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/officeDocument" Target="word/document.xml"/></Relationships>')
|
|
|
|
|
|
def _custom_docx(document_xml: str) -> bytes:
|
|
buffer = io.BytesIO()
|
|
with zipfile.ZipFile(buffer, "w", zipfile.ZIP_DEFLATED) as zf:
|
|
zf.writestr("[Content_Types].xml", CT)
|
|
zf.writestr("_rels/.rels", RELS)
|
|
zf.writestr("word/document.xml", document_xml)
|
|
return buffer.getvalue()
|
|
|
|
|
|
def test_decompression_bomb_docx_is_rejected_before_parsing() -> None:
|
|
"""~20MB decompressed XML must fail fast instead of burning CPU in the parser."""
|
|
para = "<w:p><w:r><w:t>放大攻击</w:t></w:r></w:p>"
|
|
body = para * (20 * 1024 * 1024 // len(para.encode()))
|
|
bomb = _custom_docx('<?xml version="1.0" encoding="UTF-8"?>'
|
|
'<w:document xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main">'
|
|
f"<w:body>{body}</w:body></w:document>")
|
|
assert len(bomb) < 1024 * 1024 # small compressed payload is the point of the attack
|
|
|
|
with pytest.raises(ImportExtractionError):
|
|
extract_text(extension=".docx", content=bomb)
|
|
|
|
|
|
def test_normal_docx_still_parses() -> None:
|
|
document = Document()
|
|
document.add_paragraph("张三 后端工程师")
|
|
buffer = io.BytesIO()
|
|
document.save(buffer)
|
|
|
|
assert "张三" in extract_text(extension=".docx", content=buffer.getvalue())
|