import concurrent.futures import os from pathlib import Path import sqlite3 import subprocess import sys import tempfile import threading import time import unittest from unittest import mock ROOT = Path(__file__).resolve().parents[1] APP_DIR = ROOT / 'app' sys.path.insert(0, str(APP_DIR)) import scanner import janitor class RecordingLease: def __init__(self): self.releasable = True self.children = [] self.release_calls = 0 self.non_release_calls = 0 def set_child_pid(self, pid): self.children.append(pid) return True def mark_non_releasable(self): self.non_release_calls += 1 self.releasable = False def release(self): self.release_calls += 1 class CompleteProcess: returncode = 0 job_membership_verified = True def __init__(self, pid): self.pid = pid self.payload_identity = { 'pid': pid, 'creation_time': f'fixture-{pid}', 'executable': sys.executable, } def poll(self): return 0 def private_work_dir(temp_dir): work_dir = os.path.join(temp_dir, 'work') scanner.ensure_private_directory(work_dir, reject_reparse=True) return work_dir class ArtifactLifecycleSlotTests(unittest.TestCase): @classmethod def setUpClass(cls): scanner.initialize_scanner_runtime(preflight_complete=True, register_cleanup=False) def tearDown(self): if hasattr(scanner._scan_slot_scope_local, 'scope'): del scanner._scan_slot_scope_local.scope with scanner._scan_slot_fatal_lock: scanner._scan_slot_fatal_detail = None scanner._scan_slot_fatal_event.clear() def test_existing_limiter_schema_gains_slot_kind_and_single_bonus_index(self): with tempfile.TemporaryDirectory() as temp_dir: limiter_db = os.path.join(temp_dir, 'scan-limiter.sqlite') connection = sqlite3.connect(limiter_db) try: connection.executescript(''' CREATE TABLE scan_slots ( slot_id TEXT PRIMARY KEY, owner_pid INTEGER NOT NULL, owner_thread INTEGER NOT NULL, owner_source TEXT, owner_creation_time TEXT, owner_executable TEXT, child_pid INTEGER, child_creation_time TEXT, child_executable TEXT, command TEXT, acquired_at REAL NOT NULL, updated_at REAL NOT NULL ); ''') connection.commit() finally: connection.close() scanner.ensure_scan_limiter_db(limiter_db) connection = sqlite3.connect(limiter_db) try: columns = { row[1] for row in connection.execute('PRAGMA table_info(scan_slots)') } indexes = { row[1] for row in connection.execute('PRAGMA index_list(scan_slots)') } finally: connection.close() self.assertIn('slot_kind', columns) self.assertIn('idx_scan_slots_single_bonus', indexes) def test_two_workers_cannot_enter_target_staging_concurrently(self): first_entered = threading.Event() second_entered = threading.Event() release_first = threading.Event() state_lock = threading.Lock() state = {'active': 0, 'maximum': 0, 'entries': 0} def stage_target(target, provider=None, **kwargs): with state_lock: state['active'] += 1 state['maximum'] = max(state['maximum'], state['active']) state['entries'] += 1 entry = state['entries'] try: if entry == 1: first_entered.set() if not release_first.wait(5): raise RuntimeError('test did not release first target') else: second_entered.set() return {'findings': [], 'errors': []} finally: with state_lock: state['active'] -= 1 with tempfile.TemporaryDirectory() as temp_dir: limiter_db = os.path.join(temp_dir, 'scan-limiter.sqlite') with mock.patch.object(scanner.scan_config, 'max_active_scans', 1), \ mock.patch.object(scanner.scan_config, 'scan_limiter_db', limiter_db), \ mock.patch.object(scanner.scan_config, 'scan_slot_wait_sec', 0.01), \ mock.patch.object(scanner, 'scan_git_repo', side_effect=stage_target): with concurrent.futures.ThreadPoolExecutor(max_workers=1) as caller: future = caller.submit( scanner.scan_targets_batch, ['first', 'second'], 'git', max_workers=2, persist_results=False, ) try: self.assertTrue(first_entered.wait(3)) self.assertFalse(second_entered.wait(0.3)) finally: release_first.set() results = future.result(timeout=5) self.assertTrue(second_entered.is_set()) self.assertEqual(state['maximum'], 1) self.assertEqual(len(results), 2) conn = sqlite3.connect(limiter_db) try: self.assertEqual(conn.execute('SELECT COUNT(*) FROM scan_slots').fetchone()[0], 0) finally: conn.close() def test_fair_waiters_rotate_to_the_source_least_recently_granted(self): with tempfile.TemporaryDirectory() as temp_dir, mock.patch.dict( os.environ, {'SCANNER_SOURCE': ''}, clear=False, ): limiter_db = os.path.join(temp_dir, 'scan-limiter.sqlite') order = [] order_lock = threading.Lock() acquired = {'a': threading.Event(), 'b': threading.Event()} release = {'a': threading.Event(), 'b': threading.Event()} def waiter(source): lease = scanner.acquire_scan_slot(['dispatch', source], 60) with order_lock: order.append(source) acquired[source].set() release[source].wait(5) lease.release() with mock.patch.object(scanner.scan_config, 'max_active_scans', 1), \ mock.patch.object(scanner.scan_config, 'scan_limiter_db', limiter_db), \ mock.patch.object(scanner.scan_config, 'scan_slot_wait_sec', 0.01): occupant = scanner.acquire_scan_slot(['dispatch', 'a'], 60) threads = [ threading.Thread(target=waiter, args=('a',)), threading.Thread(target=waiter, args=('b',)), ] for thread in threads: thread.start() time.sleep(0.1) occupant.release() self.assertTrue(acquired['b'].wait(3)) self.assertFalse(acquired['a'].is_set()) release['b'].set() self.assertTrue(acquired['a'].wait(3)) release['a'].set() for thread in threads: thread.join(5) self.assertEqual(order, ['b', 'a']) def test_three_base_slots_remain_guaranteed_with_one_gated_bonus(self): with tempfile.TemporaryDirectory() as temp_dir, mock.patch.dict( os.environ, {'SCANNER_SOURCE': ''}, clear=False, ): limiter_db = os.path.join(temp_dir, 'scan-limiter.sqlite') policy = ( mock.patch.object(scanner.scan_config, 'max_active_scans', 3), mock.patch.object(scanner.scan_config, 'opportunistic_scan_slots', 1), mock.patch.object(scanner.scan_config, 'opportunistic_scan_sources', ['github']), mock.patch.object(scanner.scan_config, 'scan_limiter_db', limiter_db), mock.patch.object(scanner.scan_config, 'scan_slot_wait_sec', 0.01), mock.patch.object(scanner, 'opportunistic_scan_slot_allowed', return_value=True), ) with policy[0], policy[1], policy[2], policy[3], policy[4], policy[5]: base = [ scanner.acquire_scan_slot(['dispatch', 'github'], 60, wait=False) for _ in range(3) ] bonus = scanner.acquire_scan_slot(['dispatch', 'github'], 60, wait=False) fifth = scanner.acquire_scan_slot(['dispatch', 'github'], 60, wait=False) self.assertTrue(all(base)) self.assertIsNotNone(bonus) self.assertIsNone(fifth) connection = sqlite3.connect(limiter_db) try: self.assertEqual( dict(connection.execute( 'SELECT slot_kind, COUNT(*) FROM scan_slots GROUP BY slot_kind' )), {'base': 3, 'bonus': 1}, ) finally: connection.close() base[0].release() replacement = scanner.acquire_scan_slot( ['dispatch', 'github'], 60, wait=False, ) self.assertIsNotNone(replacement) connection = sqlite3.connect(limiter_db) try: kinds = dict(connection.execute( 'SELECT slot_kind, COUNT(*) FROM scan_slots GROUP BY slot_kind' )) finally: connection.close() self.assertEqual(kinds, {'base': 3, 'bonus': 1}) for lease in base[1:] + [bonus, replacement]: lease.release() def test_bonus_slot_is_denied_to_docker_and_on_probe_failure(self): with tempfile.TemporaryDirectory() as temp_dir, mock.patch.dict( os.environ, {'SCANNER_SOURCE': ''}, clear=False, ): limiter_db = os.path.join(temp_dir, 'scan-limiter.sqlite') with mock.patch.object(scanner.scan_config, 'max_active_scans', 3), \ mock.patch.object(scanner.scan_config, 'opportunistic_scan_slots', 1), \ mock.patch.object(scanner.scan_config, 'opportunistic_scan_sources', ['github']), \ mock.patch.object(scanner.scan_config, 'scan_limiter_db', limiter_db), \ mock.patch.object(scanner.scan_config, 'scan_slot_wait_sec', 0.01), \ mock.patch.object(scanner, 'windows_scan_capacity_snapshot', side_effect=OSError('probe')): base = [ scanner.acquire_scan_slot(['dispatch', 'github'], 60, wait=False) for _ in range(3) ] self.assertIsNone(scanner.acquire_scan_slot( ['dispatch', 'dockerhub'], 60, wait=False, )) self.assertIsNone(scanner.acquire_scan_slot( ['dispatch', 'github'], 60, wait=False, )) for lease in base: lease.release() def test_preacquired_three_slots_release_before_blocked_downstream_work(self): stage_barrier = threading.Barrier(3) downstream_entered = threading.Event() release_downstream = threading.Event() state_lock = threading.Lock() state = {'active': 0, 'maximum': 0, 'sink_calls': 0} def stage_target(target, provider=None, **kwargs): with state_lock: state['active'] += 1 state['maximum'] = max(state['maximum'], state['active']) try: stage_barrier.wait(timeout=5) return {'findings': [], 'errors': []} finally: with state_lock: state['active'] -= 1 def sink(_result): with state_lock: state['sink_calls'] += 1 def blocked_downstream(completed, total, target): if target not in ('Scan completed', 'Scan completed with persistence failures') and not downstream_entered.is_set(): downstream_entered.set() if not release_downstream.wait(5): raise RuntimeError('test did not release downstream work') with tempfile.TemporaryDirectory() as temp_dir: limiter_db = os.path.join(temp_dir, 'scan-limiter.sqlite') with mock.patch.object(scanner.scan_config, 'max_active_scans', 3), \ mock.patch.object(scanner.scan_config, 'scan_limiter_db', limiter_db), \ mock.patch.object(scanner.scan_config, 'scan_slot_wait_sec', 0.01), \ mock.patch.object(scanner, 'scan_git_repo', side_effect=stage_target): leases = scanner.acquire_scan_slot_leases(['dispatch', 'git'], 3, timeout_sec=60) self.assertEqual(len(leases), 3) with concurrent.futures.ThreadPoolExecutor(max_workers=1) as caller: future = caller.submit( scanner.scan_targets_batch, ['one', 'two', 'three'], 'git', max_workers=3, progress_callback=blocked_downstream, persist_results=False, result_sink=sink, scan_slot_leases=leases, sink_within_scan_slot=True, ) self.assertTrue(downstream_entered.wait(5)) deadline = time.monotonic() + 3 active = None while time.monotonic() < deadline: conn = sqlite3.connect(limiter_db) try: active = conn.execute('SELECT COUNT(*) FROM scan_slots').fetchone()[0] finally: conn.close() if active == 0: break time.sleep(0.02) self.assertEqual(active, 0) release_downstream.set() results = future.result(timeout=10) self.assertEqual(state['maximum'], 3) self.assertEqual(state['sink_calls'], 3) self.assertEqual(len(results), 3) conn = sqlite3.connect(limiter_db) try: self.assertEqual(conn.execute('SELECT COUNT(*) FROM scan_slots').fetchone()[0], 0) finally: conn.close() def test_durable_bundle_handoff_releases_all_permits_before_db_notification(self): with tempfile.TemporaryDirectory() as temp_dir: limiter_db = os.path.join(temp_dir, 'scan-limiter.sqlite') bundle_root = os.path.join(temp_dir, 'bundles') scanner.ensure_private_directory(bundle_root, reject_reparse=True) notification_lock = threading.Lock() notification_state = {'entered': 0} all_notifications_entered = threading.Event() release_notifications = threading.Event() def stage(index, lease): event_id = f'{index + 1:032x}' bundle_id = f'{index + 100:032x}' reservation = { 'reservation_id': index + 1, 'reservation_token': f'reservation-{index}', 'bundle_id': bundle_id, 'scan_event_id': event_id, 'queue_id': index + 1, 'claim_lease_token': f'lease-{index}', 'declared_bundle_bytes': 1024 * 1024, 'ready_relative_path': f'ready/{bundle_id[:2]}/{bundle_id}.trb', 'source': 'fixture', 'platform': 'git', 'target': f'target-{index}', } with scanner.scan_slot_scope(['scan-target', 'git'], 60, lease=lease): result = { 'scan_event_id': event_id, 'target': f'target-{index}', 'scan_type': 'git', 'timestamp': '2026-07-27T00:00:00+00:00', 'findings': [], 'errors': [], } scanner.stage_result_bundle( result, reservation, bundle_root, {}, {'queue_status': 'done'}, ) with notification_lock: notification_state['entered'] += 1 if notification_state['entered'] == 3: all_notifications_entered.set() release_notifications.wait(5) with mock.patch.object(scanner.scan_config, 'max_active_scans', 3), \ mock.patch.object(scanner.scan_config, 'scan_limiter_db', limiter_db), \ mock.patch.object(scanner.scan_config, 'scan_slot_wait_sec', 0.01): leases = scanner.acquire_scan_slot_leases(['dispatch', 'git'], 3, timeout_sec=60) with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor: futures = [executor.submit(stage, index, lease) for index, lease in enumerate(leases)] deadline = time.monotonic() + 5 active = None while time.monotonic() < deadline: conn = sqlite3.connect(limiter_db) try: active = conn.execute('SELECT COUNT(*) FROM scan_slots').fetchone()[0] finally: conn.close() if active == 0 and all_notifications_entered.is_set(): break time.sleep(0.02) self.assertEqual(active, 0) self.assertTrue(all_notifications_entered.is_set()) release_notifications.set() for future in futures: future.result(timeout=5) def test_nested_commands_borrow_one_outer_physical_lease(self): lease = RecordingLease() processes = [CompleteProcess(4101), CompleteProcess(4102)] with tempfile.TemporaryDirectory() as temp_dir: work_dir = private_work_dir(temp_dir) def scan_target(target, provider=None, **kwargs): first = scanner.run_command(['fixture', 'one'], timeout_sec=5) second = scanner.run_command(['fixture', 'two'], timeout_sec=5) self.assertEqual(first, ('', '', 0)) self.assertEqual(second, ('', '', 0)) return {'findings': [], 'errors': []} with mock.patch.object(scanner, 'acquire_scan_slot', return_value=lease) as acquire, \ mock.patch.object(scanner, 'scan_git_repo', side_effect=scan_target), \ mock.patch.object(scanner.scan_config, 'work_dir', work_dir), \ mock.patch.object(scanner, 'require_trufflehog_launch_authority'), \ mock.patch.object(scanner, 'OwnedProcess', side_effect=processes), \ mock.patch.object(scanner.scan_config, 'min_free_gb', 0), \ mock.patch.object(scanner.scan_config, 'trufflehog_job_memory_limit_bytes', 4 * 1024 * 1024 * 1024): results = scanner.scan_targets_batch( ['target'], 'git', max_workers=1, persist_results=False, ) self.assertEqual(results[0]['errors'], []) self.assertEqual(acquire.call_count, 1) self.assertEqual(lease.children, [4101, 4102]) self.assertEqual(lease.release_calls, 1) def test_scope_clears_after_exception_before_executor_thread_reuse(self): lease = RecordingLease() worker_threads = [] def fail_in_scope(): worker_threads.append(threading.get_ident()) with scanner.scan_slot_scope(['scan-target', 'fixture']): raise RuntimeError('fixture failure') def inspect_reused_thread(): worker_threads.append(threading.get_ident()) return scanner.scoped_scan_slot_lease() with mock.patch.object(scanner, 'acquire_scan_slot', return_value=lease) as acquire: with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor: with self.assertRaisesRegex(RuntimeError, 'fixture failure'): executor.submit(fail_in_scope).result(timeout=3) borrowed, current = executor.submit(inspect_reused_thread).result(timeout=3) self.assertEqual(worker_threads[0], worker_threads[1]) self.assertFalse(borrowed) self.assertIsNone(current) self.assertEqual(acquire.call_count, 1) self.assertEqual(lease.release_calls, 1) def test_standalone_command_acquires_and_releases_its_own_lease(self): lease = RecordingLease() with tempfile.TemporaryDirectory() as temp_dir: work_dir = private_work_dir(temp_dir) with mock.patch.object(scanner, 'acquire_scan_slot', return_value=lease) as acquire, \ mock.patch.object(scanner.scan_config, 'work_dir', work_dir), \ mock.patch.object(scanner, 'require_trufflehog_launch_authority'), \ mock.patch.object(scanner, 'OwnedProcess', return_value=CompleteProcess(4151)), \ mock.patch.object(scanner.scan_config, 'min_free_gb', 0), \ mock.patch.object(scanner.scan_config, 'trufflehog_job_memory_limit_bytes', 4 * 1024 * 1024 * 1024): result = scanner.run_command(['fixture'], timeout_sec=5) self.assertEqual(result, ('', '', 0)) self.assertEqual(acquire.call_count, 1) self.assertEqual(lease.children, [4151]) self.assertEqual(lease.release_calls, 1) def test_uncertain_nested_child_termination_prevents_outer_release(self): class UnstoppableProcess: pid = 4201 returncode = None job_membership_verified = True payload_identity = { 'pid': pid, 'creation_time': 'fixture-4201', 'executable': sys.executable, } def poll(self): return None def kill(self): return None def wait(self, timeout=None): raise subprocess.TimeoutExpired(['fixture'], timeout) lease = RecordingLease() with tempfile.TemporaryDirectory() as temp_dir: work_dir = private_work_dir(temp_dir) with mock.patch.object(scanner.scan_config, 'work_dir', work_dir), \ mock.patch.object(scanner.scan_config, 'min_free_gb', 0): command_work_dir = scanner.create_command_work_dir() with mock.patch.object(scanner, 'acquire_scan_slot', return_value=lease) as acquire, \ mock.patch.object(scanner.scan_config, 'work_dir', work_dir), \ mock.patch.object(scanner, 'create_command_work_dir', return_value=command_work_dir), \ mock.patch.object(scanner, 'require_trufflehog_launch_authority'), \ mock.patch.object(scanner, 'OwnedProcess', return_value=UnstoppableProcess()), \ mock.patch.object(scanner.shutil, 'disk_usage', return_value=mock.Mock(free=1)), \ mock.patch.object(scanner.scan_config, 'min_free_gb', 1), \ mock.patch.object(scanner.scan_config, 'trufflehog_job_memory_limit_bytes', 4 * 1024 * 1024 * 1024): with scanner.scan_slot_scope(['scan-target', 'fixture']): with self.assertRaisesRegex(scanner.ScanSlotFatalError, 'termination was not confirmed'): scanner.run_command(['fixture'], timeout_sec=5) self.assertEqual(acquire.call_count, 1) self.assertFalse(lease.releasable) self.assertGreaterEqual(lease.non_release_calls, 1) self.assertEqual(lease.release_calls, 0) self.assertEqual(scanner.scoped_scan_slot_lease(), (False, None)) def test_command_workdir_marker_is_schema_two_and_only_janitor_removes_stale_tree(self): with tempfile.TemporaryDirectory() as temp_dir: work_dir = os.path.join(temp_dir, 'work') scanner.ensure_private_directory(work_dir, reject_reparse=True) markerless = os.path.join(work_dir, 'trufflehog-run-foreign') scanner.ensure_private_directory(markerless, reject_reparse=True) with mock.patch.object(scanner.scan_config, 'work_dir', work_dir), \ mock.patch.object(scanner.scan_config, 'min_free_gb', 0): command_dir = scanner.create_command_work_dir() marker = os.path.join(command_dir, scanner.TEMP_OWNER_FILE) owner = scanner.read_temp_owner(command_dir) self.assertEqual(owner.get('owner_pid'), os.getpid()) self.assertEqual(owner.get('schema'), 2) self.assertTrue(owner.get('owner_creation_time')) self.assertTrue(owner.get('owner_executable')) self.assertTrue(scanner.private_file_ready(marker)) old = time.time() - 7200 os.utime(command_dir, (old, old)) os.utime(markerless, (old, old)) removed = scanner.cleanup_stale_temp_dirs(age_minutes=1, log=False) self.assertEqual(removed, 0) with mock.patch.object(janitor, 'exact_process_identity_state', return_value='dead'): report = janitor.run_janitor_pass( work_dir, [owner['owner_executable'], owner['parent_executable']], minimum_age_sec=0, ) self.assertEqual(report['removed'], 1) self.assertFalse(os.path.exists(command_dir)) self.assertTrue(os.path.isdir(markerless)) before = set(os.listdir(work_dir)) with mock.patch.object(scanner, 'atomic_write_private_json', side_effect=OSError('marker denied')): with self.assertRaisesRegex(RuntimeError, 'required temp owner marker'): scanner.create_command_work_dir() self.assertEqual(set(os.listdir(work_dir)), before) if __name__ == '__main__': unittest.main()