Initial server source import
This commit is contained in:
sashatrask
@@ -0,0 +1,579 @@
|
||||
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()
|
||||
Reference in New Issue
Block a user