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()