fix(device-test): contain tool processes and add reviewed preview
This commit is contained in:
chenmingxuan committed
1 parent
c98d7527a0
commit
7d31e09cb7
59 files changed
+23722
No files matched your search
@@ -0,0 +1,262 @@
|
||||
"""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
|
||||
Reference in new issue
Block a user