mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-14 08:00:10 +00:00
e09b2d48df
normalize_filename accepts names up to 255 UTF-8 bytes, but claim_unique_filename appended _N to the stem without re-checking the budget. A duplicate at maximum length therefore produced a 257-byte name, and the write path (open_upload_file_no_symlink -> normalize_filename) rejected it with ValueError. In the Gateway upload route that error falls into the generic handler: the whole request fails with a 500 and files already written in the same batch are rolled back — including unrelated ones. The same helper backs the Feishu and DingTalk channel downloads and client-side attachment staging. Truncate the stem on a UTF-8 code-point boundary when appending the dedupe tag would exceed 255 bytes, so the result always round-trips through normalize_filename. Names short enough to fit keep the exact dedupe shape they had before. Tests: red on main, green here — - unit: max-length dedupe stays within the limit and round-trips; repeated collisions stay unique; multibyte stems truncate on a code-point boundary; short names keep the historical _N shape - router: a batch with a max-length duplicate now succeeds and keeps every file instead of failing with a 500 Co-authored-by: Terminator666666 <Terminator666666@users.noreply.github.com>
289 lines
11 KiB
Python
289 lines
11 KiB
Python
"""Tests for deerflow.uploads.manager — shared upload management logic."""
|
||
|
||
import errno
|
||
import os
|
||
from unittest.mock import patch
|
||
|
||
import pytest
|
||
|
||
from deerflow.uploads.manager import (
|
||
PathTraversalError,
|
||
UnsafeUploadPathError,
|
||
claim_unique_filename,
|
||
cleanup_stale_upload_staging_files,
|
||
delete_file_safe,
|
||
list_files_in_dir,
|
||
normalize_filename,
|
||
validate_path_traversal,
|
||
write_upload_file_no_symlink,
|
||
)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# normalize_filename
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestNormalizeFilename:
|
||
def test_safe_filename(self):
|
||
assert normalize_filename("report.pdf") == "report.pdf"
|
||
|
||
def test_strips_path_components(self):
|
||
assert normalize_filename("../../etc/passwd") == "passwd"
|
||
|
||
def test_rejects_empty(self):
|
||
with pytest.raises(ValueError, match="empty"):
|
||
normalize_filename("")
|
||
|
||
def test_rejects_dot_dot(self):
|
||
with pytest.raises(ValueError, match="unsafe"):
|
||
normalize_filename("..")
|
||
|
||
def test_strips_separators(self):
|
||
assert normalize_filename("path/to/file.txt") == "file.txt"
|
||
|
||
def test_dot_only(self):
|
||
with pytest.raises(ValueError, match="unsafe"):
|
||
normalize_filename(".")
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# claim_unique_filename
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestDeduplicateFilename:
|
||
def test_no_collision(self):
|
||
seen: set[str] = set()
|
||
assert claim_unique_filename("data.txt", seen) == "data.txt"
|
||
assert "data.txt" in seen
|
||
|
||
def test_single_collision(self):
|
||
seen = {"data.txt"}
|
||
assert claim_unique_filename("data.txt", seen) == "data_1.txt"
|
||
assert "data_1.txt" in seen
|
||
|
||
def test_triple_collision(self):
|
||
seen = {"data.txt", "data_1.txt", "data_2.txt"}
|
||
assert claim_unique_filename("data.txt", seen) == "data_3.txt"
|
||
assert "data_3.txt" in seen
|
||
|
||
def test_mutates_seen(self):
|
||
seen: set[str] = set()
|
||
claim_unique_filename("a.txt", seen)
|
||
claim_unique_filename("a.txt", seen)
|
||
assert seen == {"a.txt", "a_1.txt"}
|
||
|
||
def test_max_length_name_stays_within_filename_limit(self):
|
||
# A 255-byte name passes normalize_filename; the deduplicated name
|
||
# must not exceed that limit, or the write path rejects it.
|
||
name = "a" * 251 + ".txt"
|
||
seen = {name}
|
||
deduped = claim_unique_filename(name, seen)
|
||
assert deduped != name
|
||
assert deduped.endswith("_1.txt")
|
||
assert len(deduped.encode("utf-8")) <= 255
|
||
# The truncated result must round-trip through normalize_filename.
|
||
assert normalize_filename(deduped) == deduped
|
||
|
||
def test_max_length_collisions_stay_unique_across_truncation(self):
|
||
name = "a" * 251 + ".txt"
|
||
seen = {name}
|
||
first = claim_unique_filename(name, seen)
|
||
second = claim_unique_filename(name, seen)
|
||
assert first != second
|
||
assert len(second.encode("utf-8")) <= 255
|
||
|
||
def test_multibyte_stem_is_truncated_on_a_codepoint_boundary(self):
|
||
# 85 CJK chars × 3 bytes = 255 bytes.
|
||
name = "深" * 85
|
||
seen = {name}
|
||
deduped = claim_unique_filename(name, seen)
|
||
assert len(deduped.encode("utf-8")) <= 255
|
||
assert deduped.endswith("_1")
|
||
# No replacement characters / decode artifacts.
|
||
deduped.encode("utf-8").decode("utf-8")
|
||
|
||
def test_short_names_keep_existing_dedupe_shape(self):
|
||
seen = {"data.txt"}
|
||
assert claim_unique_filename("data.txt", seen) == "data_1.txt"
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# validate_path_traversal
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestValidatePathTraversal:
|
||
def test_inside_base_ok(self, tmp_path):
|
||
child = tmp_path / "file.txt"
|
||
child.touch()
|
||
validate_path_traversal(child, tmp_path) # no exception
|
||
|
||
def test_outside_base_raises(self, tmp_path):
|
||
outside = tmp_path / ".." / "evil.txt"
|
||
with pytest.raises(PathTraversalError, match="traversal"):
|
||
validate_path_traversal(outside, tmp_path)
|
||
|
||
def test_symlink_escape(self, tmp_path):
|
||
target = tmp_path.parent / "secret.txt"
|
||
target.touch()
|
||
link = tmp_path / "escape"
|
||
try:
|
||
link.symlink_to(target)
|
||
except OSError as exc:
|
||
if getattr(exc, "winerror", None) == 1314:
|
||
pytest.skip("symlink creation requires Developer Mode or elevated privileges on Windows")
|
||
raise
|
||
with pytest.raises(PathTraversalError, match="traversal"):
|
||
validate_path_traversal(link, tmp_path)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# write_upload_file_no_symlink
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestWriteUploadFileNoSymlink:
|
||
def test_writes_new_file(self, tmp_path):
|
||
dest = write_upload_file_no_symlink(tmp_path, "notes.txt", b"hello")
|
||
|
||
assert dest == tmp_path / "notes.txt"
|
||
assert dest.read_bytes() == b"hello"
|
||
|
||
def test_overwrites_existing_regular_file_with_single_link(self, tmp_path):
|
||
dest = tmp_path / "notes.txt"
|
||
dest.write_bytes(b"old contents")
|
||
assert os.stat(dest).st_nlink == 1
|
||
|
||
result = write_upload_file_no_symlink(tmp_path, "notes.txt", b"new contents")
|
||
|
||
assert result == dest
|
||
assert dest.read_bytes() == b"new contents"
|
||
assert os.stat(dest).st_nlink == 1
|
||
|
||
def test_fallback_without_no_follow_support_succeeds(self, tmp_path, monkeypatch):
|
||
monkeypatch.delattr(os, "O_NOFOLLOW", raising=False)
|
||
|
||
# When O_NOFOLLOW is absent (Windows), the function falls back to
|
||
# a dual-lstat + fstat approach and succeeds.
|
||
result = write_upload_file_no_symlink(tmp_path, "notes.txt", b"hello")
|
||
assert result == tmp_path / "notes.txt"
|
||
assert (tmp_path / "notes.txt").read_bytes() == b"hello"
|
||
|
||
def test_open_uses_nonblocking_flag_when_available(self, tmp_path):
|
||
if not hasattr(os, "O_NONBLOCK"):
|
||
pytest.skip("O_NONBLOCK not available on this platform")
|
||
with patch("deerflow.uploads.manager.os.open", side_effect=OSError(errno.ENXIO, "no reader")) as open_mock:
|
||
with pytest.raises(UnsafeUploadPathError, match="Unsafe upload destination"):
|
||
write_upload_file_no_symlink(tmp_path, "pipe.txt", b"hello")
|
||
|
||
flags = open_mock.call_args.args[1]
|
||
assert flags & os.O_NONBLOCK
|
||
|
||
@pytest.mark.parametrize("open_errno", [errno.ENXIO, errno.EAGAIN])
|
||
def test_nonblocking_special_file_open_errors_are_unsafe(self, tmp_path, open_errno):
|
||
if not hasattr(os, "O_NONBLOCK"):
|
||
pytest.skip("O_NONBLOCK not available on this platform")
|
||
with patch("deerflow.uploads.manager.os.open", side_effect=OSError(open_errno, "would block")):
|
||
with pytest.raises(UnsafeUploadPathError, match="Unsafe upload destination"):
|
||
write_upload_file_no_symlink(tmp_path, "pipe.txt", b"hello")
|
||
|
||
assert not (tmp_path / "pipe.txt").exists()
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# list_files_in_dir
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestListFilesInDir:
|
||
def test_empty_dir(self, tmp_path):
|
||
result = list_files_in_dir(tmp_path)
|
||
assert result == {"files": [], "count": 0}
|
||
|
||
def test_nonexistent_dir(self, tmp_path):
|
||
result = list_files_in_dir(tmp_path / "nope")
|
||
assert result == {"files": [], "count": 0}
|
||
|
||
def test_multiple_files_sorted(self, tmp_path):
|
||
(tmp_path / "b.txt").write_text("b")
|
||
(tmp_path / "a.txt").write_text("a")
|
||
result = list_files_in_dir(tmp_path)
|
||
assert result["count"] == 2
|
||
assert result["files"][0]["filename"] == "a.txt"
|
||
assert result["files"][1]["filename"] == "b.txt"
|
||
for f in result["files"]:
|
||
assert set(f.keys()) == {"filename", "size", "path", "extension", "modified"}
|
||
|
||
def test_ignores_subdirectories(self, tmp_path):
|
||
(tmp_path / "file.txt").write_text("data")
|
||
(tmp_path / "subdir").mkdir()
|
||
result = list_files_in_dir(tmp_path)
|
||
assert result["count"] == 1
|
||
assert result["files"][0]["filename"] == "file.txt"
|
||
|
||
def test_filters_only_upload_staging_files(self, tmp_path):
|
||
(tmp_path / ".env").write_text("intentional dotfile")
|
||
(tmp_path / ".upload-active.part").write_text("partial")
|
||
(tmp_path / ".upload-note.txt").write_text("intentional upload")
|
||
(tmp_path / "draft.part").write_text("intentional upload")
|
||
(tmp_path / "visible.txt").write_text("visible")
|
||
|
||
result = list_files_in_dir(tmp_path)
|
||
|
||
assert result["count"] == 4
|
||
assert [f["filename"] for f in result["files"]] == [".env", ".upload-note.txt", "draft.part", "visible.txt"]
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# cleanup_stale_upload_staging_files
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestCleanupStaleUploadStagingFiles:
|
||
def test_removes_only_stale_staging_files_from_all_upload_layouts(self, tmp_path):
|
||
legacy_uploads = tmp_path / "threads" / "thread-legacy" / "user-data" / "uploads"
|
||
user_uploads = tmp_path / "users" / "owner-1" / "threads" / "thread-owned" / "user-data" / "uploads"
|
||
unrelated_uploads = tmp_path / "misc" / "thread-other" / "user-data" / "uploads"
|
||
for uploads_dir in (legacy_uploads, user_uploads, unrelated_uploads):
|
||
uploads_dir.mkdir(parents=True)
|
||
|
||
(legacy_uploads / ".upload-old.part").write_text("legacy partial")
|
||
(user_uploads / ".upload-new.part").write_text("user partial")
|
||
(unrelated_uploads / ".upload-ignore.part").write_text("outside layout")
|
||
(legacy_uploads / ".env").write_text("intentional dotfile")
|
||
(legacy_uploads / ".upload-note.txt").write_text("intentional upload")
|
||
(legacy_uploads / "draft.part").write_text("intentional upload")
|
||
|
||
removed = cleanup_stale_upload_staging_files(tmp_path)
|
||
|
||
assert removed == 2
|
||
assert not (legacy_uploads / ".upload-old.part").exists()
|
||
assert not (user_uploads / ".upload-new.part").exists()
|
||
assert (unrelated_uploads / ".upload-ignore.part").exists()
|
||
assert (legacy_uploads / ".env").exists()
|
||
assert (legacy_uploads / ".upload-note.txt").exists()
|
||
assert (legacy_uploads / "draft.part").exists()
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# delete_file_safe
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestDeleteFileSafe:
|
||
def test_delete_existing_file(self, tmp_path):
|
||
f = tmp_path / "test.txt"
|
||
f.write_text("data")
|
||
result = delete_file_safe(tmp_path, "test.txt")
|
||
assert result["success"] is True
|
||
assert not f.exists()
|
||
|
||
def test_delete_nonexistent_raises(self, tmp_path):
|
||
with pytest.raises(FileNotFoundError):
|
||
delete_file_safe(tmp_path, "nope.txt")
|
||
|
||
def test_delete_traversal_raises(self, tmp_path):
|
||
with pytest.raises(PathTraversalError, match="traversal"):
|
||
delete_file_safe(tmp_path, "../outside.txt")
|