"""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)
|
