257 lines
10 KiB
Python
257 lines
10 KiB
Python
from __future__ import annotations
|
|
|
|
import io
|
|
import zipfile
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
from fastapi import HTTPException, UploadFile
|
|
|
|
from app.gateway.routers import uploads as uploads_router
|
|
from deerflow.persistence.upload_policy import MemoryUploadPolicyStore
|
|
from deerflow.uploads.policy import classify_upload_filename, default_upload_policy
|
|
from deerflow.uploads.validation import UploadContentError, validate_upload_content
|
|
|
|
|
|
def _ooxml_bytes(prefix: str) -> bytes:
|
|
output = io.BytesIO()
|
|
with zipfile.ZipFile(output, "w") as archive:
|
|
archive.writestr("[Content_Types].xml", "<Types />")
|
|
archive.writestr(f"{prefix}content.xml", "<root />")
|
|
return output.getvalue()
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("extension", "content"),
|
|
[
|
|
(".txt", "中文文本".encode()),
|
|
(".md", b"# heading"),
|
|
(".csv", b"a,b\n1,2\n"),
|
|
(".pdf", b"%PDF-1.7\n"),
|
|
(".png", b"\x89PNG\r\n\x1a\nrest"),
|
|
(".jpg", b"\xff\xd8\xff\xe0rest"),
|
|
(".doc", b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1rest"),
|
|
(".ppt", b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1rest"),
|
|
(".xls", b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1rest"),
|
|
(".docx", _ooxml_bytes("word/")),
|
|
(".pptx", _ooxml_bytes("ppt/")),
|
|
(".xlsx", _ooxml_bytes("xl/")),
|
|
],
|
|
)
|
|
def test_supported_signatures(extension: str, content: bytes, tmp_path: Path):
|
|
path = tmp_path / f"sample{extension}"
|
|
path.write_bytes(content)
|
|
validate_upload_content(path, extension)
|
|
|
|
|
|
def test_spoofed_extension_is_rejected(tmp_path: Path):
|
|
path = tmp_path / "fake.png"
|
|
path.write_bytes(b"this is not an image")
|
|
with pytest.raises(UploadContentError):
|
|
validate_upload_content(path, ".png")
|
|
|
|
|
|
async def _request_with_policy(policy: dict):
|
|
store = MemoryUploadPolicyStore()
|
|
await store.ensure_seed(policy)
|
|
return SimpleNamespace(app=SimpleNamespace(state=SimpleNamespace(upload_policy_store=store)))
|
|
|
|
|
|
def _patch_upload_environment(monkeypatch, uploads_dir: Path, *, sandbox=None):
|
|
uploads_dir.mkdir(parents=True, exist_ok=True)
|
|
monkeypatch.setattr(uploads_router, "ensure_uploads_dir", lambda _thread_id: uploads_dir)
|
|
monkeypatch.setattr(
|
|
uploads_router,
|
|
"get_paths",
|
|
lambda: SimpleNamespace(sandbox_uploads_dir=lambda _thread_id, user_id=None: uploads_dir),
|
|
)
|
|
|
|
async def path_user(_thread_id):
|
|
return "user-1"
|
|
|
|
monkeypatch.setattr(uploads_router, "aresolve_path_user_id", path_user)
|
|
if sandbox is None:
|
|
provider = SimpleNamespace(uses_thread_data_mounts=True)
|
|
else:
|
|
provider = SimpleNamespace(
|
|
uses_thread_data_mounts=False,
|
|
acquire=lambda _thread_id: "sandbox-1",
|
|
get=lambda _sandbox_id: sandbox,
|
|
)
|
|
monkeypatch.setattr(uploads_router, "get_sandbox_provider", lambda: provider)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_upload_uses_category_limit_and_returns_category(monkeypatch, tmp_path: Path):
|
|
policy = default_upload_policy()
|
|
policy["groups"]["document"]["max_bytes"] = 3
|
|
request = await _request_with_policy(policy)
|
|
uploads_dir = tmp_path / "uploads"
|
|
_patch_upload_environment(monkeypatch, uploads_dir)
|
|
config = SimpleNamespace(uploads={"auto_convert_documents": False})
|
|
|
|
too_large = UploadFile(filename="large.txt", file=io.BytesIO(b"four"))
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
await uploads_router.upload_files.__wrapped__("thread-1", request, [too_large], config)
|
|
assert exc_info.value.status_code == 413
|
|
assert not (uploads_dir / "large.txt").exists()
|
|
|
|
valid = UploadFile(filename="ok.txt", file=io.BytesIO(b"abc"))
|
|
response = await uploads_router.upload_files.__wrapped__("thread-1", request, [valid], config)
|
|
assert response.files[0]["category"] == "document"
|
|
assert response.files[0]["detected_type"] == "txt"
|
|
assert response.policy_revision == 1
|
|
assert (uploads_dir / "ok.txt").read_bytes() == b"abc"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_batch_validation_preserves_existing_same_name(monkeypatch, tmp_path: Path):
|
|
request = await _request_with_policy(default_upload_policy())
|
|
uploads_dir = tmp_path / "uploads"
|
|
_patch_upload_environment(monkeypatch, uploads_dir)
|
|
(uploads_dir / "same.txt").write_bytes(b"old")
|
|
config = SimpleNamespace(uploads={"auto_convert_documents": False})
|
|
files = [
|
|
UploadFile(filename="same.txt", file=io.BytesIO(b"new")),
|
|
UploadFile(filename="bad.png", file=io.BytesIO(b"not-png")),
|
|
]
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
await uploads_router.upload_files.__wrapped__("thread-1", request, files, config)
|
|
assert exc_info.value.status_code == 415
|
|
assert (uploads_dir / "same.txt").read_bytes() == b"old"
|
|
assert not (uploads_dir / "bad.png").exists()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_remote_sync_failure_rolls_back_committed_files(monkeypatch, tmp_path: Path):
|
|
request = await _request_with_policy(default_upload_policy())
|
|
uploads_dir = tmp_path / "uploads"
|
|
(uploads_dir / "same.txt").parent.mkdir(parents=True)
|
|
(uploads_dir / "same.txt").write_bytes(b"old")
|
|
|
|
class FailingSandbox:
|
|
def update_file(self, _path, _content):
|
|
raise OSError("remote write failed")
|
|
|
|
_patch_upload_environment(monkeypatch, uploads_dir, sandbox=FailingSandbox())
|
|
config = SimpleNamespace(uploads={"auto_convert_documents": False})
|
|
file = UploadFile(filename="same.txt", file=io.BytesIO(b"new"))
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
await uploads_router.upload_files.__wrapped__("thread-1", request, [file], config)
|
|
assert exc_info.value.status_code == 500
|
|
assert (uploads_dir / "same.txt").read_bytes() == b"old"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_count_and_extension_are_enforced_before_writing(monkeypatch, tmp_path: Path):
|
|
policy = default_upload_policy()
|
|
policy["max_files_per_message"] = 2
|
|
request = await _request_with_policy(policy)
|
|
uploads_dir = tmp_path / "uploads"
|
|
_patch_upload_environment(monkeypatch, uploads_dir)
|
|
config = SimpleNamespace(uploads={"auto_convert_documents": False})
|
|
|
|
files = [UploadFile(filename=f"{index}.txt", file=io.BytesIO(b"x")) for index in range(3)]
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
await uploads_router.upload_files.__wrapped__("thread-1", request, files, config)
|
|
assert exc_info.value.status_code == 413
|
|
|
|
forbidden = UploadFile(filename="malware.exe", file=io.BytesIO(b"MZ"))
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
await uploads_router.upload_files.__wrapped__("thread-1", request, [forbidden], config)
|
|
assert exc_info.value.status_code == 415
|
|
assert list(uploads_dir.iterdir()) == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_default_accepts_twenty_and_rejects_twenty_one(monkeypatch, tmp_path: Path):
|
|
request = await _request_with_policy(default_upload_policy())
|
|
uploads_dir = tmp_path / "uploads"
|
|
_patch_upload_environment(monkeypatch, uploads_dir)
|
|
config = SimpleNamespace(uploads={"auto_convert_documents": False})
|
|
|
|
twenty = [UploadFile(filename=f"ok-{index}.txt", file=io.BytesIO(b"x")) for index in range(20)]
|
|
response = await uploads_router.upload_files.__wrapped__("thread-1", request, twenty, config)
|
|
assert len(response.files) == 20
|
|
|
|
twenty_one = [UploadFile(filename=f"no-{index}.txt", file=io.BytesIO(b"x")) for index in range(21)]
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
await uploads_router.upload_files.__wrapped__("thread-1", request, twenty_one, config)
|
|
assert exc_info.value.status_code == 413
|
|
assert not any((uploads_dir / f"no-{index}.txt").exists() for index in range(21))
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
("filename", "prefix"),
|
|
[
|
|
("document.txt", b"a"),
|
|
("picture.PNG", b"\x89PNG\r\n\x1a\n"),
|
|
("table.csv", b"a,b\n"),
|
|
],
|
|
)
|
|
async def test_each_category_exact_limit_and_limit_plus_one(filename: str, prefix: bytes, tmp_path: Path):
|
|
policy = default_upload_policy()
|
|
category, limit = classify_upload_filename(filename, policy) or (None, None)
|
|
assert category is not None and limit is not None
|
|
|
|
exact_dir = tmp_path / f"exact-{category}"
|
|
exact_dir.mkdir()
|
|
exact_body = prefix + b"x" * (limit - len(prefix))
|
|
_, size, total = await uploads_router._write_upload_file_with_limits(
|
|
UploadFile(filename=filename, file=io.BytesIO(exact_body)),
|
|
uploads_dir=exact_dir,
|
|
display_filename=filename,
|
|
max_single_file_size=limit,
|
|
max_total_size=100 * 1024 * 1024,
|
|
total_size=0,
|
|
)
|
|
assert size == total == limit
|
|
|
|
over_dir = tmp_path / f"over-{category}"
|
|
over_dir.mkdir()
|
|
over_body = prefix + b"x" * (limit + 1 - len(prefix))
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
await uploads_router._write_upload_file_with_limits(
|
|
UploadFile(filename=filename, file=io.BytesIO(over_body)),
|
|
uploads_dir=over_dir,
|
|
display_filename=filename,
|
|
max_single_file_size=limit,
|
|
max_total_size=100 * 1024 * 1024,
|
|
total_size=0,
|
|
)
|
|
assert exc_info.value.status_code == 413
|
|
assert not (over_dir / filename).exists()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_total_size_exact_boundary_and_plus_one(monkeypatch, tmp_path: Path):
|
|
policy = default_upload_policy()
|
|
policy["max_total_bytes"] = 10
|
|
policy["groups"]["document"]["max_bytes"] = 10
|
|
policy["groups"]["image"]["max_bytes"] = 10
|
|
policy["groups"]["spreadsheet"]["max_bytes"] = 5
|
|
policy["large_file"]["threshold_bytes"] = 10
|
|
request = await _request_with_policy(policy)
|
|
uploads_dir = tmp_path / "uploads"
|
|
_patch_upload_environment(monkeypatch, uploads_dir)
|
|
config = SimpleNamespace(uploads={"auto_convert_documents": False})
|
|
|
|
exact = [
|
|
UploadFile(filename="first.txt", file=io.BytesIO(b"1234")),
|
|
UploadFile(filename="second.txt", file=io.BytesIO(b"123456")),
|
|
]
|
|
response = await uploads_router.upload_files.__wrapped__("thread-1", request, exact, config)
|
|
assert len(response.files) == 2
|
|
|
|
over = [
|
|
UploadFile(filename="third.txt", file=io.BytesIO(b"1234")),
|
|
UploadFile(filename="fourth.txt", file=io.BytesIO(b"1234567")),
|
|
]
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
await uploads_router.upload_files.__wrapped__("thread-1", request, over, config)
|
|
assert exc_info.value.status_code == 413
|
|
assert not (uploads_dir / "third.txt").exists()
|
|
assert not (uploads_dir / "fourth.txt").exists()
|