67 lines
3.4 KiB
Python
67 lines
3.4 KiB
Python
|
"""Windows resource containment for the inspector and all descendant workers.
|
||
|
|
|
||
|
|
This is a Job Object, not a restricted-token or filesystem/network sandbox.
|
||
|
|
The inspector owns the handle: terminating it also terminates its model workers.
|
||
|
|
"""
|
||
|
|
import ctypes
|
||
|
|
import sys
|
||
|
|
from ctypes import wintypes
|
||
|
|
|
||
|
|
_job_handle = None
|
||
|
|
MEMORY_LIMIT = 4 * 1024 * 1024 * 1024
|
||
|
|
PROCESS_LIMIT = 4
|
||
|
|
|
||
|
|
|
||
|
|
def contain_inspector():
|
||
|
|
global _job_handle
|
||
|
|
if sys.platform != "win32":
|
||
|
|
return {"active": False, "kind": "unsupported-platform"}
|
||
|
|
if _job_handle:
|
||
|
|
return {"active": True, "kind": "windows-job-v1"}
|
||
|
|
|
||
|
|
class BasicLimits(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 IOCounters(ctypes.Structure):
|
||
|
|
_fields_ = [(name, ctypes.c_uint64) for name in
|
||
|
|
("ReadOperationCount", "WriteOperationCount", "OtherOperationCount",
|
||
|
|
"ReadTransferCount", "WriteTransferCount", "OtherTransferCount")]
|
||
|
|
|
||
|
|
class ExtendedLimits(ctypes.Structure):
|
||
|
|
_fields_ = [("BasicLimitInformation", BasicLimits), ("IoInfo", IOCounters),
|
||
|
|
("ProcessMemoryLimit", ctypes.c_size_t), ("JobMemoryLimit", ctypes.c_size_t),
|
||
|
|
("PeakProcessMemoryUsed", ctypes.c_size_t), ("PeakJobMemoryUsed", ctypes.c_size_t)]
|
||
|
|
|
||
|
|
kernel = ctypes.WinDLL("kernel32", use_last_error=True)
|
||
|
|
kernel.CreateJobObjectW.argtypes = [ctypes.c_void_p, wintypes.LPCWSTR]
|
||
|
|
kernel.CreateJobObjectW.restype = wintypes.HANDLE
|
||
|
|
kernel.SetInformationJobObject.argtypes = [wintypes.HANDLE, ctypes.c_int, ctypes.c_void_p, wintypes.DWORD]
|
||
|
|
kernel.SetInformationJobObject.restype = wintypes.BOOL
|
||
|
|
kernel.AssignProcessToJobObject.argtypes = [wintypes.HANDLE, wintypes.HANDLE]
|
||
|
|
kernel.AssignProcessToJobObject.restype = wintypes.BOOL
|
||
|
|
kernel.GetCurrentProcess.restype = wintypes.HANDLE
|
||
|
|
kernel.CloseHandle.argtypes = [wintypes.HANDLE]
|
||
|
|
|
||
|
|
handle = kernel.CreateJobObjectW(None, None)
|
||
|
|
if not handle:
|
||
|
|
raise RuntimeError("Inspector resource containment could not be created")
|
||
|
|
limits = ExtendedLimits()
|
||
|
|
# KILL_ON_JOB_CLOSE | JOB_MEMORY | ACTIVE_PROCESS. No BREAKAWAY flags.
|
||
|
|
limits.BasicLimitInformation.LimitFlags = 0x2000 | 0x200 | 0x8
|
||
|
|
limits.BasicLimitInformation.ActiveProcessLimit = PROCESS_LIMIT
|
||
|
|
limits.JobMemoryLimit = MEMORY_LIMIT
|
||
|
|
if not kernel.SetInformationJobObject(handle, 9, ctypes.byref(limits), ctypes.sizeof(limits)):
|
||
|
|
kernel.CloseHandle(handle)
|
||
|
|
raise RuntimeError("Inspector resource containment limits could not be applied")
|
||
|
|
if not kernel.AssignProcessToJobObject(handle, kernel.GetCurrentProcess()):
|
||
|
|
kernel.CloseHandle(handle)
|
||
|
|
raise RuntimeError("Inspector resource containment could not be attached")
|
||
|
|
# Retain a non-inheritable owner handle until process exit. Do not close early:
|
||
|
|
# KILL_ON_JOB_CLOSE also terminates the current inspector process.
|
||
|
|
_job_handle = handle
|
||
|
|
return {"active": True, "kind": "windows-job-v1", "memoryBytes": MEMORY_LIMIT,
|
||
|
|
"maxProcesses": PROCESS_LIMIT, "killsDescendants": True}
|