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