fix: remove duplicate /health route, escape extracted text, validate uploads, handle binary decode

- main.py: remove duplicate /health endpoint definition
- text_extract.py: HTML-escape extracted plain text before indexing so raw
  markup in uploaded documents can't render as live HTML via search snippets
  (stored XSS fix)
- document_service.py: base64-encode content when UTF-8 decoding fails
  instead of raising UnicodeDecodeError (500) on binary files; add upload
  validation for file extension (allowlist) and size (10MB max)
- routers/documents.py: map new validation errors to HTTP 400 with a clear
  detail message instead of letting them 500
- add/extend tests covering escaping, binary get_document, and upload
  validation rejection + acceptance paths
This commit is contained in:
Tianyang 2026-08-01 08:59:34 +08:00
parent ad16c1732d
commit 13789cc0b4
7 changed files with 150 additions and 12 deletions

View File

@ -1,3 +1,4 @@
import base64
import logging
from datetime import datetime, timezone
@ -5,11 +6,33 @@ from app.text_extract import html_to_text, markdown_to_text
logger = logging.getLogger(__name__)
ALLOWED_UPLOAD_EXTENSIONS = {".md", ".html", ".htm", ".png", ".jpg", ".jpeg", ".gif", ".webp"}
MAX_UPLOAD_SIZE_BYTES = 10 * 1024 * 1024
class DocumentNotFoundError(Exception):
pass
class UnsupportedFileTypeError(Exception):
def __init__(self, extension: str):
self.extension = extension
super().__init__(f"unsupported file type: {extension}")
class FileTooLargeError(Exception):
def __init__(self, max_size_bytes: int):
self.max_size_bytes = max_size_bytes
super().__init__(f"file too large: max {max_size_bytes // (1024 * 1024)}MB")
def _extension_of(key: str) -> str:
name = key.rsplit("/", 1)[-1]
if "." not in name:
return ""
return "." + name.rsplit(".", 1)[-1].lower()
def _guess_content_type(key: str) -> str:
if key.endswith(".md"):
return "text/markdown"
@ -49,10 +72,17 @@ class DocumentService:
data = self._minio_client.get_object(key)
except FileNotFoundError as exc:
raise DocumentNotFoundError(key) from exc
try:
content = data.decode("utf-8")
encoding = "utf-8"
except UnicodeDecodeError:
content = base64.b64encode(data).decode("ascii")
encoding = "base64"
return {
"key": key,
"content": data.decode("utf-8"),
"content": content,
"contentType": _guess_content_type(key),
"encoding": encoding,
}
def _extract_text(self, content: bytes, content_type: str) -> str:
@ -82,6 +112,11 @@ class DocumentService:
logger.exception("Failed to delete search index entry for %s", key)
def upload_document(self, key: str, data: bytes, content_type: str) -> dict:
extension = _extension_of(key)
if extension not in ALLOWED_UPLOAD_EXTENSIONS:
raise UnsupportedFileTypeError(extension)
if len(data) > MAX_UPLOAD_SIZE_BYTES:
raise FileTooLargeError(MAX_UPLOAD_SIZE_BYTES)
self.save_document(key, data, content_type)
name = key.rsplit("/", 1)[-1]
return {"key": key, "name": name, "size": len(data)}

View File

@ -18,11 +18,6 @@ app.include_router(documents.router)
app.include_router(search.router)
@app.get("/health")
def health():
return {"status": "ok"}
@app.get("/health")
def health():
return {"status": "ok"}

View File

@ -2,7 +2,12 @@ from fastapi import APIRouter, Depends, File, Form, HTTPException, UploadFile
from pydantic import BaseModel
from app.dependencies import get_document_service
from app.document_service import DocumentNotFoundError, DocumentService
from app.document_service import (
DocumentNotFoundError,
DocumentService,
FileTooLargeError,
UnsupportedFileTypeError,
)
router = APIRouter()
@ -47,4 +52,7 @@ async def upload_document(
data = await file.read()
key = f"{path.rstrip('/')}/{file.filename}" if path else file.filename
content_type = file.content_type or "application/octet-stream"
return service.upload_document(key, data, content_type)
try:
return service.upload_document(key, data, content_type)
except (UnsupportedFileTypeError, FileTooLargeError) as exc:
raise HTTPException(status_code=400, detail=str(exc))

View File

@ -1,3 +1,4 @@
import html
import re
from html.parser import HTMLParser
@ -19,7 +20,7 @@ def html_to_text(text: str) -> str:
parser.feed(text)
parser.close()
raw = parser.get_text()
return re.sub(r"\s+", " ", raw).strip()
return html.escape(re.sub(r"\s+", " ", raw).strip())
def markdown_to_text(text: str) -> str:
@ -35,4 +36,4 @@ def markdown_to_text(text: str) -> str:
result = re.sub(r"^\s*([-*+]|\d+\.)\s+", "", result, flags=re.MULTILINE)
result = re.sub(r"^\s*>\s*", "", result, flags=re.MULTILINE)
result = re.sub(r"\s+", " ", result).strip()
return result
return html.escape(result)

View File

@ -1,6 +1,12 @@
import pytest
from app.document_service import DocumentService, DocumentNotFoundError
from app.document_service import (
DocumentService,
DocumentNotFoundError,
FileTooLargeError,
MAX_UPLOAD_SIZE_BYTES,
UnsupportedFileTypeError,
)
class FakeMinioClient:
@ -83,6 +89,29 @@ def test_get_document_raises_not_found_for_missing_key():
service.get_document("不存在/文档.md")
def test_get_document_returns_base64_content_for_binary_data_without_raising():
minio_client = FakeMinioClient()
binary_data = b"\x89PNG\r\n\x1a\n\x00\x01\xff\xfe"
minio_client.stored["产品文档/图片.png"] = binary_data
service = DocumentService(minio_client, FakeSearchClient())
result = service.get_document("产品文档/图片.png")
assert result["encoding"] == "base64"
import base64
assert base64.b64decode(result["content"]) == binary_data
def test_get_document_returns_utf8_encoding_for_text_content():
minio_client = FakeMinioClient()
minio_client.stored["产品文档/架构设计.md"] = "# 架构设计".encode("utf-8")
service = DocumentService(minio_client, FakeSearchClient())
result = service.get_document("产品文档/架构设计.md")
assert result["encoding"] == "utf-8"
def test_save_document_writes_to_minio_then_indexes_content():
minio_client = FakeMinioClient()
search_client = FakeSearchClient()
@ -144,3 +173,35 @@ def test_upload_document_writes_file_and_returns_metadata():
assert result == {"key": "产品文档/新上传.md", "name": "新上传.md", "size": len(data)}
assert minio_client.stored["产品文档/新上传.md"] == data
def test_upload_document_rejects_unsupported_file_extension():
minio_client = FakeMinioClient()
service = DocumentService(minio_client, FakeSearchClient())
with pytest.raises(UnsupportedFileTypeError):
service.upload_document("产品文档/病毒.exe", b"data", "application/octet-stream")
assert "产品文档/病毒.exe" not in minio_client.stored
def test_upload_document_rejects_oversized_file():
minio_client = FakeMinioClient()
service = DocumentService(minio_client, FakeSearchClient())
data = b"x" * (MAX_UPLOAD_SIZE_BYTES + 1)
with pytest.raises(FileTooLargeError):
service.upload_document("产品文档/大文件.md", data, "text/markdown")
assert "产品文档/大文件.md" not in minio_client.stored
def test_upload_document_accepts_allowed_image_extension():
minio_client = FakeMinioClient()
service = DocumentService(minio_client, FakeSearchClient())
data = b"\x89PNG\r\n\x1a\n"
result = service.upload_document("产品文档/图片.png", data, "image/png")
assert result == {"key": "产品文档/图片.png", "name": "图片.png", "size": len(data)}
assert minio_client.stored["产品文档/图片.png"] == data

View File

@ -3,7 +3,7 @@ from fastapi.testclient import TestClient
from app.main import app
from app.dependencies import get_document_service
from app.document_service import DocumentNotFoundError
from app.document_service import DocumentNotFoundError, FileTooLargeError, UnsupportedFileTypeError
class FakeDocumentService:
@ -33,6 +33,10 @@ class FakeDocumentService:
self.deleted.append(key)
def upload_document(self, key, data, content_type):
if key.endswith(".exe"):
raise UnsupportedFileTypeError(".exe")
if len(data) > 1024:
raise FileTooLargeError(1024)
self.uploaded.append((key, data, content_type))
return {"key": key, "name": key.rsplit("/", 1)[-1], "size": len(data)}
@ -95,3 +99,23 @@ def test_upload_document_returns_201_with_metadata(client):
assert body["key"] == "产品文档/新文档.md"
assert body["name"] == "新文档.md"
assert body["size"] == len(b"# new content")
def test_upload_document_returns_400_for_unsupported_file_type(client):
response = client.post(
"/api/upload",
data={"path": "产品文档"},
files={"file": ("恶意程序.exe", b"binary", "application/octet-stream")},
)
assert response.status_code == 400
assert "unsupported file type" in response.json()["detail"]
def test_upload_document_returns_400_for_oversized_file(client):
response = client.post(
"/api/upload",
data={"path": "产品文档"},
files={"file": ("大文件.md", b"x" * 2000, "text/markdown")},
)
assert response.status_code == 400
assert "file too large" in response.json()["detail"]

View File

@ -60,6 +60,20 @@ def test_markdown_to_text_strips_links_and_list_markers():
assert "在线编辑并保存" in result
def test_html_to_text_escapes_raw_tag_like_text_content():
html_doc = "<p>点击查看 &lt;img src=x onerror=alert(1)&gt; 示例</p>"
result = html_to_text(html_doc)
assert "<img" not in result
assert "&lt;img src=x onerror=alert(1)&gt;" in result
def test_markdown_to_text_escapes_raw_html_in_markdown_body():
markdown = "# 说明\n\n请勿插入 <img src=x onerror=alert(1)> 这样的标签。"
result = markdown_to_text(markdown)
assert "<img" not in result
assert "&lt;img src=x onerror=alert(1)&gt;" in result
def test_markdown_to_text_strips_code_fences_and_inline_code():
markdown = """使用 `pip install -r requirements.txt` 安装依赖。