import multiprocessing import os from pathlib import Path from types import SimpleNamespace import sys import tempfile import threading import unittest import uuid from unittest import mock ROOT = Path(__file__).resolve().parents[1] APP_DIR = ROOT / 'app' sys.path.insert(0, str(APP_DIR)) import result_spool from result_spool import ( ResultSpool, SpoolBlockedError, SpoolCapacityError, SpoolCorruptionError, SpoolHashConflictError, ) from runtime_security import private_directory_ready, private_file_ready def event(event_id=None, payload='value'): return { 'version': 1, 'scan_event_id': event_id or str(uuid.uuid4()), 'payload': payload, } def quarantine_offer_worker(directory, start, outcomes): try: spool = ResultSpool( directory, max_event_bytes=512, max_events=2, max_total_bytes=4096, min_free_bytes=0, ) start.wait(10) invalid = event() invalid['version'] = 99 spool.write_event(invalid) outcomes.put('unexpected-success') except Exception as exc: outcomes.put(type(exc).__name__) def reservation_worker(directory, start, outcomes): try: spool = ResultSpool( directory, max_event_bytes=1024, max_events=3, max_total_bytes=4096, min_free_bytes=0, ) start.wait(10) outcomes.put(('reserved', spool.reserve_claims(str(os.getpid()), 1, 60))) except Exception as exc: outcomes.put((type(exc).__name__, '')) class ResultSpoolTests(unittest.TestCase): def make_spool(self, directory, **kwargs): return ResultSpool(directory, min_free_bytes=kwargs.pop('min_free_bytes', 0), **kwargs) def test_write_fsyncs_and_durably_renames_private_file(self): with tempfile.TemporaryDirectory() as temp_dir: spool = self.make_spool(os.path.join(temp_dir, 'spool')) with mock.patch.object(result_spool.os, 'fsync', wraps=os.fsync) as fsync, \ mock.patch.object(result_spool, 'durable_replace', wraps=result_spool.durable_replace) as replace: record = spool.write_event(event()) self.assertTrue(fsync.called) self.assertTrue(replace.called) source, destination = replace.call_args.args self.assertEqual(os.path.dirname(source), os.path.dirname(destination)) self.assertTrue(private_directory_ready(spool.directory)) self.assertTrue(private_directory_ready(spool.quarantine_directory)) self.assertTrue(private_file_ready(record.path)) def test_event_count_total_bytes_per_event_and_free_space_bounds(self): with tempfile.TemporaryDirectory() as temp_dir: count_spool = self.make_spool(os.path.join(temp_dir, 'count'), max_events=1) count_spool.write_event(event()) with self.assertRaises(SpoolCapacityError): count_spool.write_event(event()) size_spool = self.make_spool( os.path.join(temp_dir, 'size'), max_event_bytes=300, max_total_bytes=300, ) with self.assertRaises(SpoolCapacityError): size_spool.write_event(event(payload='x' * 500)) total_spool = self.make_spool( os.path.join(temp_dir, 'total'), max_event_bytes=700, max_total_bytes=700, ) total_spool.write_event(event(payload='x' * 250)) with self.assertRaises(SpoolCapacityError): total_spool.write_event(event(payload='y' * 250)) free_spool = self.make_spool(os.path.join(temp_dir, 'free'), min_free_bytes=100) with mock.patch.object(result_spool.shutil, 'disk_usage', return_value=SimpleNamespace(free=50)): with self.assertRaises(SpoolCapacityError): free_spool.write_event(event()) def test_corruption_is_quarantined_and_blocks_claim_gate(self): with tempfile.TemporaryDirectory() as temp_dir: spool = self.make_spool(os.path.join(temp_dir, 'spool')) record = spool.write_event(event()) with open(record.path, 'wb') as handle: handle.write(b'{broken') handle.flush() os.fsync(handle.fileno()) with self.assertRaises(SpoolCorruptionError): spool.pending_events() self.assertTrue(list(os.scandir(spool.quarantine_directory))) with self.assertRaises(SpoolBlockedError): spool.assert_claims_allowed() def test_same_id_hash_conflict_is_visible_and_blocks(self): with tempfile.TemporaryDirectory() as temp_dir: spool = self.make_spool(os.path.join(temp_dir, 'spool')) event_id = str(uuid.uuid4()) spool.write_event(event(event_id, 'first')) with self.assertRaises(SpoolHashConflictError): spool.write_event(event(event_id, 'second')) self.assertTrue(list(os.scandir(spool.quarantine_directory))) with self.assertRaises(SpoolBlockedError): spool.assert_claims_allowed() def test_supplied_hash_conflict_is_quarantined(self): with tempfile.TemporaryDirectory() as temp_dir: spool = self.make_spool(os.path.join(temp_dir, 'spool')) conflicting = event() conflicting['scan_event_hash'] = '0' * 64 with self.assertRaises(SpoolHashConflictError): spool.write_event(conflicting) self.assertTrue(list(os.scandir(spool.quarantine_directory))) def test_valid_pending_event_still_blocks_claim_gate_until_drained(self): with tempfile.TemporaryDirectory() as temp_dir: spool = self.make_spool(os.path.join(temp_dir, 'spool')) spool.write_event(event()) with self.assertRaises(SpoolBlockedError): spool.assert_claims_allowed() def test_ack_requires_matching_hash_and_removes_only_confirmed_event(self): with tempfile.TemporaryDirectory() as temp_dir: spool = self.make_spool(os.path.join(temp_dir, 'spool')) record = spool.write_event(event()) with self.assertRaises(SpoolHashConflictError): spool.acknowledge(record.event_id, '0' * 64) self.assertTrue(list(os.scandir(spool.quarantine_directory))) def test_reservation_consumes_capacity_atomically_with_event_write(self): with tempfile.TemporaryDirectory() as temp_dir: spool = self.make_spool( os.path.join(temp_dir, 'spool'), max_event_bytes=1024, max_events=3, max_total_bytes=4096, ) reservation = spool.reserve_claims('owner', 1, 60) claim = {'id': 7, 'lease_token': 'lease-token'} spool.bind_claims(reservation, [claim]) with self.assertRaises(SpoolCapacityError): spool.reserve_claims('other', 1, 60) record = spool.write_event( event(), reservation_id=reservation, queue_id=claim['id'], lease_token=claim['lease_token'], ) self.assertTrue(os.path.exists(record.path)) self.assertFalse(os.path.exists(spool._reservation_path(reservation))) def test_exact_logical_limit_allows_reserve_bind_renew_and_write(self): with tempfile.TemporaryDirectory() as temp_dir: spool = self.make_spool( os.path.join(temp_dir, 'spool'), max_event_bytes=1024, max_events=2, max_total_bytes=4096, ) reservation = spool.reserve_claims('owner', 1, 60, bytes_per_event=1024) spool.bind_claims(reservation, [{ 'id': 7, 'lease_token': 'lease-token', 'claim_batch': reservation, 'lease_owner': 'owner', }]) path, bound_record = spool._reservation_record(reservation) spool.max_total_bytes = len(result_spool._serialized_event(bound_record)) + 1024 self.assertTrue(spool.renew_reservation(reservation, 60)) record = spool.write_event( event(), reservation_id=reservation, queue_id=7, lease_token='lease-token', ) self.assertTrue(os.path.exists(record.path)) self.assertFalse(os.path.exists(path)) def test_reservation_replacement_charges_temporary_coexistence_to_free_space(self): with tempfile.TemporaryDirectory() as temp_dir: spool = self.make_spool( os.path.join(temp_dir, 'spool'), max_event_bytes=1024, max_events=3, max_total_bytes=4096, ) reservation = spool.reserve_claims('owner', 1, 60, bytes_per_event=1024) spool.bind_claims(reservation, [{ 'id': 8, 'lease_token': 'lease-token', 'claim_batch': reservation, 'lease_owner': 'owner', }]) _, current = spool._reservation_record(reservation) renewed = dict(current) renewed['expires_at'] += 60 renewed['recover_after'] += 120 replacement_size = len(result_spool._serialized_event(renewed)) with mock.patch.object( result_spool.shutil, 'disk_usage', return_value=SimpleNamespace(free=1024 + replacement_size - 1), ): with self.assertRaises(SpoolCapacityError): spool.renew_reservation(reservation, 60) def test_renewal_extends_crash_recovery_deadline(self): with tempfile.TemporaryDirectory() as temp_dir: with mock.patch.object(result_spool.time, 'time', return_value=1000.0): spool = self.make_spool(os.path.join(temp_dir, 'spool')) reservation = spool.reserve_claims('owner', 1, 60) with mock.patch.object(result_spool.time, 'time', return_value=1050.0): self.assertTrue(spool.renew_reservation(reservation, 60)) with mock.patch.object(result_spool.time, 'time', return_value=1121.0): self.assertTrue(os.path.exists(spool._reservation_path(reservation))) self.assertEqual(len(spool._load_reservations(recover=True)), 1) with mock.patch.object(result_spool.time, 'time', return_value=1171.0): self.assertEqual(spool._load_reservations(recover=True), []) self.assertFalse(os.path.exists(spool._reservation_path(reservation))) def test_two_workers_near_quota_admit_only_one_reservation(self): with tempfile.TemporaryDirectory() as temp_dir: directory = os.path.join(temp_dir, 'spool') first = self.make_spool(directory, max_event_bytes=1024, max_events=3, max_total_bytes=4096) second = self.make_spool(directory, max_event_bytes=1024, max_events=3, max_total_bytes=4096) barrier = threading.Barrier(2) outcomes = [] lock = threading.Lock() def reserve(spool, owner): barrier.wait() try: value = spool.reserve_claims(owner, 1, 60) except SpoolCapacityError: value = None with lock: outcomes.append(value) threads = [ threading.Thread(target=reserve, args=(first, 'one')), threading.Thread(target=reserve, args=(second, 'two')), ] for thread in threads: thread.start() for thread in threads: thread.join(5) self.assertEqual(sum(value is not None for value in outcomes), 1) def test_two_processes_near_quota_admit_only_one_reservation(self): with tempfile.TemporaryDirectory() as temp_dir: directory = os.path.join(temp_dir, 'spool') self.make_spool( directory, max_event_bytes=1024, max_events=3, max_total_bytes=4096, ) context = multiprocessing.get_context('spawn') start = context.Event() outcomes = context.Queue() processes = [ context.Process(target=reservation_worker, args=(directory, start, outcomes)) for _ in range(2) ] for process in processes: process.start() start.set() for process in processes: process.join(15) self.assertFalse(process.is_alive()) self.assertEqual(process.exitcode, 0) kinds = sorted(outcomes.get(timeout=5)[0] for _ in processes) self.assertEqual(kinds, ['SpoolTransientCapacityError', 'reserved']) def test_fsync_failure_leaves_no_event_and_keeps_reservation(self): with tempfile.TemporaryDirectory() as temp_dir: spool = self.make_spool(os.path.join(temp_dir, 'spool'), max_event_bytes=1024) reservation = spool.reserve_claims('owner', 1, 60) claim = {'id': 9, 'lease_token': 'lease-token'} spool.bind_claims(reservation, [claim]) with mock.patch.object(result_spool.os, 'fsync', side_effect=OSError('disk fsync failed')): with self.assertRaisesRegex(OSError, 'fsync'): spool.write_event( event(), reservation_id=reservation, queue_id=claim['id'], lease_token=claim['lease_token'], ) self.assertTrue(os.path.exists(spool._reservation_path(reservation))) self.assertEqual(spool._pending_paths()[0], []) def test_oversized_invalid_offer_writes_only_bounded_diagnostic(self): with tempfile.TemporaryDirectory() as temp_dir: spool = self.make_spool( os.path.join(temp_dir, 'spool'), max_event_bytes=512, max_total_bytes=4096, ) secret = 'sensitive-value-' * 10000 with self.assertRaises(SpoolCapacityError): spool.write_event(event(payload=secret)) files = [entry.path for entry in os.scandir(spool.quarantine_directory)] self.assertTrue(files) self.assertTrue(all(os.path.getsize(path) <= result_spool.MAX_DIAGNOSTIC_BYTES for path in files)) combined = b''.join(Path(path).read_bytes() for path in files) self.assertNotIn(secret[:100].encode(), combined) def test_quarantine_pair_is_not_partially_written_when_tiny_quota_cannot_fit(self): with tempfile.TemporaryDirectory() as temp_dir: spool = self.make_spool( os.path.join(temp_dir, 'spool'), max_event_bytes=256, max_events=2, max_total_bytes=256, ) invalid = event() invalid['version'] = 99 with self.assertRaises(SpoolCapacityError): spool.write_event(invalid) self.assertEqual(list(os.scandir(spool.quarantine_directory)), []) def test_repeated_invalid_offers_stop_without_growing_quarantine(self): with tempfile.TemporaryDirectory() as temp_dir: spool = self.make_spool( os.path.join(temp_dir, 'spool'), max_event_bytes=512, max_events=2, max_total_bytes=4096, ) invalid = event() invalid['version'] = 99 with self.assertRaises(SpoolCorruptionError): spool.write_event(invalid) first = sorted(entry.name for entry in os.scandir(spool.quarantine_directory)) self.assertEqual(len(first), 2) for _ in range(5): with self.assertRaises(SpoolBlockedError): spool.write_event(invalid) self.assertEqual( sorted(entry.name for entry in os.scandir(spool.quarantine_directory)), first, ) def test_cross_process_quarantine_admission_creates_only_one_complete_pair(self): with tempfile.TemporaryDirectory() as temp_dir: directory = os.path.join(temp_dir, 'spool') self.make_spool( directory, max_event_bytes=512, max_events=2, max_total_bytes=4096, ) context = multiprocessing.get_context('spawn') start = context.Event() outcomes = context.Queue() processes = [ context.Process(target=quarantine_offer_worker, args=(directory, start, outcomes)) for _ in range(2) ] for process in processes: process.start() start.set() for process in processes: process.join(15) self.assertFalse(process.is_alive()) self.assertEqual(process.exitcode, 0) found = sorted(outcomes.get(timeout=5) for _ in processes) self.assertEqual(found, ['SpoolBlockedError', 'SpoolCorruptionError']) files = list(os.scandir(os.path.join(directory, result_spool.QUARANTINE_DIRNAME))) self.assertEqual(len(files), 2) def test_pending_event_reader_is_incremental(self): with tempfile.TemporaryDirectory() as temp_dir: spool = self.make_spool(os.path.join(temp_dir, 'spool')) spool.write_event(event()) spool.write_event(event()) self.assertEqual(len(spool.pending_events()), 1) def test_hostile_symlink_is_never_followed_or_removed(self): with tempfile.TemporaryDirectory() as temp_dir: spool = self.make_spool(os.path.join(temp_dir, 'spool')) external = os.path.join(temp_dir, 'external.txt') Path(external).write_text('external-secret', encoding='ascii') link = os.path.join(spool.directory, 'hostile.json') try: os.symlink(external, link) except (OSError, NotImplementedError): self.skipTest('symlink creation is unavailable on this platform') with self.assertRaises(SpoolCorruptionError): spool.assert_claims_allowed() self.assertEqual(Path(external).read_text(encoding='ascii'), 'external-secret') self.assertTrue(os.path.lexists(link)) if __name__ == '__main__': unittest.main()