"""Windows Job Object ownership for the authenticated worker tree. The worker is created suspended, assigned to a per-task kill-on-close job, and only then resumed. There is intentionally no non-Windows fallback. """ from __future__ import annotations import ctypes import os import subprocess from ctypes import wintypes from dataclasses import dataclass from typing import Iterable, Sequence class JobError(RuntimeError): pass if os.name == "nt": kernel32 = ctypes.WinDLL("kernel32", use_last_error=True) CREATE_SUSPENDED = 0x00000004 CREATE_NO_WINDOW = 0x08000000 EXTENDED_STARTUPINFO_PRESENT = 0x00080000 STARTF_USESTDHANDLES = 0x00000100 HANDLE_FLAG_INHERIT = 0x00000001 PROC_THREAD_ATTRIBUTE_HANDLE_LIST = 0x00020002 JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE = 0x00002000 JOB_OBJECT_EXTENDED_LIMIT_INFORMATION = 9 WAIT_OBJECT_0 = 0 WAIT_TIMEOUT = 258 INFINITE = 0xFFFFFFFF GENERIC_READ = 0x80000000 GENERIC_WRITE = 0x40000000 OPEN_EXISTING = 3 FILE_ATTRIBUTE_NORMAL = 0x80 class SECURITY_ATTRIBUTES(ctypes.Structure): _fields_ = [ ("nLength", wintypes.DWORD), ("lpSecurityDescriptor", wintypes.LPVOID), ("bInheritHandle", wintypes.BOOL), ] class STARTUPINFOW(ctypes.Structure): _fields_ = [ ("cb", wintypes.DWORD), ("lpReserved", wintypes.LPWSTR), ("lpDesktop", wintypes.LPWSTR), ("lpTitle", wintypes.LPWSTR), ("dwX", wintypes.DWORD), ("dwY", wintypes.DWORD), ("dwXSize", wintypes.DWORD), ("dwYSize", wintypes.DWORD), ("dwXCountChars", wintypes.DWORD), ("dwYCountChars", wintypes.DWORD), ("dwFillAttribute", wintypes.DWORD), ("dwFlags", wintypes.DWORD), ("wShowWindow", wintypes.WORD), ("cbReserved2", wintypes.WORD), ("lpReserved2", ctypes.POINTER(ctypes.c_ubyte)), ("hStdInput", wintypes.HANDLE), ("hStdOutput", wintypes.HANDLE), ("hStdError", wintypes.HANDLE), ] class STARTUPINFOEXW(ctypes.Structure): _fields_ = [ ("StartupInfo", STARTUPINFOW), ("lpAttributeList", wintypes.LPVOID), ] class PROCESS_INFORMATION(ctypes.Structure): _fields_ = [ ("hProcess", wintypes.HANDLE), ("hThread", wintypes.HANDLE), ("dwProcessId", wintypes.DWORD), ("dwThreadId", wintypes.DWORD), ] class JOBOBJECT_BASIC_LIMIT_INFORMATION(ctypes.Structure): _fields_ = [ ("PerProcessUserTimeLimit", ctypes.c_int64), ("PerJobUserTimeLimit", ctypes.c_int64), ("LimitFlags", wintypes.DWORD), ("MinimumWorkingSetSize", ctypes.c_size_t), ("MaximumWorkingSetSize", ctypes.c_size_t), ("ActiveProcessLimit", wintypes.DWORD), ("Affinity", ctypes.c_size_t), ("PriorityClass", wintypes.DWORD), ("SchedulingClass", wintypes.DWORD), ] class IO_COUNTERS(ctypes.Structure): _fields_ = [ ("ReadOperationCount", ctypes.c_uint64), ("WriteOperationCount", ctypes.c_uint64), ("OtherOperationCount", ctypes.c_uint64), ("ReadTransferCount", ctypes.c_uint64), ("WriteTransferCount", ctypes.c_uint64), ("OtherTransferCount", ctypes.c_uint64), ] class JOBOBJECT_EXTENDED_LIMIT_INFORMATION_STRUCT(ctypes.Structure): _fields_ = [ ("BasicLimitInformation", JOBOBJECT_BASIC_LIMIT_INFORMATION), ("IoInfo", IO_COUNTERS), ("ProcessMemoryLimit", ctypes.c_size_t), ("JobMemoryLimit", ctypes.c_size_t), ("PeakProcessMemoryUsed", ctypes.c_size_t), ("PeakJobMemoryUsed", ctypes.c_size_t), ] kernel32.CreateJobObjectW.argtypes = [ctypes.POINTER(SECURITY_ATTRIBUTES), wintypes.LPCWSTR] kernel32.CreateJobObjectW.restype = wintypes.HANDLE kernel32.SetInformationJobObject.argtypes = [ wintypes.HANDLE, ctypes.c_int, wintypes.LPVOID, wintypes.DWORD, ] kernel32.SetInformationJobObject.restype = wintypes.BOOL kernel32.AssignProcessToJobObject.argtypes = [wintypes.HANDLE, wintypes.HANDLE] kernel32.AssignProcessToJobObject.restype = wintypes.BOOL kernel32.TerminateJobObject.argtypes = [wintypes.HANDLE, wintypes.UINT] kernel32.TerminateJobObject.restype = wintypes.BOOL kernel32.CloseHandle.argtypes = [wintypes.HANDLE] kernel32.CloseHandle.restype = wintypes.BOOL kernel32.ResumeThread.argtypes = [wintypes.HANDLE] kernel32.ResumeThread.restype = wintypes.DWORD kernel32.TerminateProcess.argtypes = [wintypes.HANDLE, wintypes.UINT] kernel32.TerminateProcess.restype = wintypes.BOOL kernel32.WaitForSingleObject.argtypes = [wintypes.HANDLE, wintypes.DWORD] kernel32.WaitForSingleObject.restype = wintypes.DWORD kernel32.GetExitCodeProcess.argtypes = [wintypes.HANDLE, ctypes.POINTER(wintypes.DWORD)] kernel32.GetExitCodeProcess.restype = wintypes.BOOL kernel32.SetHandleInformation.argtypes = [wintypes.HANDLE, wintypes.DWORD, wintypes.DWORD] kernel32.SetHandleInformation.restype = wintypes.BOOL kernel32.InitializeProcThreadAttributeList.argtypes = [ wintypes.LPVOID, wintypes.DWORD, wintypes.DWORD, ctypes.POINTER(ctypes.c_size_t), ] kernel32.InitializeProcThreadAttributeList.restype = wintypes.BOOL kernel32.UpdateProcThreadAttribute.argtypes = [ wintypes.LPVOID, wintypes.DWORD, ctypes.c_size_t, wintypes.LPVOID, ctypes.c_size_t, wintypes.LPVOID, wintypes.LPVOID, ] kernel32.UpdateProcThreadAttribute.restype = wintypes.BOOL kernel32.DeleteProcThreadAttributeList.argtypes = [wintypes.LPVOID] kernel32.DeleteProcThreadAttributeList.restype = None kernel32.CreateProcessW.argtypes = [ wintypes.LPCWSTR, wintypes.LPWSTR, ctypes.POINTER(SECURITY_ATTRIBUTES), ctypes.POINTER(SECURITY_ATTRIBUTES), wintypes.BOOL, wintypes.DWORD, wintypes.LPVOID, wintypes.LPCWSTR, ctypes.POINTER(STARTUPINFOW), ctypes.POINTER(PROCESS_INFORMATION), ] kernel32.CreateProcessW.restype = wintypes.BOOL kernel32.CreateFileW.argtypes = [ wintypes.LPCWSTR, wintypes.DWORD, wintypes.DWORD, ctypes.POINTER(SECURITY_ATTRIBUTES), wintypes.DWORD, wintypes.DWORD, wintypes.HANDLE, ] kernel32.CreateFileW.restype = wintypes.HANDLE def _win_error(prefix: str) -> JobError: return JobError(f"{prefix}: winerror={ctypes.get_last_error()}") def _close_handle(handle: int | None) -> None: if os.name == "nt" and handle: kernel32.CloseHandle(wintypes.HANDLE(handle)) def set_handle_inheritable(handle: int, inheritable: bool) -> None: if os.name != "nt": raise JobError("Windows is required") flags = HANDLE_FLAG_INHERIT if inheritable else 0 if not kernel32.SetHandleInformation(handle, HANDLE_FLAG_INHERIT, flags): raise _win_error("SetHandleInformation") def create_event(inheritable: bool = True) -> int: if os.name != "nt": raise JobError("Windows is required") create_event_w = kernel32.CreateEventW create_event_w.argtypes = [ctypes.POINTER(SECURITY_ATTRIBUTES), wintypes.BOOL, wintypes.BOOL, wintypes.LPCWSTR] create_event_w.restype = wintypes.HANDLE sa = SECURITY_ATTRIBUTES(ctypes.sizeof(SECURITY_ATTRIBUTES), None, bool(inheritable)) handle = create_event_w(ctypes.byref(sa), True, False, None) if not handle: raise _win_error("CreateEventW") return int(handle) def set_event(handle: int) -> None: set_event_fn = kernel32.SetEvent set_event_fn.argtypes = [wintypes.HANDLE] set_event_fn.restype = wintypes.BOOL if not set_event_fn(handle): raise _win_error("SetEvent") def is_event_set(handle: int) -> bool: result = kernel32.WaitForSingleObject(handle, 0) if result == WAIT_OBJECT_0: return True if result == WAIT_TIMEOUT: return False raise _win_error("WaitForSingleObject(event)") @dataclass class JobProcess: process_handle: int pid: int def wait(self, timeout_seconds: float | None = None) -> bool: timeout_ms = INFINITE if timeout_seconds is None else max(0, int(timeout_seconds * 1000)) result = kernel32.WaitForSingleObject(self.process_handle, timeout_ms) if result == WAIT_OBJECT_0: return True if result == WAIT_TIMEOUT: return False raise _win_error("WaitForSingleObject(process)") def exit_code(self) -> int | None: code = wintypes.DWORD() if not kernel32.GetExitCodeProcess(self.process_handle, ctypes.byref(code)): raise _win_error("GetExitCodeProcess") return None if code.value == 259 else int(code.value) def terminate(self, code: int = 125) -> None: if not kernel32.TerminateProcess(self.process_handle, code): error = ctypes.get_last_error() if error != 5: raise _win_error("TerminateProcess") def close(self) -> None: if self.process_handle: _close_handle(self.process_handle) self.process_handle = 0 def __enter__(self) -> "JobProcess": return self def __exit__(self, *_: object) -> None: self.close() class WindowsJob: """One kill-on-close job whose handle is never inherited.""" def __init__(self) -> None: if os.name != "nt": raise JobError("E_JOB_OBJECT") handle = kernel32.CreateJobObjectW(None, None) if not handle: raise _win_error("CreateJobObjectW") self.handle = int(handle) try: info = JOBOBJECT_EXTENDED_LIMIT_INFORMATION_STRUCT() info.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE if not kernel32.SetInformationJobObject( self.handle, JOB_OBJECT_EXTENDED_LIMIT_INFORMATION, ctypes.byref(info), ctypes.sizeof(info), ): raise _win_error("SetInformationJobObject") set_handle_inheritable(self.handle, False) except BaseException: self.close() raise def launch_suspended( self, command: Sequence[str], *, stdin_handle: int, stdout_handle: int, stderr_handle: int, inherited_handles: Iterable[int] = (), cwd: str | None = None, ) -> JobProcess: handles = list(dict.fromkeys([stdin_handle, stdout_handle, stderr_handle, *inherited_handles])) if not handles or any(not item for item in handles): raise JobError("invalid inherited handle") for handle in handles: set_handle_inheritable(handle, True) attribute_size = ctypes.c_size_t() kernel32.InitializeProcThreadAttributeList(None, 1, 0, ctypes.byref(attribute_size)) attribute_buffer = ctypes.create_string_buffer(attribute_size.value) attribute_list = ctypes.cast(attribute_buffer, wintypes.LPVOID) if not kernel32.InitializeProcThreadAttributeList( attribute_list, 1, 0, ctypes.byref(attribute_size) ): raise _win_error("InitializeProcThreadAttributeList") raw_handles = (wintypes.HANDLE * len(handles))(*handles) process_info = PROCESS_INFORMATION() try: if not kernel32.UpdateProcThreadAttribute( attribute_list, 0, PROC_THREAD_ATTRIBUTE_HANDLE_LIST, ctypes.cast(raw_handles, wintypes.LPVOID), ctypes.sizeof(raw_handles), None, None, ): raise _win_error("UpdateProcThreadAttribute") startup = STARTUPINFOEXW() startup.StartupInfo.cb = ctypes.sizeof(startup) startup.StartupInfo.dwFlags = STARTF_USESTDHANDLES startup.StartupInfo.hStdInput = stdin_handle startup.StartupInfo.hStdOutput = stdout_handle startup.StartupInfo.hStdError = stderr_handle startup.lpAttributeList = attribute_list command_line = ctypes.create_unicode_buffer(subprocess.list2cmdline(list(command))) flags = CREATE_SUSPENDED | CREATE_NO_WINDOW | EXTENDED_STARTUPINFO_PRESENT if not kernel32.CreateProcessW( None, command_line, None, None, True, flags, None, cwd, ctypes.byref(startup.StartupInfo), ctypes.byref(process_info), ): raise _win_error("CreateProcessW") try: if not kernel32.AssignProcessToJobObject(self.handle, process_info.hProcess): raise _win_error("AssignProcessToJobObject") if kernel32.ResumeThread(process_info.hThread) == 0xFFFFFFFF: raise _win_error("ResumeThread") return JobProcess(int(process_info.hProcess), int(process_info.dwProcessId)) except BaseException: kernel32.TerminateProcess(process_info.hProcess, 125) kernel32.WaitForSingleObject(process_info.hProcess, 5000) _close_handle(int(process_info.hProcess)) raise finally: _close_handle(int(process_info.hThread)) finally: kernel32.DeleteProcThreadAttributeList(attribute_list) for handle in handles: try: set_handle_inheritable(handle, False) except JobError: pass def terminate(self, code: int = 125) -> None: if self.handle and not kernel32.TerminateJobObject(self.handle, code): raise _win_error("TerminateJobObject") def close(self) -> None: if self.handle: _close_handle(self.handle) self.handle = 0 def __enter__(self) -> "WindowsJob": return self def __exit__(self, *_: object) -> None: self.close() def open_nul_handles() -> tuple[int, int, int]: """Open inheritable NUL handles suitable for STARTUPINFOEX.""" if os.name != "nt": raise JobError("Windows is required") sa = SECURITY_ATTRIBUTES(ctypes.sizeof(SECURITY_ATTRIBUTES), None, True) read_handle = kernel32.CreateFileW( "NUL", GENERIC_READ, 3, ctypes.byref(sa), OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, None ) write_handle = kernel32.CreateFileW( "NUL", GENERIC_WRITE, 3, ctypes.byref(sa), OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, None ) error_handle = kernel32.CreateFileW( "NUL", GENERIC_WRITE, 3, ctypes.byref(sa), OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, None ) if not read_handle or not write_handle or not error_handle: for handle in (read_handle, write_handle, error_handle): if handle: _close_handle(int(handle)) raise _win_error("CreateFileW(NUL)") return int(read_handle), int(write_handle), int(error_handle) def close_handles(*handles: int) -> None: for handle in handles: _close_handle(handle)