Mercurial > ~astiob > upreckon > hgweb

comparison win32.py @ 80:809b77302b21

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
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
comparison
equal deleted inserted replaced
79:ee8a99dcaaed 80:809b77302b21
1 # Copyright (c) 2011 Chortos-2 <chortos@inbox.lv>
2
3 from __future__ import division, with_statement
4
5 try:
6 from compat import *
7 except ImportError:
8 import __main__
9 __main__.import_error(sys.exc_info()[1])
10
11 from __main__ import clock
12 from ctypes import *
13 from ctypes.wintypes import *
14 from subprocess import Popen
15
16 # Defaults that may be overwritten by values from _subprocess
17 INFINITE = -1
18 STD_INPUT_HANDLE = -10
19 WAIT_OBJECT_0 = 0
20
21 try:
22 from _subprocess import *
23 except ImportError:
24 pass
25
26 try:
27 from numbers import Integral
28 except ImportError:
29 Integral = int, long
30
31 try:
32 from collections import namedtuple
33 except ImportError:
34 from operator import itemgetter
35 class ProcessTimes(tuple):
36 __slots__ = ()
37 __new__ = lambda cls, kernel, user: tuple.__new__(cls, (kernel, user))
38 __getnewargs__ = lambda self: tuple(self)
39 kernel, user = (property(itemgetter(i)) for i in (0, 1))
40 else:
41 ProcessTimes = namedtuple('ProcessTimes', 'kernel user')
42
43 __all__ = 'call', 'kill', 'terminate'
44
45
46 # Automatically convert _subprocess handle objects into low-level HANDLEs
47 # and replicate their functionality for our own use
48 try:
49 _subprocess_handle = type(GetCurrentProcess())
50 except NameError:
51 _subprocess_handle = Integral
52 class Handle(object):
53 @staticmethod
54 def from_param(handle):
55 if isinstance(handle, (_subprocess_handle, Integral)):
56 return HANDLE(int(handle))
57 elif isinstance(handle, Handle):
58 return HANDLE(handle.handle)
59 elif isinstance(handle, HANDLE):
60 return handle
61 else:
62 raise TypeError('cannot convert %s to a handle' %
63 type(handle).__name__)
64
65 __slots__ = 'handle'
66
67 def __init__(self, handle):
68 if isinstance(handle, Integral):
69 self.handle = handle
70 elif isinstance(handle, HANDLE):
71 self.handle = handle.value
72 elif isinstance(handle, Handle):
73 self.handle = handle.handle
74 elif isinstance(handle, _subprocess_handle):
75 handle = HANDLE(int(handle))
76 flags = DWORD()
77 try:
78 if windll.kernel32.GetHandleInformation(handle, byref(flags)):
79 flags = flags.value
80 else:
81 flags = 0
82 except AttributeError:
83 # Available on NT 3.51 and up, NT line only
84 flags = 0
85 proc = HANDLE(int(GetCurrentProcess()))
86 handle = DuplicateHandle(proc, handle, proc, 0, flags & 1, 2)
87 self.handle = handle.Detach()
88 else:
89 raise TypeError("Handle() argument must be a handle, not '%s'" %
90 type(name).__name__)
91
92 def __int__(self):
93 return int(self.handle)
94
95 def Detach(self):
96 handle = self.handle
97 self.handle = None
98 return handle
99
100 # This is also __del__, so only locals are accessed
101 def Close(self, _CloseHandle=windll.kernel32.CloseHandle, _HANDLE=HANDLE):
102 if self.handle:
103 _CloseHandle(_HANDLE(self.handle))
104 self.handle = None
105 __del__ = Close
106
107 CHAR = c_char
108 INVALID_HANDLE_VALUE = HANDLE(-1).value
109 LPDWORD = POINTER(DWORD)
110 LPFILETIME = POINTER(FILETIME)
111 SIZE_T = ULONG_PTR = WPARAM
112 ULONGLONG = c_ulonglong
113
114 try:
115 unicode
116 except NameError:
117 LPCTSTR = LPCWSTR
118 unisuffix = 'W'
119 else:
120 LPCTSTR = LPCSTR
121 unisuffix = 'A'
122
123
124 prototype = WINFUNCTYPE(BOOL, Handle,
125 LPFILETIME, LPFILETIME, LPFILETIME, LPFILETIME)
126 flags = ((1, 'process'),
127 (2, 'creation'), (2, 'exit'), (2, 'kernel'), (2, 'user'))
128 try:
129 GetProcessTimes = prototype(('GetProcessTimes', windll.kernel32), flags)
130 except AttributeError:
131 # Available on NT 3.5 and up, NT line only
132 GetProcessTimes = None
133 else:
134 def errcheck(result, func, args):
135 if not result: raise WinError()
136 ftimes = [t.dwHighDateTime << 32 | t.dwLowDateTime for t in args[3:]]
137 kernel = ftimes[0] / 10000000
138 user = ftimes[1] / 10000000
139 return ProcessTimes(kernel, user)
140 GetProcessTimes.errcheck = errcheck
141
142
143 class PROCESS_MEMORY_COUNTERS(Structure):
144 _fields_ = (('cb', DWORD),
145 ('PageFaultCount', DWORD),
146 ('PeakWorkingSetSize', SIZE_T),
147 ('WorkingSetSize', SIZE_T),
148 ('QuotaPeakPagedPoolUsage', SIZE_T),
149 ('QuotaPagedPoolUsage', SIZE_T),
150 ('QuotaPeakNonPagedPoolUsage', SIZE_T),
151 ('QuotaNonPagedPoolUsage', SIZE_T),
152 ('PagefileUsage', SIZE_T),
153 ('PeakPagefileUsage', SIZE_T))
154
155 prototype = WINFUNCTYPE(BOOL, Handle, POINTER(PROCESS_MEMORY_COUNTERS), DWORD)
156 flags = ((1, 'process'), (2, 'counters'),
157 (5, 'cb', sizeof(PROCESS_MEMORY_COUNTERS)))
158 try:
159 GetProcessMemoryInfo = prototype(('GetProcessMemoryInfo', windll.psapi),
160 flags)
161 except AttributeError:
162 # Available on NT 4.0 and up, NT line only
163 GetProcessMemoryInfo = None
164 else:
165 def errcheck(result, func, args):
166 if not result: raise WinError()
167 return args
168 GetProcessMemoryInfo.errcheck = errcheck
169
170
171 class _uChar_union(Union):
172 _fields_ = (('UnicodeChar', WCHAR),
173 ('AsciiChar', CHAR))
174
175 class KEY_EVENT_RECORD(Structure):
176 _fields_ = (('bKeyDown', BOOL),
177 ('wRepeatCount', WORD),
178 ('wVirtualKeyCode', WORD),
179 ('wVirtualScanCode', WORD),
180 ('uChar', _uChar_union),
181 ('dwControlKeyState', DWORD))
182
183 RIGHT_ALT_PRESSED = 0x001
184 LEFT_ALT_PRESSED = 0x002
185 RIGHT_CTRL_PRESSED = 0x004
186 LEFT_CTRL_PRESSED = 0x008
187 SHIFT_PRESSED = 0x010
188 NUMLOCK_ON = 0x020
189 SCROLLLOCK_ON = 0x040
190 CAPSLOCK_ON = 0x080
191 ENHANCED_KEY = 0x100
192
193 class _Event_union(Union):
194 _fields_ = ('KeyEvent', KEY_EVENT_RECORD),
195
196 class INPUT_RECORD(Structure):
197 _fields_ = (('EventType', WORD),
198 ('Event', _Event_union))
199
200 KEY_EVENT = 0x01
201 MOUSE_EVENT = 0x02
202 WINDOW_BUFFER_SIZE_EVENT = 0x04
203 MENU_EVENT = 0x08
204 FOCUS_EVENT = 0x10
205
206 prototype = WINFUNCTYPE(BOOL, Handle, POINTER(INPUT_RECORD), DWORD, LPDWORD)
207 flags = (1, 'input'), (2, 'buffer'), (5, 'length', 1), (2, 'number_read')
208 ReadConsoleInput = prototype(('ReadConsoleInputA', windll.kernel32), flags)
209 def errcheck(result, func, args):
210 if not result: raise WinError()
211 return args[1] if args[3] else None
212 ReadConsoleInput.errcheck = errcheck
213
214
215 prototype = WINFUNCTYPE(BOOL, Handle)
216 flags = (1, 'input'),
217 FlushConsoleInputBuffer = prototype(('FlushConsoleInputBuffer',
218 windll.kernel32), flags)
219 def errcheck(result, func, args):
220 if not result: raise WinError()
221 FlushConsoleInputBuffer.errcheck = errcheck
222
223
224 prototype = WINFUNCTYPE(BOOL, Handle, DWORD)
225 flags = (1, 'console'), (1, 'mode')
226 SetConsoleMode = prototype(('SetConsoleMode', windll.kernel32), flags)
227 def errcheck(result, func, args):
228 if not result: raise WinError()
229 SetConsoleMode.errcheck = errcheck
230
231 ENABLE_PROCESSED_INPUT = 0x001
232 ENABLE_LINE_INPUT = 0x002
233 ENABLE_ECHO_INPUT = 0x004
234 ENABLE_WINDOW_INPUT = 0x008
235 ENABLE_MOUSE_INPUT = 0x010
236 ENABLE_INSERT_MODE = 0x020
237 ENABLE_QUICK_EDIT_MODE = 0x040
238 ENABLE_EXTENDED_FLAGS = 0x080
239
240 ENABLE_PROCESSED_OUTPUT = 1
241 ENABLE_WRAP_AT_EOL_OUTPUT = 2
242
243
244 prototype = WINFUNCTYPE(HANDLE, c_void_p, LPCTSTR)
245 flags = (5, 'attributes'), (1, 'name')
246 try:
247 CreateJobObject = prototype(('CreateJobObject'+unisuffix, windll.kernel32),
248 flags)
249 except AttributeError:
250 # Available on 2000 and up, NT line only
251 CreateJobObject = lambda name: None
252 else:
253 def errcheck(result, func, args):
254 if not result: raise WinError()
255 return Handle(result)
256 CreateJobObject.errcheck = errcheck
257
258
259 prototype = WINFUNCTYPE(BOOL, Handle, Handle)
260 flags = (1, 'job'), (1, 'handle')
261 try:
262 AssignProcessToJobObject = prototype(('AssignProcessToJobObject',
263 windll.kernel32), flags)
264 except AttributeError:
265 # Available on 2000 and up, NT line only
266 AssignProcessToJobObject = lambda job, handle: None
267 else:
268 def errcheck(result, func, args):
269 if not result: raise WinError()
270 AssignProcessToJobObject.errcheck = errcheck
271
272
273 class JOBOBJECT_BASIC_LIMIT_INFORMATION(Structure):
274 _fields_ = (('PerProcessUserTimeLimit', LARGE_INTEGER),
275 ('PerJobUserTimeLimit', LARGE_INTEGER),
276 ('LimitFlags', DWORD),
277 ('MinimumWorkingSetSize', SIZE_T),
278 ('MaximumWorkingSetSize', SIZE_T),
279 ('ActiveProcessLimit', DWORD),
280 ('Affinity', ULONG_PTR),
281 ('PriorityClass', DWORD),
282 ('SchedulingClass', DWORD))
283
284 JOB_OBJECT_LIMIT_WORKINGSET = 0x0001
285 JOB_OBJECT_LIMIT_PROCESS_TIME = 0x0002
286 JOB_OBJECT_LIMIT_JOB_TIME = 0x0004
287 JOB_OBJECT_LIMIT_ACTIVE_PROCESS = 0x0008
288 JOB_OBJECT_LIMIT_AFFINITY = 0x0010
289 JOB_OBJECT_LIMIT_PRIORITY_CLASS = 0x0020
290 JOB_OBJECT_LIMIT_PRESERVE_JOB_TIME = 0x0040
291 JOB_OBJECT_LIMIT_SCHEDULING_CLASS = 0x0080
292 JOB_OBJECT_LIMIT_PROCESS_MEMORY = 0x0100
293 JOB_OBJECT_LIMIT_JOB_MEMORY = 0x0200
294 JOB_OBJECT_LIMIT_DIE_ON_UNHANDLED_EXCEPTION = 0x0400
295 JOB_OBJECT_LIMIT_BREAKAWAY_OK = 0x0800
296 JOB_OBJECT_LIMIT_SILENT_BREAKAWAY_OK = 0x1000
297 JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE = 0x2000
298 JOB_OBJECT_LIMIT_SUBSET_AFFINITY = 0x4000
299
300 class IO_COUNTERS(Structure):
301 _fields_ = (('ReadOperationCount', ULONGLONG),
302 ('WriteOperationCount', ULONGLONG),
303 ('OtherOperationCount', ULONGLONG),
304 ('ReadTransferCount', ULONGLONG),
305 ('WriteTransferCount', ULONGLONG),
306 ('OtherTransferCount', ULONGLONG))
307
308 class JOBOBJECT_EXTENDED_LIMIT_INFORMATION(Structure):
309 _fields_ = (('BasicLimitInformation', JOBOBJECT_BASIC_LIMIT_INFORMATION),
310 ('IoInfo', IO_COUNTERS),
311 ('ProcessMemoryLimit', SIZE_T),
312 ('JobMemoryLimit', SIZE_T),
313 ('PeakProcessMemoryUsed', SIZE_T),
314 ('PeakJobMemoryUsed', SIZE_T))
315
316 prototype = WINFUNCTYPE(BOOL, HANDLE, c_int, c_void_p, DWORD)
317 flags = (1, 'job'), (1, 'infoclass'), (1, 'info'), (1, 'infosize')
318 try:
319 _setjobinfo = prototype(('SetInformationJobObject',windll.kernel32), flags)
320 except AttributeError:
321 # Available on 2000 and up, NT line only
322 SetInformationJobObject = lambda job, infoclass, info: None
323 else:
324 def errcheck(result, func, args):
325 if not result: raise WinError()
326 _setjobinfo.errcheck = errcheck
327 def SetInformationJobObject(job, infoclass, info):
328 return _setjobinfo(job, infoclass, info, sizeof(info))
329
330 (
331 JobObjectBasicAccountingInformation,
332 JobObjectBasicLimitInformation,
333 JobObjectBasicProcessIdList,
334 JobObjectBasicUIRestrictions,
335 JobObjectSecurityLimitInformation,
336 JobObjectEndOfJobTimeInformation,
337 JobObjectAssociateCompletionPortInformation,
338 JobObjectBasicAndIoAccountingInformation,
339 JobObjectExtendedLimitInformation,
340 JobObjectJobSetInformation,
341 MaxJobObjectInfoClass
342 ) = range(1, 12)
343
344
345 prototype = WINFUNCTYPE(DWORD, DWORD, POINTER(HANDLE), BOOL, DWORD)
346 flags = (1, 'count'), (1, 'handles'), (1, 'wait_all'), (1, 'milliseconds')
347 _wait_multiple = prototype(('WaitForMultipleObjects', windll.kernel32), flags)
348 def errcheck(result, func, args):
349 if result == WAIT_FAILED: raise WinError()
350 return args
351 _wait_multiple.errcheck = errcheck
352 def WaitForMultipleObjects(handles, wait_all, timeout):
353 n = len(handles)
354 handles = (Handle.from_param(handle) for handle in handles)
355 timeout = ceil(timeout * 1000)
356 return _wait_multiple(n, (HANDLE * n)(*handles), wait_all, timeout)
357
358 # WAIT_OBJECT_0 defined at the top of the file
359 WAIT_ABANDONED_0 = 0x00000080
360 WAIT_TIMEOUT = 0x00000102
361 WAIT_FAILED = 0xFFFFFFFF
362
363
364 try:
365 _wait_single = WaitForSingleObject
366 except NameError:
367 prototype = WINFUNCTYPE(DWORD, Handle, DWORD)
368 flags = (1, 'handle'), (1, 'milliseconds')
369 _wait_single = prototype(('WaitForSingleObject', windll.kernel32), flags)
370 def errcheck(result, func, args):
371 if result == WAIT_FAILED: raise WinError()
372 return args
373 _wait_single.errcheck = errcheck
374 def WaitForSingleObject(handle, timeout):
375 return _wait_single(handle, ceil(timeout * 1000))
376
377
378 try:
379 GetStdHandle
380 except NameError:
381 prototype = WINFUNCTYPE(HANDLE, DWORD)
382 flags = (1, 'which'),
383 GetStdHandle = prototype(('GetStdHandle', windll.kernel32), flags)
384 def errcheck(result, func, args):
385 if result == INVALID_HANDLE_VALUE: raise WinError()
386 return args if result else None
387 GetStdHandle.errcheck = errcheck
388
389
390 try:
391 TerminateProcess
392 except NameError:
393 prototype = WINFUNCTYPE(BOOL, Handle, UINT)
394 flags = (1, 'process'), (1, 'exitcode')
395 TerminateProcess = prototype(('TerminateProcess', windll.kernel32), flags)
396 def errcheck(result, func, args):
397 if not result: raise WinError()
398 TerminateProcess.errcheck = errcheck
399
400
401 # Do not show error messages due to errors in the program being tested
402 try:
403 errmode = ctypes.windll.kernel32.GetErrorMode()
404 except AttributeError:
405 # GetErrorMode is available on Vista/2008 and up
406 errmode = ctypes.windll.kernel32.SetErrorMode(0)
407 ctypes.windll.kernel32.SetErrorMode(errmode | 0x8003)
408
409 stdin = GetStdHandle(STD_INPUT_HANDLE)
410 try:
411 SetConsoleMode(stdin, ENABLE_PROCESSED_INPUT)
412 except WindowsError:
413 console_input = False
414 else:
415 console_input = True
416 FlushConsoleInputBuffer(stdin)
417
418 def kill(process):
419 try:
420 process.terminate()
421 except AttributeError:
422 TerminateProcess(process._handle)
423 terminate = kill
424
425 def call(*args, **kwargs):
426 case = kwargs.pop('case')
427 job = CreateJobObject(None)
428 flags = 0
429 if case.maxcputime:
430 flags |= JOB_OBJECT_LIMIT_PROCESS_TIME
431 if case.maxmemory:
432 flags |= JOB_OBJECT_LIMIT_PROCESS_MEMORY
433 limits = JOBOBJECT_EXTENDED_LIMIT_INFORMATION(
434 JOBOBJECT_BASIC_LIMIT_INFORMATION(
435 PerProcessUserTimeLimit=ceil((case.maxcputime or 0)*10000000),
436 LimitFlags=flags,
437 ),
438 ProcessMemoryLimit=ceil((case.maxmemory or 0)*1048576),
439 )
440 SetInformationJobObject(job, JobObjectExtendedLimitInformation, limits)
441 try:
442 case.process = Popen(*args, **kwargs)
443 except OSError:
444 raise CannotStartTestee(sys.exc_info()[1])
445 case.time_started = clock()
446 AssignProcessToJobObject(job, case.process._handle)
447 if not console_input:
448 if case.maxwalltime:
449 if (WaitForSingleObject(case.process._handle, case.maxwalltime) !=
450 WAIT_OBJECT_0):
451 raise TimeLimitExceeded
452 else:
453 case.process.wait()
454 else:
455 handles = stdin, case.process._handle
456 if case.maxwalltime:
457 time_end = clock() + case.maxwalltime
458 while case.process.poll() is None:
459 remaining = time_end - clock()
460 if remaining > 0:
461 if (WaitForMultipleObjects(handles, False, remaining) ==
462 WAIT_OBJECT_0):
463 ir = ReadConsoleInput(stdin)
464 if (ir and
465 ir.EventType == 1 and
466 ir.Event.KeyEvent.bKeyDown and
467 ir.Event.KeyEvent.wVirtualKeyCode == 27):
468 raise CanceledByUser
469 else:
470 raise TimeLimitExceeded
471 else:
472 while case.process.poll() is None:
473 if (WaitForMultipleObjects(handles, False, INFINITE) ==
474 WAIT_OBJECT_0):
475 ir = ReadConsoleInput(stdin)
476 if (ir and
477 ir.EventType == 1 and
478 ir.Event.KeyEvent.bKeyDown and
479 ir.Event.KeyEvent.wVirtualKeyCode == 27):
480 raise CanceledByUser
481 case.time_stopped = clock()
482 if case.maxcputime and GetProcessTimes:
483 try:
484 times = GetProcessTimes(case.process._handle)
485 except WindowsError:
486 pass
487 else:
488 if times.kernel + times.user > case.maxcputime:
489 raise TimeLimitExceeded
490 if case.maxmemory and GetProcessMemoryInfo:
491 try:
492 counters = GetProcessMemoryInfo(case.process._handle)
493 except WindowsError:
494 pass
495 else:
496 if counters.PeakPagefileUsage > case.maxmemory * 1048576:
497 raise MemoryLimitExceeded