347 lines
17 KiB
Python
347 lines
17 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import replace
|
|
import json
|
|
import inspect
|
|
import os
|
|
from pathlib import Path
|
|
import subprocess
|
|
import tarfile
|
|
import tempfile
|
|
import time
|
|
import tracemalloc
|
|
from types import SimpleNamespace
|
|
import unittest
|
|
from unittest.mock import patch
|
|
import zlib
|
|
|
|
from govoplan_core.security import bounded_process
|
|
from govoplan_core.settings import settings
|
|
from govoplan_files.backend.storage import archive_workers as workers
|
|
from govoplan_files.backend.storage.archives import extract_archive_upload, inspect_archive
|
|
from govoplan_files.backend.storage.common import FileStorageError
|
|
from test_archives import _zip_bytes
|
|
|
|
|
|
def _cpu_exhaustion_worker(payload: bytes) -> bytes:
|
|
# A real imported child operation proves the Files wrapper's CPU error path;
|
|
# the metadata tests below execute the actual archive parser, not this probe.
|
|
while True:
|
|
pass
|
|
|
|
|
|
def _extract(payload, **options):
|
|
return extract_archive_upload(
|
|
object(), tenant_id="tenant", owner_type="user", owner_id="owner",
|
|
user_id="owner", filename="fixture.zip", archive_data=payload,
|
|
folder="destination", campaign_id=None, **options,
|
|
)
|
|
|
|
|
|
class ArchiveWorkerTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.directories = []
|
|
self.children = []
|
|
real_directory = tempfile.TemporaryDirectory
|
|
real_popen = subprocess.Popen
|
|
|
|
def directory(*args, **kwargs):
|
|
context = real_directory(*args, **kwargs)
|
|
self.directories.append(Path(context.name))
|
|
return context
|
|
|
|
def popen(*args, **kwargs):
|
|
child = real_popen(*args, **kwargs)
|
|
self.children.append(child)
|
|
return child
|
|
|
|
self.enterContext(patch.object(workers.tempfile, "TemporaryDirectory", side_effect=directory))
|
|
self.enterContext(patch.object(bounded_process.subprocess, "Popen", side_effect=popen))
|
|
self.store = self.enterContext(patch(
|
|
"govoplan_files.backend.storage.archives.create_file_asset",
|
|
return_value=SimpleNamespace(asset=SimpleNamespace(id="asset")),
|
|
))
|
|
|
|
def tearDown(self):
|
|
self.assertTrue(all(not directory.exists() for directory in self.directories))
|
|
self.assertTrue(all(child.returncode is not None for child in self.children))
|
|
|
|
def test_parent_never_parses_metadata_or_extracts_and_one_child_handles_all_members(self):
|
|
payload = _zip_bytes({"one.txt": b"one", "two.txt": b"two"})
|
|
with (
|
|
patch("govoplan_files.backend.storage.archives._inspect_archive_content", side_effect=AssertionError("parent parser")),
|
|
patch("govoplan_files.backend.storage.archives._read_selected_zip_members", side_effect=AssertionError("parent decoder")),
|
|
):
|
|
self.assertEqual(2, inspect_archive(payload, filename="fixture.zip").file_count)
|
|
self.assertEqual(2, len(_extract(payload)))
|
|
self.assertEqual(2, len(self.children)) # one preview, one whole extraction
|
|
self.assertEqual([b"one", b"two"], [call.kwargs["data"] for call in self.store.call_args_list])
|
|
|
|
def test_hostile_pax_metadata_allocation_is_confined_before_any_store(self):
|
|
header = tarfile.TarInfo("pax")
|
|
header.type = tarfile.XHDTYPE
|
|
header.size = 1024 * 1024 * 1024
|
|
payload = header.tobuf() + b"\0" * 1024
|
|
with patch("govoplan_files.backend.storage.archives._open_tar", side_effect=AssertionError("parent TAR parser")):
|
|
with self.assertRaisesRegex(FileStorageError, "memory_limit|Invalid TAR|could not complete safely"):
|
|
inspect_archive(payload, filename="hostile.tar")
|
|
self.store.assert_not_called()
|
|
self.assertEqual(1, len(self.children))
|
|
# A rejected archive neither consumes the admission slot permanently nor
|
|
# prevents a later ordinary parse in the same parent process.
|
|
self.assertEqual(1, inspect_archive(_zip_bytes({"ok": b"ok"}), filename="ok.zip").file_count)
|
|
|
|
def test_compressed_pax_metadata_hits_real_child_memory_limit(self):
|
|
header = tarfile.TarInfo("pax")
|
|
header.type = tarfile.XHDTYPE
|
|
header.size = 1024 * 1024 * 1024
|
|
compressor = zlib.compressobj(1, wbits=31)
|
|
chunks = [compressor.compress(header.tobuf())]
|
|
block = b"\0" * (1024 * 1024)
|
|
# Build the hostile fixture incrementally: ~2.3 MiB compressed, never a
|
|
# 512 MiB parent allocation. TAR consumes PAX metadata before ordinary
|
|
# returned-member limits can inspect it; the child AS limit must win.
|
|
for _ in range(512):
|
|
chunks.append(compressor.compress(block))
|
|
chunks.append(compressor.flush())
|
|
payload = b"".join(chunks)
|
|
self.assertLess(len(payload), 3 * 1024 * 1024)
|
|
with patch("govoplan_files.backend.storage.archives._open_tar", side_effect=AssertionError("parent TAR parser")):
|
|
with self.assertRaisesRegex(FileStorageError, "memory_limit"):
|
|
inspect_archive(payload, filename="hostile.tar.gz")
|
|
self.store.assert_not_called()
|
|
self.assertEqual(1, len(self.children))
|
|
|
|
def test_real_child_timeout_reaps_and_removes_snapshot(self):
|
|
with patch.object(workers, "INSPECTION_LIMITS", replace(workers.INSPECTION_LIMITS, wall_seconds=0.001)):
|
|
with self.assertRaisesRegex(FileStorageError, "timeout"):
|
|
inspect_archive(_zip_bytes({"one": b"one"}), filename="fixture.zip")
|
|
self.assertEqual(1, len(self.children))
|
|
self.store.assert_not_called()
|
|
|
|
def test_real_child_cpu_limit_has_sanitized_files_error(self):
|
|
with (
|
|
patch.object(workers, "_inspect_worker", _cpu_exhaustion_worker),
|
|
patch.object(workers, "INSPECTION_LIMITS", replace(workers.INSPECTION_LIMITS, cpu_seconds=1, wall_seconds=10)),
|
|
):
|
|
with self.assertRaisesRegex(FileStorageError, "cpu_limit"):
|
|
inspect_archive(_zip_bytes({"one": b"one"}), filename="fixture.zip")
|
|
|
|
def test_real_preview_output_is_bounded(self):
|
|
with patch.object(workers, "INSPECTION_LIMITS", replace(workers.INSPECTION_LIMITS, output_bytes=256)):
|
|
with self.assertRaisesRegex(FileStorageError, "output_limit"):
|
|
inspect_archive(_zip_bytes({f"file-{index}": b"x" for index in range(20)}), filename="fixture.zip")
|
|
|
|
def test_destination_failure_kills_waiting_child_and_preserves_error(self):
|
|
self.store.side_effect = FileStorageError("Destination failure")
|
|
with self.assertRaisesRegex(FileStorageError, "Destination failure"):
|
|
_extract(_zip_bytes({"one": b"one", "two": b"two"}))
|
|
self.assertEqual(1, self.store.call_count)
|
|
self.assertEqual(1, len(self.children))
|
|
|
|
def test_missing_ack_times_out_and_wrong_ack_fails_in_real_child(self):
|
|
real_send = workers._send_ack
|
|
for send, error in ((lambda descriptor, sequence: None, "timeout"),
|
|
(lambda descriptor, sequence: real_send(descriptor, sequence + 1), "invalid staging record")):
|
|
with self.subTest(error=error), patch.object(workers, "_send_ack", side_effect=send):
|
|
with patch.object(workers, "EXTRACTION_LIMITS", replace(workers.EXTRACTION_LIMITS, wall_seconds=2)):
|
|
with self.assertRaisesRegex(FileStorageError, error):
|
|
_extract(_zip_bytes({"one": b"one", "two": b"two"}))
|
|
|
|
def test_parent_progress_exception_is_not_reclassified_as_transport_error(self):
|
|
def progress(stage, *counters):
|
|
if stage == "extracting":
|
|
raise ValueError("parent progress failure")
|
|
|
|
with self.assertRaisesRegex(ValueError, "parent progress failure"):
|
|
_extract(_zip_bytes({"one": b"one"}), progress=progress)
|
|
self.store.assert_not_called()
|
|
|
|
def test_source_path_snapshot_rejects_growth_and_shrink_without_spawning(self):
|
|
# A source path is server-owned; it must still not turn into an unbounded
|
|
# copy when another writer changes it during snapshot creation.
|
|
for replacement in (b"fixture plus unexpected growth", b"x"):
|
|
with self.subTest(replacement=replacement), tempfile.TemporaryDirectory() as temporary:
|
|
source = Path(temporary) / "archive"
|
|
source.write_bytes(b"fixture")
|
|
real_fstat = os.fstat
|
|
changed = False
|
|
|
|
def fstat(descriptor):
|
|
nonlocal changed
|
|
result = real_fstat(descriptor)
|
|
if not changed:
|
|
changed = True
|
|
source.write_bytes(replacement)
|
|
return result
|
|
|
|
with patch.object(workers.os, "fstat", side_effect=fstat):
|
|
with self.assertRaisesRegex(FileStorageError, "source changed"):
|
|
inspect_archive(source, filename="fixture.zip")
|
|
self.assertEqual([], self.children)
|
|
|
|
def test_busy_rejects_before_source_snapshot(self):
|
|
with patch.object(settings, "isolated_process_concurrency", 1), bounded_process.bounded_operation_admission():
|
|
with patch.object(workers, "_source_stage", side_effect=AssertionError("must admit before copying")):
|
|
with self.assertRaisesRegex(FileStorageError, "busy"):
|
|
inspect_archive(b"fixture", filename="fixture.zip")
|
|
self.assertEqual([], self.children)
|
|
self.assertEqual([], self.directories)
|
|
|
|
def test_unsupported_format_rejects_before_snapshot(self):
|
|
with self.assertRaisesRegex(FileStorageError, "Unsupported archive format"):
|
|
inspect_archive(b"fixture", filename="fixture.exe")
|
|
self.assertEqual([], self.directories)
|
|
|
|
def test_symlinked_staged_member_is_never_read_or_stored(self):
|
|
real_read = workers._read_regular
|
|
attacked = False
|
|
|
|
def read(path, maximum, **options):
|
|
nonlocal attacked
|
|
if path.name.startswith("member-") and not attacked:
|
|
attacked = True
|
|
path.unlink()
|
|
path.symlink_to(path.parent / "source")
|
|
return real_read(path, maximum, **options)
|
|
|
|
with patch.object(workers, "_read_regular", side_effect=read):
|
|
with self.assertRaisesRegex(FileStorageError, "invalid staging record"):
|
|
_extract(_zip_bytes({"one": b"one"}))
|
|
self.assertTrue(attacked)
|
|
self.store.assert_not_called()
|
|
|
|
def test_forged_member_records_reject_before_store(self):
|
|
real_read = workers._read_regular
|
|
for mutation in (
|
|
{"path": "outside/one"}, {"path": "../escape"}, {"size": -1},
|
|
{"sha256": "incorrect"}, {"sequence": True}, {"unexpected": True},
|
|
{"progress": ["extracting", 1, 1, 4, 3]},
|
|
):
|
|
with self.subTest(mutation=mutation):
|
|
self.store.reset_mock()
|
|
|
|
def read(path, maximum, **options):
|
|
value = real_read(path, maximum, **options)
|
|
if path.name == "status":
|
|
record = json.loads(value)
|
|
if record.get("kind") == "member":
|
|
record.update(mutation)
|
|
return json.dumps(record).encode()
|
|
return value
|
|
|
|
with patch.object(workers, "_read_regular", side_effect=read):
|
|
with self.assertRaises(FileStorageError):
|
|
_extract(_zip_bytes({"selected/one": b"one"}), selected_paths=("selected",))
|
|
self.store.assert_not_called()
|
|
|
|
def test_regressing_sequence_or_changing_totals_fail_before_second_store(self):
|
|
real_read = workers._read_regular
|
|
for mutation in ({"sequence": 1}, {"progress": ["extracting", 2, 3, 6, 6]}):
|
|
with self.subTest(mutation=mutation):
|
|
self.store.reset_mock()
|
|
|
|
def read(path, maximum, **options):
|
|
value = real_read(path, maximum, **options)
|
|
if path.name == "status":
|
|
record = json.loads(value)
|
|
if record.get("kind") == "member" and record["progress"][1] == 2:
|
|
record.update(mutation)
|
|
return json.dumps(record).encode()
|
|
return value
|
|
|
|
with patch.object(workers, "_read_regular", side_effect=read):
|
|
with self.assertRaisesRegex(FileStorageError, "invalid staging record"):
|
|
_extract(_zip_bytes({"one": b"one", "two": b"two"}))
|
|
self.assertEqual(1, self.store.call_count)
|
|
|
|
def test_record_reads_are_bounded_and_reject_non_regular_files(self):
|
|
with tempfile.TemporaryDirectory() as temporary:
|
|
directory = Path(temporary)
|
|
status = directory / "status"
|
|
status.write_bytes(b"x" * (workers._STATUS_BYTES + 1))
|
|
with self.assertRaises(FileStorageError):
|
|
workers._read_regular(status, workers._STATUS_BYTES)
|
|
with self.assertRaises(FileStorageError):
|
|
workers._read_regular(directory, workers._STATUS_BYTES)
|
|
|
|
def test_tiny_member_read_does_not_allocate_the_configured_gibibyte_cap(self):
|
|
with tempfile.TemporaryDirectory() as temporary:
|
|
member = Path(temporary) / "member"
|
|
member.write_bytes(b"x")
|
|
tracemalloc.start()
|
|
try:
|
|
self.assertEqual(b"x", workers._read_regular(member, 2 * 1024 * 1024 * 1024))
|
|
_current, peak = tracemalloc.get_traced_memory()
|
|
finally:
|
|
tracemalloc.stop()
|
|
self.assertLess(peak, 1024 * 1024)
|
|
|
|
def test_atomic_status_replacement_keeps_open_snapshot_valid_but_member_identity_is_strict(self):
|
|
for atomic_record in (True, False):
|
|
with self.subTest(atomic_record=atomic_record), tempfile.TemporaryDirectory() as temporary:
|
|
status = Path(temporary) / "status"
|
|
status.write_bytes(b"old")
|
|
replacement = Path(temporary) / "replacement"
|
|
replacement.write_bytes(b"new")
|
|
real_fstat = os.fstat
|
|
replaced = False
|
|
|
|
def fstat(descriptor):
|
|
nonlocal replaced
|
|
info = real_fstat(descriptor)
|
|
if not replaced:
|
|
replaced = True
|
|
os.replace(replacement, status)
|
|
return info
|
|
|
|
with patch.object(workers.os, "fstat", side_effect=fstat):
|
|
if atomic_record:
|
|
self.assertEqual(b"old", workers._read_regular(status, 16, atomic_record=True))
|
|
else:
|
|
with self.assertRaises(FileStorageError):
|
|
workers._read_regular(status, 16)
|
|
self.assertEqual(b"new", status.read_bytes())
|
|
|
|
def test_private_ack_channel_is_bounded_and_validates_type_permissions_and_sequence(self):
|
|
with tempfile.TemporaryDirectory() as temporary:
|
|
directory = Path(temporary)
|
|
with workers._ack_channel(directory, create=True) as writer:
|
|
with workers._ack_channel(directory) as reader:
|
|
workers._send_ack(writer, 7)
|
|
workers._receive_ack(reader, 7)
|
|
workers._send_ack(writer, 8)
|
|
with self.assertRaises(FileStorageError):
|
|
workers._receive_ack(reader, 9)
|
|
with patch.object(workers.os, "write", side_effect=BlockingIOError()):
|
|
with self.assertRaisesRegex(FileStorageError, "staging channel is unavailable"):
|
|
workers._send_ack(writer, 10)
|
|
ack = directory / "ack"
|
|
os.chmod(ack, 0o644)
|
|
with self.assertRaises(FileStorageError), workers._ack_channel(directory):
|
|
pass
|
|
ack.unlink()
|
|
ack.write_bytes(b"not a FIFO")
|
|
with self.assertRaises(FileStorageError), workers._ack_channel(directory):
|
|
pass
|
|
ack.unlink()
|
|
ack.symlink_to(directory / "missing")
|
|
with self.assertRaises(FileStorageError), workers._ack_channel(directory):
|
|
pass
|
|
|
|
def test_tiny_member_import_has_no_per_member_sleep_or_polling_delay(self):
|
|
self.assertNotIn("sleep(", inspect.getsource(workers._extract_worker))
|
|
for count in (100, 1000):
|
|
with self.subTest(count=count):
|
|
payload = _zip_bytes({f"member-{index}.txt": b"x" for index in range(count)})
|
|
started = time.monotonic()
|
|
self.assertEqual(count, len(_extract(payload)))
|
|
elapsed = time.monotonic() - started
|
|
# Broad synthetic regression ceiling, not a storage-throughput
|
|
# promise: 50 ms per member would take >=50 s for 1,000 files.
|
|
self.assertLess(elapsed, 15)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|