sashatrask
261 lines
12 KiB
Python
261 lines
12 KiB
Python
import os
|
|
from pathlib import Path
|
|
import sqlite3
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
from types import SimpleNamespace
|
|
import unittest
|
|
from unittest import mock
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
APP_DIR = ROOT / 'app'
|
|
sys.path.insert(0, str(APP_DIR))
|
|
|
|
import console_runner
|
|
import owned_process
|
|
import scanner
|
|
|
|
|
|
class ScanSlotCleanupTests(unittest.TestCase):
|
|
def test_fresh_dead_slot_is_removed_but_live_and_unknown_slots_remain(self):
|
|
now = time.time()
|
|
conn = sqlite3.connect(':memory:')
|
|
self.addCleanup(conn.close)
|
|
conn.executescript(scanner.SCAN_SLOT_SCHEMA)
|
|
|
|
for slot_id, owner_pid, child_pid in (
|
|
('dead', 101, 102),
|
|
('live', 201, 202),
|
|
('unknown', 301, 302),
|
|
):
|
|
conn.execute(
|
|
'''INSERT INTO scan_slots(
|
|
slot_id, owner_pid, owner_thread, owner_source, owner_creation_time,
|
|
owner_executable, child_pid, child_creation_time, child_executable,
|
|
acquired_at, updated_at
|
|
) VALUES (?, ?, 1, 'test', 'owner-created', 'owner.exe', ?,
|
|
'child-created', 'child.exe', ?, ?)''',
|
|
(slot_id, owner_pid, child_pid, now - 1, now - 1),
|
|
)
|
|
|
|
identity_state = {
|
|
101: False,
|
|
102: False,
|
|
201: True,
|
|
202: False,
|
|
301: False,
|
|
302: None,
|
|
}
|
|
with mock.patch.object(
|
|
scanner,
|
|
'exact_process_identity_live',
|
|
side_effect=lambda pid, creation, executable: identity_state[pid],
|
|
):
|
|
scanner.cleanup_stale_scan_slots(conn, now, stale_sec=60)
|
|
|
|
remaining = {row[0] for row in conn.execute('SELECT slot_id FROM scan_slots')}
|
|
self.assertEqual(remaining, {'live', 'unknown'})
|
|
|
|
|
|
class WindowsJobMemoryConfigTests(unittest.TestCase):
|
|
def test_job_limit_structure_combines_memory_and_kill_on_close_flags(self):
|
|
limit = 4096 * 1024 * 1024
|
|
limits = owned_process._windows_job_limit_information(limit)
|
|
flags = limits.BasicLimitInformation.LimitFlags
|
|
|
|
self.assertEqual(limits.JobMemoryLimit, limit)
|
|
self.assertEqual(
|
|
flags,
|
|
owned_process._JOB_OBJECT_LIMIT_JOB_MEMORY
|
|
| owned_process._JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE,
|
|
)
|
|
|
|
unlimited = owned_process._windows_job_limit_information(0)
|
|
self.assertEqual(
|
|
unlimited.BasicLimitInformation.LimitFlags,
|
|
owned_process._JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE,
|
|
)
|
|
self.assertEqual(unlimited.JobMemoryLimit, 0)
|
|
|
|
def test_owned_process_rejects_invalid_job_memory_limits_before_launch(self):
|
|
with mock.patch.object(owned_process.subprocess, 'Popen') as popen:
|
|
for value in (-1, None, True, 1.5, '4096'):
|
|
with self.subTest(value=value):
|
|
with self.assertRaises(ValueError):
|
|
owned_process.OwnedProcess(['fixture'], job_memory_limit_bytes=value)
|
|
popen.assert_not_called()
|
|
|
|
def test_owned_process_rejects_invalid_resource_policy_before_launch(self):
|
|
with mock.patch.object(owned_process.subprocess, 'Popen') as popen:
|
|
for name, values in (
|
|
('job_cpu_weight', (-1, 10, True, '2')),
|
|
('process_memory_priority', (-1, 6, True, '4')),
|
|
):
|
|
for value in values:
|
|
with self.subTest(name=name, value=value):
|
|
with self.assertRaises(ValueError):
|
|
owned_process.OwnedProcess(['fixture'], **{name: value})
|
|
popen.assert_not_called()
|
|
|
|
def test_canonical_job_memory_limit_is_applied(self):
|
|
limit = 4096 * 1024 * 1024
|
|
config = console_runner.load_config(str(APP_DIR / 'config.yaml'))
|
|
self.assertEqual(config['global']['trufflehog_job_memory_limit_bytes'], limit)
|
|
self.assertEqual(config['global']['trufflehog_windows_job_cpu_weight'], 2)
|
|
self.assertEqual(config['global']['trufflehog_windows_memory_priority'], 4)
|
|
docker_args = console_runner.build_args_from_source_config(
|
|
'dockerhub', config['sources']['dockerhub'], config['global'], 'fixture',
|
|
)
|
|
self.assertEqual(
|
|
docker_args.trufflehog_job_memory_limit_bytes,
|
|
6 * 1024 * 1024 * 1024,
|
|
)
|
|
|
|
previous = scanner.scan_config.__dict__.copy()
|
|
try:
|
|
console_runner.apply_global_config({
|
|
'trufflehog_job_memory_limit_bytes': limit,
|
|
'trufflehog_windows_job_cpu_weight': 2,
|
|
'trufflehog_windows_memory_priority': 4,
|
|
})
|
|
self.assertEqual(scanner.scan_config.trufflehog_job_memory_limit_bytes, limit)
|
|
self.assertEqual(scanner.scan_config.trufflehog_windows_job_cpu_weight, 2)
|
|
self.assertEqual(scanner.scan_config.trufflehog_windows_memory_priority, 4)
|
|
finally:
|
|
scanner.scan_config.__dict__.clear()
|
|
scanner.scan_config.__dict__.update(previous)
|
|
|
|
|
|
class RunCommandDiskPressureTests(unittest.TestCase):
|
|
def test_disk_pressure_kills_owned_process_and_reports_source_resource(self):
|
|
class MockProcess:
|
|
pid = 4321
|
|
job_membership_verified = True
|
|
|
|
def __init__(self):
|
|
self.returncode = None
|
|
self.killed = False
|
|
|
|
def poll(self):
|
|
return self.returncode
|
|
|
|
def kill(self):
|
|
self.killed = True
|
|
self.returncode = 1
|
|
|
|
def wait(self, timeout=None):
|
|
return self.returncode
|
|
|
|
process = MockProcess()
|
|
memory_limit = 4096 * 1024 * 1024
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
command_work_dir = os.path.join(temp_dir, 'trufflehog-run-test')
|
|
os.mkdir(command_work_dir)
|
|
with mock.patch.object(scanner, 'acquire_scan_slot', return_value=None), \
|
|
mock.patch.object(scanner, 'create_command_work_dir', return_value=command_work_dir), \
|
|
mock.patch.object(scanner, '_shared_staging_owners', return_value=[]), \
|
|
mock.patch.object(scanner, 'cleanup_command_work_dir'), \
|
|
mock.patch.object(scanner, 'require_trufflehog_launch_authority'), \
|
|
mock.patch.object(scanner, 'write_temp_owner'), \
|
|
mock.patch.object(scanner, 'OwnedProcess', return_value=process) as process_factory, \
|
|
mock.patch.object(scanner.shutil, 'disk_usage', return_value=SimpleNamespace(free=1)), \
|
|
mock.patch.object(scanner.scan_config, 'min_free_gb', 1), \
|
|
mock.patch.object(scanner.scan_config, 'trufflehog_job_memory_limit_bytes', memory_limit), \
|
|
mock.patch.object(scanner.scan_config, 'trufflehog_windows_job_cpu_weight', 2), \
|
|
mock.patch.object(scanner.scan_config, 'trufflehog_windows_memory_priority', 4):
|
|
stdout, stderr, returncode = scanner.run_command(['fixture'], timeout_sec=30)
|
|
|
|
process_factory.assert_called_once()
|
|
self.assertTrue(process.killed)
|
|
self.assertEqual(stdout, '')
|
|
self.assertEqual(returncode, -1)
|
|
self.assertIn('Not enough free space on configured TruffleHog work volume', stderr)
|
|
if os.name == 'nt':
|
|
self.assertEqual(process_factory.call_args.kwargs['job_memory_limit_bytes'], memory_limit)
|
|
self.assertEqual(process_factory.call_args.kwargs['job_cpu_weight'], 2)
|
|
self.assertEqual(process_factory.call_args.kwargs['process_memory_priority'], 4)
|
|
creationflags = process_factory.call_args.kwargs['creationflags']
|
|
self.assertTrue(creationflags & subprocess.CREATE_NEW_PROCESS_GROUP)
|
|
self.assertTrue(creationflags & subprocess.CREATE_NO_WINDOW)
|
|
self.assertFalse(
|
|
creationflags & getattr(subprocess, 'BELOW_NORMAL_PRIORITY_CLASS', 0x00004000)
|
|
)
|
|
|
|
result = scanner.apply_trufflehog_diagnostics({}, stderr, returncode, 'git')
|
|
self.assertEqual(result['error_class'], 'source_resource')
|
|
self.assertEqual(result['source_failure_category'], 'source_resource')
|
|
|
|
def test_opportunistic_capacity_gate_reserves_job_and_overhead(self):
|
|
gib = 1024 ** 3
|
|
previous = scanner.scan_config.__dict__.copy()
|
|
try:
|
|
scanner.scan_config.opportunistic_scan_slots = 1
|
|
scanner.scan_config.opportunistic_scan_sources = ['github']
|
|
scanner.scan_config.trufflehog_job_memory_limit_bytes = 4 * gib
|
|
scanner.scan_config.opportunistic_scan_reserve_overhead_bytes = 1 * gib
|
|
scanner.scan_config.opportunistic_scan_min_available_after_reserve_bytes = 4 * gib
|
|
scanner.scan_config.opportunistic_scan_min_commit_after_reserve_bytes = 6 * gib
|
|
with mock.patch.object(scanner, 'windows_scan_capacity_snapshot', return_value={
|
|
'available_physical_bytes': 9 * gib,
|
|
'commit_headroom_bytes': 11 * gib,
|
|
}):
|
|
self.assertTrue(scanner.opportunistic_scan_slot_allowed('github'))
|
|
with mock.patch.object(scanner, 'windows_scan_capacity_snapshot', return_value={
|
|
'available_physical_bytes': 9 * gib - 1,
|
|
'commit_headroom_bytes': 11 * gib,
|
|
}):
|
|
self.assertFalse(scanner.opportunistic_scan_slot_allowed('github'))
|
|
with mock.patch.object(scanner, 'windows_scan_capacity_snapshot', side_effect=OSError('probe')):
|
|
self.assertFalse(scanner.opportunistic_scan_slot_allowed('github'))
|
|
self.assertFalse(scanner.opportunistic_scan_slot_allowed('dockerhub'))
|
|
finally:
|
|
scanner.scan_config.__dict__.clear()
|
|
scanner.scan_config.__dict__.update(previous)
|
|
|
|
def test_disk_monitor_failure_kills_owned_process_fail_closed(self):
|
|
class MockProcess:
|
|
pid = 4322
|
|
job_membership_verified = True
|
|
|
|
def __init__(self):
|
|
self.returncode = None
|
|
self.killed = False
|
|
|
|
def poll(self):
|
|
return self.returncode
|
|
|
|
def kill(self):
|
|
self.killed = True
|
|
self.returncode = 1
|
|
|
|
def wait(self, timeout=None):
|
|
return self.returncode
|
|
|
|
process = MockProcess()
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
command_work_dir = os.path.join(temp_dir, 'trufflehog-run-test')
|
|
os.mkdir(command_work_dir)
|
|
with mock.patch.object(scanner, 'acquire_scan_slot', return_value=None), \
|
|
mock.patch.object(scanner, 'create_command_work_dir', return_value=command_work_dir), \
|
|
mock.patch.object(scanner, '_shared_staging_owners', return_value=[]), \
|
|
mock.patch.object(scanner, 'cleanup_command_work_dir'), \
|
|
mock.patch.object(scanner, 'require_trufflehog_launch_authority'), \
|
|
mock.patch.object(scanner, 'write_temp_owner'), \
|
|
mock.patch.object(scanner, 'OwnedProcess', return_value=process) as process_factory, \
|
|
mock.patch.object(scanner.shutil, 'disk_usage', side_effect=OSError('probe failed')), \
|
|
mock.patch.object(scanner.scan_config, 'min_free_gb', 1), \
|
|
mock.patch.object(scanner.scan_config, 'trufflehog_job_memory_limit_bytes', 4096 * 1024 * 1024):
|
|
_, stderr, returncode = scanner.run_command(['fixture'], timeout_sec=30)
|
|
|
|
process_factory.assert_called_once()
|
|
self.assertTrue(process.killed)
|
|
self.assertEqual(returncode, -1)
|
|
self.assertIn('Unable to monitor configured TruffleHog work volume', stderr)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|