sashatrask
397 lines
18 KiB
Python
397 lines
18 KiB
Python
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()
|