Mercurial > ~astiob > upreckon > hgweb
diff win32.py @ 80:809b77302b21
Win32-specific module with memory and CPU time limits
The Win32-specific implementation of call() and friends now lives
in module win32, looks clean and in addition is able to enforce memory
and CPU time limits on NT kernels, in particular on Windows 2000 and up
asking the system to terminate the process as soon as or (in the case
of CPU time) almost as soon as the limits are broken. According to my
observations, malloc() in the limited process does not return NULL
when memory usage is close to the limit and instead crashes the process
(which Upreckon happily translates into 'memory limit exceeded').
The catch is that the module is not actually used yet; coming soon.
author | Oleg Oshmyan <chortos@inbox.lv> |
---|---|
date | Wed, 16 Feb 2011 00:01:33 +0000 |
parents | |
children | 24752db487c5 |
line wrap: on
line diff
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/win32.py Wed Feb 16 00:01:33 2011 +0000 @@ -0,0 +1,497 @@ +# Copyright (c) 2011 Chortos-2 <chortos@inbox.lv> + +from __future__ import division, with_statement + +try: + from compat import * +except ImportError: + import __main__ + __main__.import_error(sys.exc_info()[1]) + +from __main__ import clock +from ctypes import * +from ctypes.wintypes import * +from subprocess import Popen + +# Defaults that may be overwritten by values from _subprocess +INFINITE = -1 +STD_INPUT_HANDLE = -10 +WAIT_OBJECT_0 = 0 + +try: + from _subprocess import * +except ImportError: + pass + +try: + from numbers import Integral +except ImportError: + Integral = int, long + +try: + from collections import namedtuple +except ImportError: + from operator import itemgetter + class ProcessTimes(tuple): + __slots__ = () + __new__ = lambda cls, kernel, user: tuple.__new__(cls, (kernel, user)) + __getnewargs__ = lambda self: tuple(self) + kernel, user = (property(itemgetter(i)) for i in (0, 1)) +else: + ProcessTimes = namedtuple('ProcessTimes', 'kernel user') + +__all__ = 'call', 'kill', 'terminate' + + +# Automatically convert _subprocess handle objects into low-level HANDLEs +# and replicate their functionality for our own use +try: + _subprocess_handle = type(GetCurrentProcess()) +except NameError: + _subprocess_handle = Integral +class Handle(object): + @staticmethod + def from_param(handle): + if isinstance(handle, (_subprocess_handle, Integral)): + return HANDLE(int(handle)) + elif isinstance(handle, Handle): + return HANDLE(handle.handle) + elif isinstance(handle, HANDLE): + return handle + else: + raise TypeError('cannot convert %s to a handle' % + type(handle).__name__) + + __slots__ = 'handle' + + def __init__(self, handle): + if isinstance(handle, Integral): + self.handle = handle + elif isinstance(handle, HANDLE): + self.handle = handle.value + elif isinstance(handle, Handle): + self.handle = handle.handle + elif isinstance(handle, _subprocess_handle): + handle = HANDLE(int(handle)) + flags = DWORD() + try: + if windll.kernel32.GetHandleInformation(handle, byref(flags)): + flags = flags.value + else: + flags = 0 + except AttributeError: + # Available on NT 3.51 and up, NT line only + flags = 0 + proc = HANDLE(int(GetCurrentProcess())) + handle = DuplicateHandle(proc, handle, proc, 0, flags & 1, 2) + self.handle = handle.Detach() + else: + raise TypeError("Handle() argument must be a handle, not '%s'" % + type(name).__name__) + + def __int__(self): + return int(self.handle) + + def Detach(self): + handle = self.handle + self.handle = None + return handle + + # This is also __del__, so only locals are accessed + def Close(self, _CloseHandle=windll.kernel32.CloseHandle, _HANDLE=HANDLE): + if self.handle: + _CloseHandle(_HANDLE(self.handle)) + self.handle = None + __del__ = Close + +CHAR = c_char +INVALID_HANDLE_VALUE = HANDLE(-1).value +LPDWORD = POINTER(DWORD) +LPFILETIME = POINTER(FILETIME) +SIZE_T = ULONG_PTR = WPARAM +ULONGLONG = c_ulonglong + +try: + unicode +except NameError: + LPCTSTR = LPCWSTR + unisuffix = 'W' +else: + LPCTSTR = LPCSTR + unisuffix = 'A' + + +prototype = WINFUNCTYPE(BOOL, Handle, + LPFILETIME, LPFILETIME, LPFILETIME, LPFILETIME) +flags = ((1, 'process'), + (2, 'creation'), (2, 'exit'), (2, 'kernel'), (2, 'user')) +try: + GetProcessTimes = prototype(('GetProcessTimes', windll.kernel32), flags) +except AttributeError: + # Available on NT 3.5 and up, NT line only + GetProcessTimes = None +else: + def errcheck(result, func, args): + if not result: raise WinError() + ftimes = [t.dwHighDateTime << 32 | t.dwLowDateTime for t in args[3:]] + kernel = ftimes[0] / 10000000 + user = ftimes[1] / 10000000 + return ProcessTimes(kernel, user) + GetProcessTimes.errcheck = errcheck + + +class PROCESS_MEMORY_COUNTERS(Structure): + _fields_ = (('cb', DWORD), + ('PageFaultCount', DWORD), + ('PeakWorkingSetSize', SIZE_T), + ('WorkingSetSize', SIZE_T), + ('QuotaPeakPagedPoolUsage', SIZE_T), + ('QuotaPagedPoolUsage', SIZE_T), + ('QuotaPeakNonPagedPoolUsage', SIZE_T), + ('QuotaNonPagedPoolUsage', SIZE_T), + ('PagefileUsage', SIZE_T), + ('PeakPagefileUsage', SIZE_T)) + +prototype = WINFUNCTYPE(BOOL, Handle, POINTER(PROCESS_MEMORY_COUNTERS), DWORD) +flags = ((1, 'process'), (2, 'counters'), + (5, 'cb', sizeof(PROCESS_MEMORY_COUNTERS))) +try: + GetProcessMemoryInfo = prototype(('GetProcessMemoryInfo', windll.psapi), + flags) +except AttributeError: + # Available on NT 4.0 and up, NT line only + GetProcessMemoryInfo = None +else: + def errcheck(result, func, args): + if not result: raise WinError() + return args + GetProcessMemoryInfo.errcheck = errcheck + + +class _uChar_union(Union): + _fields_ = (('UnicodeChar', WCHAR), + ('AsciiChar', CHAR)) + +class KEY_EVENT_RECORD(Structure): + _fields_ = (('bKeyDown', BOOL), + ('wRepeatCount', WORD), + ('wVirtualKeyCode', WORD), + ('wVirtualScanCode', WORD), + ('uChar', _uChar_union), + ('dwControlKeyState', DWORD)) + +RIGHT_ALT_PRESSED = 0x001 +LEFT_ALT_PRESSED = 0x002 +RIGHT_CTRL_PRESSED = 0x004 +LEFT_CTRL_PRESSED = 0x008 +SHIFT_PRESSED = 0x010 +NUMLOCK_ON = 0x020 +SCROLLLOCK_ON = 0x040 +CAPSLOCK_ON = 0x080 +ENHANCED_KEY = 0x100 + +class _Event_union(Union): + _fields_ = ('KeyEvent', KEY_EVENT_RECORD), + +class INPUT_RECORD(Structure): + _fields_ = (('EventType', WORD), + ('Event', _Event_union)) + +KEY_EVENT = 0x01 +MOUSE_EVENT = 0x02 +WINDOW_BUFFER_SIZE_EVENT = 0x04 +MENU_EVENT = 0x08 +FOCUS_EVENT = 0x10 + +prototype = WINFUNCTYPE(BOOL, Handle, POINTER(INPUT_RECORD), DWORD, LPDWORD) +flags = (1, 'input'), (2, 'buffer'), (5, 'length', 1), (2, 'number_read') +ReadConsoleInput = prototype(('ReadConsoleInputA', windll.kernel32), flags) +def errcheck(result, func, args): + if not result: raise WinError() + return args[1] if args[3] else None +ReadConsoleInput.errcheck = errcheck + + +prototype = WINFUNCTYPE(BOOL, Handle) +flags = (1, 'input'), +FlushConsoleInputBuffer = prototype(('FlushConsoleInputBuffer', + windll.kernel32), flags) +def errcheck(result, func, args): + if not result: raise WinError() +FlushConsoleInputBuffer.errcheck = errcheck + + +prototype = WINFUNCTYPE(BOOL, Handle, DWORD) +flags = (1, 'console'), (1, 'mode') +SetConsoleMode = prototype(('SetConsoleMode', windll.kernel32), flags) +def errcheck(result, func, args): + if not result: raise WinError() +SetConsoleMode.errcheck = errcheck + +ENABLE_PROCESSED_INPUT = 0x001 +ENABLE_LINE_INPUT = 0x002 +ENABLE_ECHO_INPUT = 0x004 +ENABLE_WINDOW_INPUT = 0x008 +ENABLE_MOUSE_INPUT = 0x010 +ENABLE_INSERT_MODE = 0x020 +ENABLE_QUICK_EDIT_MODE = 0x040 +ENABLE_EXTENDED_FLAGS = 0x080 + +ENABLE_PROCESSED_OUTPUT = 1 +ENABLE_WRAP_AT_EOL_OUTPUT = 2 + + +prototype = WINFUNCTYPE(HANDLE, c_void_p, LPCTSTR) +flags = (5, 'attributes'), (1, 'name') +try: + CreateJobObject = prototype(('CreateJobObject'+unisuffix, windll.kernel32), + flags) +except AttributeError: + # Available on 2000 and up, NT line only + CreateJobObject = lambda name: None +else: + def errcheck(result, func, args): + if not result: raise WinError() + return Handle(result) + CreateJobObject.errcheck = errcheck + + +prototype = WINFUNCTYPE(BOOL, Handle, Handle) +flags = (1, 'job'), (1, 'handle') +try: + AssignProcessToJobObject = prototype(('AssignProcessToJobObject', + windll.kernel32), flags) +except AttributeError: + # Available on 2000 and up, NT line only + AssignProcessToJobObject = lambda job, handle: None +else: + def errcheck(result, func, args): + if not result: raise WinError() + AssignProcessToJobObject.errcheck = errcheck + + +class JOBOBJECT_BASIC_LIMIT_INFORMATION(Structure): + _fields_ = (('PerProcessUserTimeLimit', LARGE_INTEGER), + ('PerJobUserTimeLimit', LARGE_INTEGER), + ('LimitFlags', DWORD), + ('MinimumWorkingSetSize', SIZE_T), + ('MaximumWorkingSetSize', SIZE_T), + ('ActiveProcessLimit', DWORD), + ('Affinity', ULONG_PTR), + ('PriorityClass', DWORD), + ('SchedulingClass', DWORD)) + +JOB_OBJECT_LIMIT_WORKINGSET = 0x0001 +JOB_OBJECT_LIMIT_PROCESS_TIME = 0x0002 +JOB_OBJECT_LIMIT_JOB_TIME = 0x0004 +JOB_OBJECT_LIMIT_ACTIVE_PROCESS = 0x0008 +JOB_OBJECT_LIMIT_AFFINITY = 0x0010 +JOB_OBJECT_LIMIT_PRIORITY_CLASS = 0x0020 +JOB_OBJECT_LIMIT_PRESERVE_JOB_TIME = 0x0040 +JOB_OBJECT_LIMIT_SCHEDULING_CLASS = 0x0080 +JOB_OBJECT_LIMIT_PROCESS_MEMORY = 0x0100 +JOB_OBJECT_LIMIT_JOB_MEMORY = 0x0200 +JOB_OBJECT_LIMIT_DIE_ON_UNHANDLED_EXCEPTION = 0x0400 +JOB_OBJECT_LIMIT_BREAKAWAY_OK = 0x0800 +JOB_OBJECT_LIMIT_SILENT_BREAKAWAY_OK = 0x1000 +JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE = 0x2000 +JOB_OBJECT_LIMIT_SUBSET_AFFINITY = 0x4000 + +class IO_COUNTERS(Structure): + _fields_ = (('ReadOperationCount', ULONGLONG), + ('WriteOperationCount', ULONGLONG), + ('OtherOperationCount', ULONGLONG), + ('ReadTransferCount', ULONGLONG), + ('WriteTransferCount', ULONGLONG), + ('OtherTransferCount', ULONGLONG)) + +class JOBOBJECT_EXTENDED_LIMIT_INFORMATION(Structure): + _fields_ = (('BasicLimitInformation', JOBOBJECT_BASIC_LIMIT_INFORMATION), + ('IoInfo', IO_COUNTERS), + ('ProcessMemoryLimit', SIZE_T), + ('JobMemoryLimit', SIZE_T), + ('PeakProcessMemoryUsed', SIZE_T), + ('PeakJobMemoryUsed', SIZE_T)) + +prototype = WINFUNCTYPE(BOOL, HANDLE, c_int, c_void_p, DWORD) +flags = (1, 'job'), (1, 'infoclass'), (1, 'info'), (1, 'infosize') +try: + _setjobinfo = prototype(('SetInformationJobObject',windll.kernel32), flags) +except AttributeError: + # Available on 2000 and up, NT line only + SetInformationJobObject = lambda job, infoclass, info: None +else: + def errcheck(result, func, args): + if not result: raise WinError() + _setjobinfo.errcheck = errcheck + def SetInformationJobObject(job, infoclass, info): + return _setjobinfo(job, infoclass, info, sizeof(info)) + +( + JobObjectBasicAccountingInformation, + JobObjectBasicLimitInformation, + JobObjectBasicProcessIdList, + JobObjectBasicUIRestrictions, + JobObjectSecurityLimitInformation, + JobObjectEndOfJobTimeInformation, + JobObjectAssociateCompletionPortInformation, + JobObjectBasicAndIoAccountingInformation, + JobObjectExtendedLimitInformation, + JobObjectJobSetInformation, + MaxJobObjectInfoClass +) = range(1, 12) + + +prototype = WINFUNCTYPE(DWORD, DWORD, POINTER(HANDLE), BOOL, DWORD) +flags = (1, 'count'), (1, 'handles'), (1, 'wait_all'), (1, 'milliseconds') +_wait_multiple = prototype(('WaitForMultipleObjects', windll.kernel32), flags) +def errcheck(result, func, args): + if result == WAIT_FAILED: raise WinError() + return args +_wait_multiple.errcheck = errcheck +def WaitForMultipleObjects(handles, wait_all, timeout): + n = len(handles) + handles = (Handle.from_param(handle) for handle in handles) + timeout = ceil(timeout * 1000) + return _wait_multiple(n, (HANDLE * n)(*handles), wait_all, timeout) + +# WAIT_OBJECT_0 defined at the top of the file +WAIT_ABANDONED_0 = 0x00000080 +WAIT_TIMEOUT = 0x00000102 +WAIT_FAILED = 0xFFFFFFFF + + +try: + _wait_single = WaitForSingleObject +except NameError: + prototype = WINFUNCTYPE(DWORD, Handle, DWORD) + flags = (1, 'handle'), (1, 'milliseconds') + _wait_single = prototype(('WaitForSingleObject', windll.kernel32), flags) + def errcheck(result, func, args): + if result == WAIT_FAILED: raise WinError() + return args + _wait_single.errcheck = errcheck +def WaitForSingleObject(handle, timeout): + return _wait_single(handle, ceil(timeout * 1000)) + + +try: + GetStdHandle +except NameError: + prototype = WINFUNCTYPE(HANDLE, DWORD) + flags = (1, 'which'), + GetStdHandle = prototype(('GetStdHandle', windll.kernel32), flags) + def errcheck(result, func, args): + if result == INVALID_HANDLE_VALUE: raise WinError() + return args if result else None + GetStdHandle.errcheck = errcheck + + +try: + TerminateProcess +except NameError: + prototype = WINFUNCTYPE(BOOL, Handle, UINT) + flags = (1, 'process'), (1, 'exitcode') + TerminateProcess = prototype(('TerminateProcess', windll.kernel32), flags) + def errcheck(result, func, args): + if not result: raise WinError() + TerminateProcess.errcheck = errcheck + + +# Do not show error messages due to errors in the program being tested +try: + errmode = ctypes.windll.kernel32.GetErrorMode() +except AttributeError: + # GetErrorMode is available on Vista/2008 and up + errmode = ctypes.windll.kernel32.SetErrorMode(0) +ctypes.windll.kernel32.SetErrorMode(errmode | 0x8003) + +stdin = GetStdHandle(STD_INPUT_HANDLE) +try: + SetConsoleMode(stdin, ENABLE_PROCESSED_INPUT) +except WindowsError: + console_input = False +else: + console_input = True + FlushConsoleInputBuffer(stdin) + +def kill(process): + try: + process.terminate() + except AttributeError: + TerminateProcess(process._handle) +terminate = kill + +def call(*args, **kwargs): + case = kwargs.pop('case') + job = CreateJobObject(None) + flags = 0 + if case.maxcputime: + flags |= JOB_OBJECT_LIMIT_PROCESS_TIME + if case.maxmemory: + flags |= JOB_OBJECT_LIMIT_PROCESS_MEMORY + limits = JOBOBJECT_EXTENDED_LIMIT_INFORMATION( + JOBOBJECT_BASIC_LIMIT_INFORMATION( + PerProcessUserTimeLimit=ceil((case.maxcputime or 0)*10000000), + LimitFlags=flags, + ), + ProcessMemoryLimit=ceil((case.maxmemory or 0)*1048576), + ) + SetInformationJobObject(job, JobObjectExtendedLimitInformation, limits) + try: + case.process = Popen(*args, **kwargs) + except OSError: + raise CannotStartTestee(sys.exc_info()[1]) + case.time_started = clock() + AssignProcessToJobObject(job, case.process._handle) + if not console_input: + if case.maxwalltime: + if (WaitForSingleObject(case.process._handle, case.maxwalltime) != + WAIT_OBJECT_0): + raise TimeLimitExceeded + else: + case.process.wait() + else: + handles = stdin, case.process._handle + if case.maxwalltime: + time_end = clock() + case.maxwalltime + while case.process.poll() is None: + remaining = time_end - clock() + if remaining > 0: + if (WaitForMultipleObjects(handles, False, remaining) == + WAIT_OBJECT_0): + ir = ReadConsoleInput(stdin) + if (ir and + ir.EventType == 1 and + ir.Event.KeyEvent.bKeyDown and + ir.Event.KeyEvent.wVirtualKeyCode == 27): + raise CanceledByUser + else: + raise TimeLimitExceeded + else: + while case.process.poll() is None: + if (WaitForMultipleObjects(handles, False, INFINITE) == + WAIT_OBJECT_0): + ir = ReadConsoleInput(stdin) + if (ir and + ir.EventType == 1 and + ir.Event.KeyEvent.bKeyDown and + ir.Event.KeyEvent.wVirtualKeyCode == 27): + raise CanceledByUser + case.time_stopped = clock() + if case.maxcputime and GetProcessTimes: + try: + times = GetProcessTimes(case.process._handle) + except WindowsError: + pass + else: + if times.kernel + times.user > case.maxcputime: + raise TimeLimitExceeded + if case.maxmemory and GetProcessMemoryInfo: + try: + counters = GetProcessMemoryInfo(case.process._handle) + except WindowsError: + pass + else: + if counters.PeakPagefileUsage > case.maxmemory * 1048576: + raise MemoryLimitExceeded \ No newline at end of file