sashatrask
91 lines
3.4 KiB
Python
91 lines
3.4 KiB
Python
import ctypes
|
|
import os
|
|
from pathlib import Path
|
|
import sys
|
|
import tempfile
|
|
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 process_identity
|
|
import runtime_security
|
|
|
|
|
|
@unittest.skipUnless(os.name == 'nt', 'Windows native binding stability')
|
|
class NativeBindingStabilityTests(unittest.TestCase):
|
|
def test_hot_native_paths_are_warm_stable(self):
|
|
from ctypes import wintypes
|
|
|
|
kernel32 = ctypes.WinDLL('kernel32', use_last_error=True)
|
|
pointer_dword = ctypes.POINTER(wintypes.DWORD)
|
|
get_handle_count = kernel32.GetProcessHandleCount
|
|
get_handle_count.argtypes = [wintypes.HANDLE, pointer_dword]
|
|
get_handle_count.restype = wintypes.BOOL
|
|
get_current_process = kernel32.GetCurrentProcess
|
|
get_current_process.argtypes = []
|
|
get_current_process.restype = wintypes.HANDLE
|
|
|
|
def handle_count():
|
|
value = wintypes.DWORD()
|
|
self.assertTrue(get_handle_count(get_current_process(), ctypes.byref(value)))
|
|
return int(value.value)
|
|
|
|
type_ids = (
|
|
id(runtime_security._OVERLAPPED),
|
|
id(runtime_security._SECURITY_ATTRIBUTES),
|
|
id(runtime_security._TOKEN_USER),
|
|
id(process_identity._FILETIME),
|
|
id(process_identity._UNICODE_STRING),
|
|
id(console_runner._PROCESS_MEMORY_COUNTERS_EX),
|
|
id(owned_process._JOBOBJECT_EXTENDED_LIMIT_INFORMATION),
|
|
)
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
runtime_security.ensure_private_directory(temp_dir, reject_reparse=True)
|
|
lock_path = os.path.join(temp_dir, 'warm.lock')
|
|
|
|
def exercise():
|
|
lock = runtime_security.PrivateFileLock(lock_path).acquire()
|
|
lock.release()
|
|
self.assertTrue(runtime_security.private_file_ready(lock_path))
|
|
process_identity.current_process_identity()
|
|
console_runner.current_process_private_bytes()
|
|
owned_process._windows_job_limit_information(1024 * 1024)
|
|
|
|
for _ in range(2000):
|
|
exercise()
|
|
memory_before = console_runner.current_process_private_bytes()
|
|
handles_before = handle_count()
|
|
with mock.patch.object(
|
|
ctypes, 'POINTER', side_effect=AssertionError('hot ctypes.POINTER call'),
|
|
), mock.patch.object(
|
|
ctypes, 'WinDLL', side_effect=AssertionError('hot ctypes.WinDLL call'),
|
|
):
|
|
for _ in range(20000):
|
|
exercise()
|
|
memory_after = console_runner.current_process_private_bytes()
|
|
handles_after = handle_count()
|
|
|
|
self.assertEqual(type_ids, (
|
|
id(runtime_security._OVERLAPPED),
|
|
id(runtime_security._SECURITY_ATTRIBUTES),
|
|
id(runtime_security._TOKEN_USER),
|
|
id(process_identity._FILETIME),
|
|
id(process_identity._UNICODE_STRING),
|
|
id(console_runner._PROCESS_MEMORY_COUNTERS_EX),
|
|
id(owned_process._JOBOBJECT_EXTENDED_LIMIT_INFORMATION),
|
|
))
|
|
growth = max(0, memory_after - memory_before)
|
|
self.assertLessEqual(growth, 16 * 1024 * 1024)
|
|
self.assertLessEqual(growth / 20000, 256)
|
|
self.assertLessEqual(abs(handles_after - handles_before), 8)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|