263 lines
9.7 KiB
Python
263 lines
9.7 KiB
Python
"""Shared Windows process ownership for the plugin's fixed external tools."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import subprocess
|
|
import threading
|
|
import time
|
|
|
|
|
|
class ManagedProcessError(OSError):
|
|
"""Stable error with execution facts; never includes tool arguments/output."""
|
|
|
|
def __init__(self, code, *, may_have_started=False):
|
|
super().__init__(code)
|
|
self.error_code = code
|
|
self.may_have_started = may_have_started
|
|
|
|
|
|
class WindowsJob:
|
|
def __init__(self):
|
|
import ctypes
|
|
from ctypes import wintypes
|
|
|
|
class IO_COUNTERS(ctypes.Structure):
|
|
_fields_ = [
|
|
(n, ctypes.c_ulonglong)
|
|
for n in (
|
|
"ReadOperationCount",
|
|
"WriteOperationCount",
|
|
"OtherOperationCount",
|
|
"ReadTransferCount",
|
|
"WriteTransferCount",
|
|
"OtherTransferCount",
|
|
)
|
|
]
|
|
|
|
class BASIC_LIMIT(ctypes.Structure):
|
|
_fields_ = [
|
|
("PerProcessUserTimeLimit", ctypes.c_longlong),
|
|
("PerJobUserTimeLimit", ctypes.c_longlong),
|
|
("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 EXTENDED_LIMIT(ctypes.Structure):
|
|
_fields_ = [
|
|
("BasicLimitInformation", BASIC_LIMIT),
|
|
("IoInfo", IO_COUNTERS),
|
|
("ProcessMemoryLimit", ctypes.c_size_t),
|
|
("JobMemoryLimit", ctypes.c_size_t),
|
|
("PeakProcessMemoryUsed", ctypes.c_size_t),
|
|
("PeakJobMemoryUsed", ctypes.c_size_t),
|
|
]
|
|
|
|
class ACCOUNTING(ctypes.Structure):
|
|
_fields_ = [
|
|
("TotalUserTime", ctypes.c_longlong),
|
|
("TotalKernelTime", ctypes.c_longlong),
|
|
("ThisPeriodTotalUserTime", ctypes.c_longlong),
|
|
("ThisPeriodTotalKernelTime", ctypes.c_longlong),
|
|
("TotalPageFaultCount", wintypes.DWORD),
|
|
("TotalProcesses", wintypes.DWORD),
|
|
("ActiveProcesses", wintypes.DWORD),
|
|
("TotalTerminatedProcesses", wintypes.DWORD),
|
|
]
|
|
|
|
self.ctypes = ctypes
|
|
self.accounting = ACCOUNTING
|
|
self.kernel = k = ctypes.WinDLL("kernel32", use_last_error=True)
|
|
self.ntdll = ctypes.WinDLL("ntdll")
|
|
k.CreateJobObjectW.argtypes = [wintypes.LPVOID, wintypes.LPCWSTR]
|
|
k.CreateJobObjectW.restype = wintypes.HANDLE
|
|
k.SetInformationJobObject.argtypes = [
|
|
wintypes.HANDLE,
|
|
ctypes.c_int,
|
|
wintypes.LPVOID,
|
|
wintypes.DWORD,
|
|
]
|
|
k.SetInformationJobObject.restype = wintypes.BOOL
|
|
k.AssignProcessToJobObject.argtypes = [wintypes.HANDLE, wintypes.HANDLE]
|
|
k.AssignProcessToJobObject.restype = wintypes.BOOL
|
|
k.TerminateJobObject.argtypes = [wintypes.HANDLE, wintypes.UINT]
|
|
k.TerminateJobObject.restype = wintypes.BOOL
|
|
k.QueryInformationJobObject.argtypes = [
|
|
wintypes.HANDLE,
|
|
ctypes.c_int,
|
|
wintypes.LPVOID,
|
|
wintypes.DWORD,
|
|
wintypes.LPVOID,
|
|
]
|
|
k.QueryInformationJobObject.restype = wintypes.BOOL
|
|
k.CloseHandle.argtypes = [wintypes.HANDLE]
|
|
k.CloseHandle.restype = wintypes.BOOL
|
|
self.ntdll.NtResumeProcess.argtypes = [wintypes.HANDLE]
|
|
self.ntdll.NtResumeProcess.restype = ctypes.c_long
|
|
self.handle = k.CreateJobObjectW(None, None)
|
|
if not self.handle:
|
|
raise ManagedProcessError("process_ownership_failed")
|
|
limits = EXTENDED_LIMIT()
|
|
limits.BasicLimitInformation.LimitFlags = (
|
|
0x2000 # KILL_ON_JOB_CLOSE; handle is not inherited.
|
|
)
|
|
if not k.SetInformationJobObject(
|
|
self.handle, 9, ctypes.byref(limits), ctypes.sizeof(limits)
|
|
):
|
|
self.close()
|
|
raise ManagedProcessError("process_ownership_failed")
|
|
|
|
def assign(self, process):
|
|
if not self.kernel.AssignProcessToJobObject(self.handle, int(process._handle)):
|
|
raise ManagedProcessError("process_ownership_failed")
|
|
|
|
def resume(self, process):
|
|
if self.ntdll.NtResumeProcess(int(process._handle)) != 0:
|
|
raise ManagedProcessError(
|
|
"process_resume_unconfirmed", may_have_started=True
|
|
)
|
|
|
|
def stop(self):
|
|
c = self.ctypes
|
|
if not self.kernel.TerminateJobObject(self.handle, 1):
|
|
raise OSError("job_termination_unconfirmed")
|
|
deadline = time.monotonic() + 10
|
|
while True:
|
|
account = self.accounting()
|
|
if not self.kernel.QueryInformationJobObject(
|
|
self.handle, 1, c.byref(account), c.sizeof(account), None
|
|
):
|
|
raise OSError("job_cleanup_unconfirmed")
|
|
if account.ActiveProcesses == 0:
|
|
return
|
|
if time.monotonic() >= deadline:
|
|
raise OSError("job_cleanup_timeout")
|
|
time.sleep(0.02)
|
|
|
|
def close(self):
|
|
if self.handle:
|
|
handle, self.handle = self.handle, None
|
|
if not self.kernel.CloseHandle(handle):
|
|
raise OSError("job_handle_close_unconfirmed")
|
|
|
|
|
|
def run(argv, *, timeout=None, check=False, max_output_bytes=4 * 1024 * 1024, **kwargs):
|
|
"""Contain before resuming; bound pipes; reap the tree on every exit path."""
|
|
if os.name != "nt":
|
|
raise ManagedProcessError("process_ownership_requires_windows")
|
|
text = bool(
|
|
kwargs.get("text")
|
|
or kwargs.get("universal_newlines")
|
|
or kwargs.get("encoding")
|
|
or kwargs.get("errors")
|
|
)
|
|
process = None
|
|
assigned = started = False
|
|
job = WindowsJob()
|
|
buffers = {}
|
|
readers = []
|
|
output_limit = threading.Event()
|
|
read_error = threading.Event()
|
|
lock = threading.Lock()
|
|
size = 0
|
|
|
|
def read_pipe(name, pipe):
|
|
nonlocal size
|
|
buffers[name] = bytearray()
|
|
try:
|
|
while data := os.read(pipe.fileno(), 65536):
|
|
with lock:
|
|
size += len(data)
|
|
if size > max_output_bytes:
|
|
output_limit.set()
|
|
elif not output_limit.is_set():
|
|
buffers[name].extend(data)
|
|
except OSError:
|
|
read_error.set()
|
|
|
|
try:
|
|
kwargs["creationflags"] = (
|
|
kwargs.get("creationflags", 0) | subprocess.CREATE_NO_WINDOW | 0x4
|
|
)
|
|
try:
|
|
process = subprocess.Popen(argv, **kwargs)
|
|
except OSError:
|
|
raise ManagedProcessError("process_launch_failed") from None
|
|
job.assign(process)
|
|
assigned = True
|
|
started = True # Set before resume: failure cannot prove the tool did not run.
|
|
job.resume(process)
|
|
for name in ("stdout", "stderr"):
|
|
if pipe := getattr(process, name):
|
|
thread = threading.Thread(
|
|
target=read_pipe, args=(name, pipe), daemon=True
|
|
)
|
|
thread.start()
|
|
readers.append(thread)
|
|
deadline = time.monotonic() + timeout if timeout is not None else None
|
|
while process.poll() is None:
|
|
if output_limit.is_set() or read_error.is_set():
|
|
break
|
|
remaining = deadline - time.monotonic() if deadline is not None else 0.05
|
|
if remaining <= 0:
|
|
raise subprocess.TimeoutExpired(argv, timeout)
|
|
try:
|
|
process.wait(timeout=min(0.05, remaining))
|
|
except subprocess.TimeoutExpired:
|
|
continue
|
|
finally:
|
|
cleanup_failed = False
|
|
try:
|
|
if process is not None:
|
|
if assigned:
|
|
job.stop()
|
|
elif process.poll() is None:
|
|
process.kill() # Still suspended; no unowned descendants can exist.
|
|
process.wait(timeout=10)
|
|
except (OSError, subprocess.TimeoutExpired):
|
|
cleanup_failed = True
|
|
finally:
|
|
try:
|
|
job.close()
|
|
except OSError:
|
|
cleanup_failed = True
|
|
for thread in readers:
|
|
thread.join(timeout=2)
|
|
cleanup_failed |= thread.is_alive()
|
|
for name in ("stdin", "stdout", "stderr"):
|
|
if (pipe := getattr(process, name, None)) and not cleanup_failed:
|
|
try:
|
|
pipe.close()
|
|
except OSError:
|
|
cleanup_failed = True
|
|
if cleanup_failed:
|
|
raise ManagedProcessError(
|
|
"process_cleanup_unconfirmed", may_have_started=started
|
|
) from None
|
|
if output_limit.is_set():
|
|
raise ManagedProcessError("process_output_limit", may_have_started=True)
|
|
if read_error.is_set():
|
|
raise ManagedProcessError("process_output_interrupted", may_have_started=True)
|
|
output = {}
|
|
for name in ("stdout", "stderr"):
|
|
data = bytes(buffers[name]) if name in buffers else None
|
|
if text and data is not None:
|
|
pipe = getattr(process, name)
|
|
data = (
|
|
data.decode(pipe.encoding, pipe.errors)
|
|
.replace("\r\n", "\n")
|
|
.replace("\r", "\n")
|
|
)
|
|
output[name] = data
|
|
result = subprocess.CompletedProcess(
|
|
argv, process.returncode, output["stdout"], output["stderr"]
|
|
)
|
|
if check:
|
|
result.check_returncode()
|
|
return result
|