| Benjamin Peterson | e711caf | 2008-06-11 16:44:04 +0000 | [diff] [blame] | 1 | # | 
 | 2 | # A higher level module for using sockets (or Windows named pipes) | 
 | 3 | # | 
 | 4 | # multiprocessing/connection.py | 
 | 5 | # | 
| R. David Murray | 3fc969a | 2010-12-14 01:38:16 +0000 | [diff] [blame] | 6 | # Copyright (c) 2006-2008, R Oudkerk | 
| Richard Oudkerk | 3e268aa | 2012-04-30 12:13:55 +0100 | [diff] [blame] | 7 | # Licensed to PSF under a Contributor Agreement. | 
| Benjamin Peterson | e711caf | 2008-06-11 16:44:04 +0000 | [diff] [blame] | 8 | # | 
 | 9 |  | 
| Antoine Pitrou | bdb1cf1 | 2012-03-05 19:28:37 +0100 | [diff] [blame] | 10 | __all__ = [ 'Client', 'Listener', 'Pipe', 'wait' ] | 
| Benjamin Peterson | e711caf | 2008-06-11 16:44:04 +0000 | [diff] [blame] | 11 |  | 
| Antoine Pitrou | 87cf220 | 2011-05-09 17:04:27 +0200 | [diff] [blame] | 12 | import io | 
| Benjamin Peterson | e711caf | 2008-06-11 16:44:04 +0000 | [diff] [blame] | 13 | import os | 
 | 14 | import sys | 
| Antoine Pitrou | 87cf220 | 2011-05-09 17:04:27 +0200 | [diff] [blame] | 15 | import pickle | 
 | 16 | import select | 
| Benjamin Peterson | e711caf | 2008-06-11 16:44:04 +0000 | [diff] [blame] | 17 | import socket | 
| Antoine Pitrou | 87cf220 | 2011-05-09 17:04:27 +0200 | [diff] [blame] | 18 | import struct | 
| Georg Brandl | 6aa2d1f | 2008-08-12 08:35:52 +0000 | [diff] [blame] | 19 | import errno | 
| Benjamin Peterson | e711caf | 2008-06-11 16:44:04 +0000 | [diff] [blame] | 20 | import time | 
 | 21 | import tempfile | 
 | 22 | import itertools | 
 | 23 |  | 
 | 24 | import _multiprocessing | 
| Antoine Pitrou | 87cf220 | 2011-05-09 17:04:27 +0200 | [diff] [blame] | 25 | from multiprocessing import current_process, AuthenticationError, BufferTooShort | 
| Richard Oudkerk | 59d5404 | 2012-05-10 16:11:12 +0100 | [diff] [blame^] | 26 | from multiprocessing.util import get_temp_dir, Finalize, sub_debug, debug | 
| Antoine Pitrou | 5438ed1 | 2012-04-24 22:56:57 +0200 | [diff] [blame] | 27 | from multiprocessing.forking import ForkingPickler | 
| Antoine Pitrou | 87cf220 | 2011-05-09 17:04:27 +0200 | [diff] [blame] | 28 | try: | 
| Antoine Pitrou | 23bba4c | 2012-04-18 20:51:15 +0200 | [diff] [blame] | 29 |     import _winapi | 
 | 30 |     from _winapi import WAIT_OBJECT_0, WAIT_TIMEOUT, INFINITE | 
| Antoine Pitrou | 87cf220 | 2011-05-09 17:04:27 +0200 | [diff] [blame] | 31 | except ImportError: | 
 | 32 |     if sys.platform == 'win32': | 
 | 33 |         raise | 
| Antoine Pitrou | 23bba4c | 2012-04-18 20:51:15 +0200 | [diff] [blame] | 34 |     _winapi = None | 
| Benjamin Peterson | e711caf | 2008-06-11 16:44:04 +0000 | [diff] [blame] | 35 |  | 
 | 36 | # | 
 | 37 | # | 
 | 38 | # | 
 | 39 |  | 
 | 40 | BUFSIZE = 8192 | 
| Antoine Pitrou | 45d61a3 | 2009-11-13 22:35:18 +0000 | [diff] [blame] | 41 | # A very generous timeout when it comes to local connections... | 
 | 42 | CONNECTION_TIMEOUT = 20. | 
| Benjamin Peterson | e711caf | 2008-06-11 16:44:04 +0000 | [diff] [blame] | 43 |  | 
 | 44 | _mmap_counter = itertools.count() | 
 | 45 |  | 
 | 46 | default_family = 'AF_INET' | 
 | 47 | families = ['AF_INET'] | 
 | 48 |  | 
 | 49 | if hasattr(socket, 'AF_UNIX'): | 
 | 50 |     default_family = 'AF_UNIX' | 
 | 51 |     families += ['AF_UNIX'] | 
 | 52 |  | 
 | 53 | if sys.platform == 'win32': | 
 | 54 |     default_family = 'AF_PIPE' | 
 | 55 |     families += ['AF_PIPE'] | 
 | 56 |  | 
| Antoine Pitrou | 45d61a3 | 2009-11-13 22:35:18 +0000 | [diff] [blame] | 57 |  | 
 | 58 | def _init_timeout(timeout=CONNECTION_TIMEOUT): | 
 | 59 |     return time.time() + timeout | 
 | 60 |  | 
 | 61 | def _check_timeout(t): | 
 | 62 |     return time.time() > t | 
 | 63 |  | 
| Benjamin Peterson | e711caf | 2008-06-11 16:44:04 +0000 | [diff] [blame] | 64 | # | 
 | 65 | # | 
 | 66 | # | 
 | 67 |  | 
 | 68 | def arbitrary_address(family): | 
 | 69 |     ''' | 
 | 70 |     Return an arbitrary free address for the given family | 
 | 71 |     ''' | 
 | 72 |     if family == 'AF_INET': | 
 | 73 |         return ('localhost', 0) | 
 | 74 |     elif family == 'AF_UNIX': | 
 | 75 |         return tempfile.mktemp(prefix='listener-', dir=get_temp_dir()) | 
 | 76 |     elif family == 'AF_PIPE': | 
 | 77 |         return tempfile.mktemp(prefix=r'\\.\pipe\pyc-%d-%d-' % | 
 | 78 |                                (os.getpid(), next(_mmap_counter))) | 
 | 79 |     else: | 
 | 80 |         raise ValueError('unrecognized family') | 
 | 81 |  | 
| Antoine Pitrou | 709176f | 2012-04-01 17:19:09 +0200 | [diff] [blame] | 82 | def _validate_family(family): | 
 | 83 |     ''' | 
 | 84 |     Checks if the family is valid for the current environment. | 
 | 85 |     ''' | 
 | 86 |     if sys.platform != 'win32' and family == 'AF_PIPE': | 
 | 87 |         raise ValueError('Family %s is not recognized.' % family) | 
 | 88 |  | 
| Antoine Pitrou | 6d20cba | 2012-04-03 20:12:23 +0200 | [diff] [blame] | 89 |     if sys.platform == 'win32' and family == 'AF_UNIX': | 
 | 90 |         # double check | 
 | 91 |         if not hasattr(socket, family): | 
 | 92 |             raise ValueError('Family %s is not recognized.' % family) | 
| Benjamin Peterson | e711caf | 2008-06-11 16:44:04 +0000 | [diff] [blame] | 93 |  | 
 | 94 | def address_type(address): | 
 | 95 |     ''' | 
 | 96 |     Return the types of the address | 
 | 97 |  | 
 | 98 |     This can be 'AF_INET', 'AF_UNIX', or 'AF_PIPE' | 
 | 99 |     ''' | 
 | 100 |     if type(address) == tuple: | 
 | 101 |         return 'AF_INET' | 
 | 102 |     elif type(address) is str and address.startswith('\\\\'): | 
 | 103 |         return 'AF_PIPE' | 
 | 104 |     elif type(address) is str: | 
 | 105 |         return 'AF_UNIX' | 
 | 106 |     else: | 
 | 107 |         raise ValueError('address type of %r unrecognized' % address) | 
 | 108 |  | 
 | 109 | # | 
| Antoine Pitrou | 87cf220 | 2011-05-09 17:04:27 +0200 | [diff] [blame] | 110 | # Connection classes | 
 | 111 | # | 
 | 112 |  | 
 | 113 | class _ConnectionBase: | 
 | 114 |     _handle = None | 
 | 115 |  | 
 | 116 |     def __init__(self, handle, readable=True, writable=True): | 
 | 117 |         handle = handle.__index__() | 
 | 118 |         if handle < 0: | 
 | 119 |             raise ValueError("invalid handle") | 
 | 120 |         if not readable and not writable: | 
 | 121 |             raise ValueError( | 
 | 122 |                 "at least one of `readable` and `writable` must be True") | 
 | 123 |         self._handle = handle | 
 | 124 |         self._readable = readable | 
 | 125 |         self._writable = writable | 
 | 126 |  | 
| Antoine Pitrou | 6000120 | 2011-07-09 01:03:46 +0200 | [diff] [blame] | 127 |     # XXX should we use util.Finalize instead of a __del__? | 
 | 128 |  | 
| Antoine Pitrou | 87cf220 | 2011-05-09 17:04:27 +0200 | [diff] [blame] | 129 |     def __del__(self): | 
 | 130 |         if self._handle is not None: | 
 | 131 |             self._close() | 
 | 132 |  | 
 | 133 |     def _check_closed(self): | 
 | 134 |         if self._handle is None: | 
 | 135 |             raise IOError("handle is closed") | 
 | 136 |  | 
 | 137 |     def _check_readable(self): | 
 | 138 |         if not self._readable: | 
 | 139 |             raise IOError("connection is write-only") | 
 | 140 |  | 
 | 141 |     def _check_writable(self): | 
 | 142 |         if not self._writable: | 
 | 143 |             raise IOError("connection is read-only") | 
 | 144 |  | 
 | 145 |     def _bad_message_length(self): | 
 | 146 |         if self._writable: | 
 | 147 |             self._readable = False | 
 | 148 |         else: | 
 | 149 |             self.close() | 
 | 150 |         raise IOError("bad message length") | 
 | 151 |  | 
 | 152 |     @property | 
 | 153 |     def closed(self): | 
 | 154 |         """True if the connection is closed""" | 
 | 155 |         return self._handle is None | 
 | 156 |  | 
 | 157 |     @property | 
 | 158 |     def readable(self): | 
 | 159 |         """True if the connection is readable""" | 
 | 160 |         return self._readable | 
 | 161 |  | 
 | 162 |     @property | 
 | 163 |     def writable(self): | 
 | 164 |         """True if the connection is writable""" | 
 | 165 |         return self._writable | 
 | 166 |  | 
 | 167 |     def fileno(self): | 
 | 168 |         """File descriptor or handle of the connection""" | 
 | 169 |         self._check_closed() | 
 | 170 |         return self._handle | 
 | 171 |  | 
 | 172 |     def close(self): | 
 | 173 |         """Close the connection""" | 
 | 174 |         if self._handle is not None: | 
 | 175 |             try: | 
 | 176 |                 self._close() | 
 | 177 |             finally: | 
 | 178 |                 self._handle = None | 
 | 179 |  | 
 | 180 |     def send_bytes(self, buf, offset=0, size=None): | 
 | 181 |         """Send the bytes data from a bytes-like object""" | 
 | 182 |         self._check_closed() | 
 | 183 |         self._check_writable() | 
 | 184 |         m = memoryview(buf) | 
 | 185 |         # HACK for byte-indexing of non-bytewise buffers (e.g. array.array) | 
 | 186 |         if m.itemsize > 1: | 
 | 187 |             m = memoryview(bytes(m)) | 
 | 188 |         n = len(m) | 
 | 189 |         if offset < 0: | 
 | 190 |             raise ValueError("offset is negative") | 
 | 191 |         if n < offset: | 
 | 192 |             raise ValueError("buffer length < offset") | 
 | 193 |         if size is None: | 
 | 194 |             size = n - offset | 
 | 195 |         elif size < 0: | 
 | 196 |             raise ValueError("size is negative") | 
 | 197 |         elif offset + size > n: | 
 | 198 |             raise ValueError("buffer length < offset + size") | 
 | 199 |         self._send_bytes(m[offset:offset + size]) | 
 | 200 |  | 
 | 201 |     def send(self, obj): | 
 | 202 |         """Send a (picklable) object""" | 
 | 203 |         self._check_closed() | 
 | 204 |         self._check_writable() | 
| Antoine Pitrou | 5438ed1 | 2012-04-24 22:56:57 +0200 | [diff] [blame] | 205 |         buf = io.BytesIO() | 
 | 206 |         ForkingPickler(buf, pickle.HIGHEST_PROTOCOL).dump(obj) | 
 | 207 |         self._send_bytes(buf.getbuffer()) | 
| Antoine Pitrou | 87cf220 | 2011-05-09 17:04:27 +0200 | [diff] [blame] | 208 |  | 
 | 209 |     def recv_bytes(self, maxlength=None): | 
 | 210 |         """ | 
 | 211 |         Receive bytes data as a bytes object. | 
 | 212 |         """ | 
 | 213 |         self._check_closed() | 
 | 214 |         self._check_readable() | 
 | 215 |         if maxlength is not None and maxlength < 0: | 
 | 216 |             raise ValueError("negative maxlength") | 
 | 217 |         buf = self._recv_bytes(maxlength) | 
 | 218 |         if buf is None: | 
 | 219 |             self._bad_message_length() | 
 | 220 |         return buf.getvalue() | 
 | 221 |  | 
 | 222 |     def recv_bytes_into(self, buf, offset=0): | 
 | 223 |         """ | 
 | 224 |         Receive bytes data into a writeable buffer-like object. | 
 | 225 |         Return the number of bytes read. | 
 | 226 |         """ | 
 | 227 |         self._check_closed() | 
 | 228 |         self._check_readable() | 
 | 229 |         with memoryview(buf) as m: | 
 | 230 |             # Get bytesize of arbitrary buffer | 
 | 231 |             itemsize = m.itemsize | 
 | 232 |             bytesize = itemsize * len(m) | 
 | 233 |             if offset < 0: | 
 | 234 |                 raise ValueError("negative offset") | 
 | 235 |             elif offset > bytesize: | 
 | 236 |                 raise ValueError("offset too large") | 
 | 237 |             result = self._recv_bytes() | 
 | 238 |             size = result.tell() | 
 | 239 |             if bytesize < offset + size: | 
 | 240 |                 raise BufferTooShort(result.getvalue()) | 
 | 241 |             # Message can fit in dest | 
 | 242 |             result.seek(0) | 
 | 243 |             result.readinto(m[offset // itemsize : | 
 | 244 |                               (offset + size) // itemsize]) | 
 | 245 |             return size | 
 | 246 |  | 
| Antoine Pitrou | bdb1cf1 | 2012-03-05 19:28:37 +0100 | [diff] [blame] | 247 |     def recv(self): | 
| Antoine Pitrou | 87cf220 | 2011-05-09 17:04:27 +0200 | [diff] [blame] | 248 |         """Receive a (picklable) object""" | 
 | 249 |         self._check_closed() | 
 | 250 |         self._check_readable() | 
| Antoine Pitrou | bdb1cf1 | 2012-03-05 19:28:37 +0100 | [diff] [blame] | 251 |         buf = self._recv_bytes() | 
| Antoine Pitrou | 87cf220 | 2011-05-09 17:04:27 +0200 | [diff] [blame] | 252 |         return pickle.loads(buf.getbuffer()) | 
 | 253 |  | 
 | 254 |     def poll(self, timeout=0.0): | 
 | 255 |         """Whether there is any input available to be read""" | 
 | 256 |         self._check_closed() | 
 | 257 |         self._check_readable() | 
| Antoine Pitrou | 87cf220 | 2011-05-09 17:04:27 +0200 | [diff] [blame] | 258 |         return self._poll(timeout) | 
 | 259 |  | 
 | 260 |  | 
| Antoine Pitrou | 23bba4c | 2012-04-18 20:51:15 +0200 | [diff] [blame] | 261 | if _winapi: | 
| Antoine Pitrou | 87cf220 | 2011-05-09 17:04:27 +0200 | [diff] [blame] | 262 |  | 
 | 263 |     class PipeConnection(_ConnectionBase): | 
 | 264 |         """ | 
 | 265 |         Connection class based on a Windows named pipe. | 
| Antoine Pitrou | dd69649 | 2011-06-08 17:21:55 +0200 | [diff] [blame] | 266 |         Overlapped I/O is used, so the handles must have been created | 
 | 267 |         with FILE_FLAG_OVERLAPPED. | 
| Antoine Pitrou | 87cf220 | 2011-05-09 17:04:27 +0200 | [diff] [blame] | 268 |         """ | 
| Antoine Pitrou | bdb1cf1 | 2012-03-05 19:28:37 +0100 | [diff] [blame] | 269 |         _got_empty_message = False | 
| Antoine Pitrou | 87cf220 | 2011-05-09 17:04:27 +0200 | [diff] [blame] | 270 |  | 
| Antoine Pitrou | 23bba4c | 2012-04-18 20:51:15 +0200 | [diff] [blame] | 271 |         def _close(self, _CloseHandle=_winapi.CloseHandle): | 
| Antoine Pitrou | 71a28a9 | 2011-07-09 01:03:00 +0200 | [diff] [blame] | 272 |             _CloseHandle(self._handle) | 
| Antoine Pitrou | 87cf220 | 2011-05-09 17:04:27 +0200 | [diff] [blame] | 273 |  | 
 | 274 |         def _send_bytes(self, buf): | 
| Antoine Pitrou | 23bba4c | 2012-04-18 20:51:15 +0200 | [diff] [blame] | 275 |             ov, err = _winapi.WriteFile(self._handle, buf, overlapped=True) | 
| Antoine Pitrou | bdb1cf1 | 2012-03-05 19:28:37 +0100 | [diff] [blame] | 276 |             try: | 
| Antoine Pitrou | 23bba4c | 2012-04-18 20:51:15 +0200 | [diff] [blame] | 277 |                 if err == _winapi.ERROR_IO_PENDING: | 
 | 278 |                     waitres = _winapi.WaitForMultipleObjects( | 
| Antoine Pitrou | bdb1cf1 | 2012-03-05 19:28:37 +0100 | [diff] [blame] | 279 |                         [ov.event], False, INFINITE) | 
 | 280 |                     assert waitres == WAIT_OBJECT_0 | 
 | 281 |             except: | 
 | 282 |                 ov.cancel() | 
 | 283 |                 raise | 
 | 284 |             finally: | 
 | 285 |                 nwritten, err = ov.GetOverlappedResult(True) | 
 | 286 |             assert err == 0 | 
| Antoine Pitrou | 87cf220 | 2011-05-09 17:04:27 +0200 | [diff] [blame] | 287 |             assert nwritten == len(buf) | 
 | 288 |  | 
| Antoine Pitrou | bdb1cf1 | 2012-03-05 19:28:37 +0100 | [diff] [blame] | 289 |         def _recv_bytes(self, maxsize=None): | 
 | 290 |             if self._got_empty_message: | 
 | 291 |                 self._got_empty_message = False | 
 | 292 |                 return io.BytesIO() | 
| Antoine Pitrou | dd69649 | 2011-06-08 17:21:55 +0200 | [diff] [blame] | 293 |             else: | 
| Antoine Pitrou | bdb1cf1 | 2012-03-05 19:28:37 +0100 | [diff] [blame] | 294 |                 bsize = 128 if maxsize is None else min(maxsize, 128) | 
| Antoine Pitrou | dd69649 | 2011-06-08 17:21:55 +0200 | [diff] [blame] | 295 |                 try: | 
| Antoine Pitrou | 23bba4c | 2012-04-18 20:51:15 +0200 | [diff] [blame] | 296 |                     ov, err = _winapi.ReadFile(self._handle, bsize, | 
 | 297 |                                                 overlapped=True) | 
| Antoine Pitrou | bdb1cf1 | 2012-03-05 19:28:37 +0100 | [diff] [blame] | 298 |                     try: | 
| Antoine Pitrou | 23bba4c | 2012-04-18 20:51:15 +0200 | [diff] [blame] | 299 |                         if err == _winapi.ERROR_IO_PENDING: | 
 | 300 |                             waitres = _winapi.WaitForMultipleObjects( | 
| Antoine Pitrou | bdb1cf1 | 2012-03-05 19:28:37 +0100 | [diff] [blame] | 301 |                                 [ov.event], False, INFINITE) | 
 | 302 |                             assert waitres == WAIT_OBJECT_0 | 
 | 303 |                     except: | 
 | 304 |                         ov.cancel() | 
 | 305 |                         raise | 
 | 306 |                     finally: | 
 | 307 |                         nread, err = ov.GetOverlappedResult(True) | 
 | 308 |                         if err == 0: | 
 | 309 |                             f = io.BytesIO() | 
 | 310 |                             f.write(ov.getbuffer()) | 
 | 311 |                             return f | 
| Antoine Pitrou | 23bba4c | 2012-04-18 20:51:15 +0200 | [diff] [blame] | 312 |                         elif err == _winapi.ERROR_MORE_DATA: | 
| Antoine Pitrou | bdb1cf1 | 2012-03-05 19:28:37 +0100 | [diff] [blame] | 313 |                             return self._get_more_data(ov, maxsize) | 
| Antoine Pitrou | dd69649 | 2011-06-08 17:21:55 +0200 | [diff] [blame] | 314 |                 except IOError as e: | 
| Antoine Pitrou | 23bba4c | 2012-04-18 20:51:15 +0200 | [diff] [blame] | 315 |                     if e.winerror == _winapi.ERROR_BROKEN_PIPE: | 
| Antoine Pitrou | dd69649 | 2011-06-08 17:21:55 +0200 | [diff] [blame] | 316 |                         raise EOFError | 
| Antoine Pitrou | bdb1cf1 | 2012-03-05 19:28:37 +0100 | [diff] [blame] | 317 |                     else: | 
 | 318 |                         raise | 
 | 319 |             raise RuntimeError("shouldn't get here; expected KeyboardInterrupt") | 
| Antoine Pitrou | 87cf220 | 2011-05-09 17:04:27 +0200 | [diff] [blame] | 320 |  | 
| Antoine Pitrou | bdb1cf1 | 2012-03-05 19:28:37 +0100 | [diff] [blame] | 321 |         def _poll(self, timeout): | 
 | 322 |             if (self._got_empty_message or | 
| Antoine Pitrou | 23bba4c | 2012-04-18 20:51:15 +0200 | [diff] [blame] | 323 |                         _winapi.PeekNamedPipe(self._handle)[0] != 0): | 
| Antoine Pitrou | 87cf220 | 2011-05-09 17:04:27 +0200 | [diff] [blame] | 324 |                 return True | 
| Antoine Pitrou | bdb1cf1 | 2012-03-05 19:28:37 +0100 | [diff] [blame] | 325 |             return bool(wait([self], timeout)) | 
 | 326 |  | 
 | 327 |         def _get_more_data(self, ov, maxsize): | 
 | 328 |             buf = ov.getbuffer() | 
 | 329 |             f = io.BytesIO() | 
 | 330 |             f.write(buf) | 
| Antoine Pitrou | 23bba4c | 2012-04-18 20:51:15 +0200 | [diff] [blame] | 331 |             left = _winapi.PeekNamedPipe(self._handle)[1] | 
| Antoine Pitrou | bdb1cf1 | 2012-03-05 19:28:37 +0100 | [diff] [blame] | 332 |             assert left > 0 | 
 | 333 |             if maxsize is not None and len(buf) + left > maxsize: | 
 | 334 |                 self._bad_message_length() | 
| Antoine Pitrou | 23bba4c | 2012-04-18 20:51:15 +0200 | [diff] [blame] | 335 |             ov, err = _winapi.ReadFile(self._handle, left, overlapped=True) | 
| Antoine Pitrou | bdb1cf1 | 2012-03-05 19:28:37 +0100 | [diff] [blame] | 336 |             rbytes, err = ov.GetOverlappedResult(True) | 
 | 337 |             assert err == 0 | 
 | 338 |             assert rbytes == left | 
 | 339 |             f.write(ov.getbuffer()) | 
 | 340 |             return f | 
| Antoine Pitrou | 87cf220 | 2011-05-09 17:04:27 +0200 | [diff] [blame] | 341 |  | 
 | 342 |  | 
 | 343 | class Connection(_ConnectionBase): | 
 | 344 |     """ | 
 | 345 |     Connection class based on an arbitrary file descriptor (Unix only), or | 
 | 346 |     a socket handle (Windows). | 
 | 347 |     """ | 
 | 348 |  | 
| Antoine Pitrou | 23bba4c | 2012-04-18 20:51:15 +0200 | [diff] [blame] | 349 |     if _winapi: | 
 | 350 |         def _close(self, _close=_multiprocessing.closesocket): | 
| Antoine Pitrou | 71a28a9 | 2011-07-09 01:03:00 +0200 | [diff] [blame] | 351 |             _close(self._handle) | 
| Antoine Pitrou | 23bba4c | 2012-04-18 20:51:15 +0200 | [diff] [blame] | 352 |         _write = _multiprocessing.send | 
 | 353 |         _read = _multiprocessing.recv | 
| Antoine Pitrou | 87cf220 | 2011-05-09 17:04:27 +0200 | [diff] [blame] | 354 |     else: | 
| Antoine Pitrou | 71a28a9 | 2011-07-09 01:03:00 +0200 | [diff] [blame] | 355 |         def _close(self, _close=os.close): | 
 | 356 |             _close(self._handle) | 
| Antoine Pitrou | 87cf220 | 2011-05-09 17:04:27 +0200 | [diff] [blame] | 357 |         _write = os.write | 
 | 358 |         _read = os.read | 
 | 359 |  | 
 | 360 |     def _send(self, buf, write=_write): | 
 | 361 |         remaining = len(buf) | 
 | 362 |         while True: | 
 | 363 |             n = write(self._handle, buf) | 
 | 364 |             remaining -= n | 
 | 365 |             if remaining == 0: | 
 | 366 |                 break | 
 | 367 |             buf = buf[n:] | 
 | 368 |  | 
| Antoine Pitrou | bdb1cf1 | 2012-03-05 19:28:37 +0100 | [diff] [blame] | 369 |     def _recv(self, size, read=_read): | 
| Antoine Pitrou | 87cf220 | 2011-05-09 17:04:27 +0200 | [diff] [blame] | 370 |         buf = io.BytesIO() | 
| Antoine Pitrou | dd69649 | 2011-06-08 17:21:55 +0200 | [diff] [blame] | 371 |         handle = self._handle | 
| Antoine Pitrou | 87cf220 | 2011-05-09 17:04:27 +0200 | [diff] [blame] | 372 |         remaining = size | 
 | 373 |         while remaining > 0: | 
| Antoine Pitrou | dd69649 | 2011-06-08 17:21:55 +0200 | [diff] [blame] | 374 |             chunk = read(handle, remaining) | 
| Antoine Pitrou | 87cf220 | 2011-05-09 17:04:27 +0200 | [diff] [blame] | 375 |             n = len(chunk) | 
 | 376 |             if n == 0: | 
 | 377 |                 if remaining == size: | 
 | 378 |                     raise EOFError | 
 | 379 |                 else: | 
 | 380 |                     raise IOError("got end of file during message") | 
 | 381 |             buf.write(chunk) | 
 | 382 |             remaining -= n | 
 | 383 |         return buf | 
 | 384 |  | 
 | 385 |     def _send_bytes(self, buf): | 
 | 386 |         # For wire compatibility with 3.2 and lower | 
 | 387 |         n = len(buf) | 
| Charles-François Natali | 225aa4f | 2011-09-20 19:27:39 +0200 | [diff] [blame] | 388 |         self._send(struct.pack("!i", n)) | 
| Antoine Pitrou | 87cf220 | 2011-05-09 17:04:27 +0200 | [diff] [blame] | 389 |         # The condition is necessary to avoid "broken pipe" errors | 
 | 390 |         # when sending a 0-length buffer if the other end closed the pipe. | 
 | 391 |         if n > 0: | 
 | 392 |             self._send(buf) | 
 | 393 |  | 
| Antoine Pitrou | bdb1cf1 | 2012-03-05 19:28:37 +0100 | [diff] [blame] | 394 |     def _recv_bytes(self, maxsize=None): | 
 | 395 |         buf = self._recv(4) | 
| Charles-François Natali | 225aa4f | 2011-09-20 19:27:39 +0200 | [diff] [blame] | 396 |         size, = struct.unpack("!i", buf.getvalue()) | 
| Antoine Pitrou | 87cf220 | 2011-05-09 17:04:27 +0200 | [diff] [blame] | 397 |         if maxsize is not None and size > maxsize: | 
 | 398 |             return None | 
| Antoine Pitrou | bdb1cf1 | 2012-03-05 19:28:37 +0100 | [diff] [blame] | 399 |         return self._recv(size) | 
| Antoine Pitrou | 87cf220 | 2011-05-09 17:04:27 +0200 | [diff] [blame] | 400 |  | 
 | 401 |     def _poll(self, timeout): | 
| Antoine Pitrou | bdb1cf1 | 2012-03-05 19:28:37 +0100 | [diff] [blame] | 402 |         r = wait([self._handle], timeout) | 
| Antoine Pitrou | 87cf220 | 2011-05-09 17:04:27 +0200 | [diff] [blame] | 403 |         return bool(r) | 
 | 404 |  | 
 | 405 |  | 
 | 406 | # | 
| Benjamin Peterson | e711caf | 2008-06-11 16:44:04 +0000 | [diff] [blame] | 407 | # Public functions | 
 | 408 | # | 
 | 409 |  | 
 | 410 | class Listener(object): | 
 | 411 |     ''' | 
 | 412 |     Returns a listener object. | 
 | 413 |  | 
 | 414 |     This is a wrapper for a bound socket which is 'listening' for | 
 | 415 |     connections, or for a Windows named pipe. | 
 | 416 |     ''' | 
 | 417 |     def __init__(self, address=None, family=None, backlog=1, authkey=None): | 
 | 418 |         family = family or (address and address_type(address)) \ | 
 | 419 |                  or default_family | 
 | 420 |         address = address or arbitrary_address(family) | 
 | 421 |  | 
| Antoine Pitrou | 709176f | 2012-04-01 17:19:09 +0200 | [diff] [blame] | 422 |         _validate_family(family) | 
| Benjamin Peterson | e711caf | 2008-06-11 16:44:04 +0000 | [diff] [blame] | 423 |         if family == 'AF_PIPE': | 
 | 424 |             self._listener = PipeListener(address, backlog) | 
 | 425 |         else: | 
 | 426 |             self._listener = SocketListener(address, family, backlog) | 
 | 427 |  | 
 | 428 |         if authkey is not None and not isinstance(authkey, bytes): | 
 | 429 |             raise TypeError('authkey should be a byte string') | 
 | 430 |  | 
 | 431 |         self._authkey = authkey | 
 | 432 |  | 
 | 433 |     def accept(self): | 
 | 434 |         ''' | 
 | 435 |         Accept a connection on the bound socket or named pipe of `self`. | 
 | 436 |  | 
 | 437 |         Returns a `Connection` object. | 
 | 438 |         ''' | 
 | 439 |         c = self._listener.accept() | 
 | 440 |         if self._authkey: | 
 | 441 |             deliver_challenge(c, self._authkey) | 
 | 442 |             answer_challenge(c, self._authkey) | 
 | 443 |         return c | 
 | 444 |  | 
 | 445 |     def close(self): | 
 | 446 |         ''' | 
 | 447 |         Close the bound socket or named pipe of `self`. | 
 | 448 |         ''' | 
 | 449 |         return self._listener.close() | 
 | 450 |  | 
 | 451 |     address = property(lambda self: self._listener._address) | 
 | 452 |     last_accepted = property(lambda self: self._listener._last_accepted) | 
 | 453 |  | 
 | 454 |  | 
 | 455 | def Client(address, family=None, authkey=None): | 
 | 456 |     ''' | 
 | 457 |     Returns a connection to the address of a `Listener` | 
 | 458 |     ''' | 
 | 459 |     family = family or address_type(address) | 
| Antoine Pitrou | 709176f | 2012-04-01 17:19:09 +0200 | [diff] [blame] | 460 |     _validate_family(family) | 
| Benjamin Peterson | e711caf | 2008-06-11 16:44:04 +0000 | [diff] [blame] | 461 |     if family == 'AF_PIPE': | 
 | 462 |         c = PipeClient(address) | 
 | 463 |     else: | 
 | 464 |         c = SocketClient(address) | 
 | 465 |  | 
 | 466 |     if authkey is not None and not isinstance(authkey, bytes): | 
 | 467 |         raise TypeError('authkey should be a byte string') | 
 | 468 |  | 
 | 469 |     if authkey is not None: | 
 | 470 |         answer_challenge(c, authkey) | 
 | 471 |         deliver_challenge(c, authkey) | 
 | 472 |  | 
 | 473 |     return c | 
 | 474 |  | 
 | 475 |  | 
 | 476 | if sys.platform != 'win32': | 
 | 477 |  | 
 | 478 |     def Pipe(duplex=True): | 
 | 479 |         ''' | 
 | 480 |         Returns pair of connection objects at either end of a pipe | 
 | 481 |         ''' | 
 | 482 |         if duplex: | 
 | 483 |             s1, s2 = socket.socketpair() | 
| Antoine Pitrou | 5aa878c | 2011-05-09 21:00:28 +0200 | [diff] [blame] | 484 |             c1 = Connection(s1.detach()) | 
 | 485 |             c2 = Connection(s2.detach()) | 
| Benjamin Peterson | e711caf | 2008-06-11 16:44:04 +0000 | [diff] [blame] | 486 |         else: | 
 | 487 |             fd1, fd2 = os.pipe() | 
| Antoine Pitrou | 87cf220 | 2011-05-09 17:04:27 +0200 | [diff] [blame] | 488 |             c1 = Connection(fd1, writable=False) | 
 | 489 |             c2 = Connection(fd2, readable=False) | 
| Benjamin Peterson | e711caf | 2008-06-11 16:44:04 +0000 | [diff] [blame] | 490 |  | 
 | 491 |         return c1, c2 | 
 | 492 |  | 
 | 493 | else: | 
 | 494 |  | 
| Benjamin Peterson | e711caf | 2008-06-11 16:44:04 +0000 | [diff] [blame] | 495 |     def Pipe(duplex=True): | 
 | 496 |         ''' | 
 | 497 |         Returns pair of connection objects at either end of a pipe | 
 | 498 |         ''' | 
 | 499 |         address = arbitrary_address('AF_PIPE') | 
 | 500 |         if duplex: | 
| Antoine Pitrou | 23bba4c | 2012-04-18 20:51:15 +0200 | [diff] [blame] | 501 |             openmode = _winapi.PIPE_ACCESS_DUPLEX | 
 | 502 |             access = _winapi.GENERIC_READ | _winapi.GENERIC_WRITE | 
| Benjamin Peterson | e711caf | 2008-06-11 16:44:04 +0000 | [diff] [blame] | 503 |             obsize, ibsize = BUFSIZE, BUFSIZE | 
 | 504 |         else: | 
| Antoine Pitrou | 23bba4c | 2012-04-18 20:51:15 +0200 | [diff] [blame] | 505 |             openmode = _winapi.PIPE_ACCESS_INBOUND | 
 | 506 |             access = _winapi.GENERIC_WRITE | 
| Benjamin Peterson | e711caf | 2008-06-11 16:44:04 +0000 | [diff] [blame] | 507 |             obsize, ibsize = 0, BUFSIZE | 
 | 508 |  | 
| Antoine Pitrou | 23bba4c | 2012-04-18 20:51:15 +0200 | [diff] [blame] | 509 |         h1 = _winapi.CreateNamedPipe( | 
 | 510 |             address, openmode | _winapi.FILE_FLAG_OVERLAPPED | | 
 | 511 |             _winapi.FILE_FLAG_FIRST_PIPE_INSTANCE, | 
 | 512 |             _winapi.PIPE_TYPE_MESSAGE | _winapi.PIPE_READMODE_MESSAGE | | 
 | 513 |             _winapi.PIPE_WAIT, | 
 | 514 |             1, obsize, ibsize, _winapi.NMPWAIT_WAIT_FOREVER, _winapi.NULL | 
| Benjamin Peterson | e711caf | 2008-06-11 16:44:04 +0000 | [diff] [blame] | 515 |             ) | 
| Antoine Pitrou | 23bba4c | 2012-04-18 20:51:15 +0200 | [diff] [blame] | 516 |         h2 = _winapi.CreateFile( | 
 | 517 |             address, access, 0, _winapi.NULL, _winapi.OPEN_EXISTING, | 
 | 518 |             _winapi.FILE_FLAG_OVERLAPPED, _winapi.NULL | 
| Benjamin Peterson | e711caf | 2008-06-11 16:44:04 +0000 | [diff] [blame] | 519 |             ) | 
| Antoine Pitrou | 23bba4c | 2012-04-18 20:51:15 +0200 | [diff] [blame] | 520 |         _winapi.SetNamedPipeHandleState( | 
 | 521 |             h2, _winapi.PIPE_READMODE_MESSAGE, None, None | 
| Benjamin Peterson | e711caf | 2008-06-11 16:44:04 +0000 | [diff] [blame] | 522 |             ) | 
 | 523 |  | 
| Antoine Pitrou | 23bba4c | 2012-04-18 20:51:15 +0200 | [diff] [blame] | 524 |         overlapped = _winapi.ConnectNamedPipe(h1, overlapped=True) | 
| Antoine Pitrou | bdb1cf1 | 2012-03-05 19:28:37 +0100 | [diff] [blame] | 525 |         _, err = overlapped.GetOverlappedResult(True) | 
 | 526 |         assert err == 0 | 
| Benjamin Peterson | e711caf | 2008-06-11 16:44:04 +0000 | [diff] [blame] | 527 |  | 
| Antoine Pitrou | 87cf220 | 2011-05-09 17:04:27 +0200 | [diff] [blame] | 528 |         c1 = PipeConnection(h1, writable=duplex) | 
 | 529 |         c2 = PipeConnection(h2, readable=duplex) | 
| Benjamin Peterson | e711caf | 2008-06-11 16:44:04 +0000 | [diff] [blame] | 530 |  | 
 | 531 |         return c1, c2 | 
 | 532 |  | 
 | 533 | # | 
 | 534 | # Definitions for connections based on sockets | 
 | 535 | # | 
 | 536 |  | 
 | 537 | class SocketListener(object): | 
 | 538 |     ''' | 
| Georg Brandl | 734e268 | 2008-08-12 08:18:18 +0000 | [diff] [blame] | 539 |     Representation of a socket which is bound to an address and listening | 
| Benjamin Peterson | e711caf | 2008-06-11 16:44:04 +0000 | [diff] [blame] | 540 |     ''' | 
 | 541 |     def __init__(self, address, family, backlog=1): | 
 | 542 |         self._socket = socket.socket(getattr(socket, family)) | 
| Charles-François Natali | edc67fe | 2012-02-04 15:12:08 +0100 | [diff] [blame] | 543 |         try: | 
| Charles-François Natali | ed4a8fc | 2012-02-08 21:15:58 +0100 | [diff] [blame] | 544 |             # SO_REUSEADDR has different semantics on Windows (issue #2550). | 
 | 545 |             if os.name == 'posix': | 
 | 546 |                 self._socket.setsockopt(socket.SOL_SOCKET, | 
 | 547 |                                         socket.SO_REUSEADDR, 1) | 
| Charles-François Natali | edc67fe | 2012-02-04 15:12:08 +0100 | [diff] [blame] | 548 |             self._socket.bind(address) | 
 | 549 |             self._socket.listen(backlog) | 
 | 550 |             self._address = self._socket.getsockname() | 
 | 551 |         except OSError: | 
 | 552 |             self._socket.close() | 
 | 553 |             raise | 
| Benjamin Peterson | e711caf | 2008-06-11 16:44:04 +0000 | [diff] [blame] | 554 |         self._family = family | 
 | 555 |         self._last_accepted = None | 
 | 556 |  | 
| Benjamin Peterson | e711caf | 2008-06-11 16:44:04 +0000 | [diff] [blame] | 557 |         if family == 'AF_UNIX': | 
 | 558 |             self._unlink = Finalize( | 
| Georg Brandl | 2ee470f | 2008-07-16 12:55:28 +0000 | [diff] [blame] | 559 |                 self, os.unlink, args=(address,), exitpriority=0 | 
| Benjamin Peterson | e711caf | 2008-06-11 16:44:04 +0000 | [diff] [blame] | 560 |                 ) | 
 | 561 |         else: | 
 | 562 |             self._unlink = None | 
 | 563 |  | 
 | 564 |     def accept(self): | 
 | 565 |         s, self._last_accepted = self._socket.accept() | 
| Antoine Pitrou | df97cbe | 2012-04-07 22:38:52 +0200 | [diff] [blame] | 566 |         return Connection(s.detach()) | 
| Benjamin Peterson | e711caf | 2008-06-11 16:44:04 +0000 | [diff] [blame] | 567 |  | 
 | 568 |     def close(self): | 
 | 569 |         self._socket.close() | 
 | 570 |         if self._unlink is not None: | 
 | 571 |             self._unlink() | 
 | 572 |  | 
 | 573 |  | 
 | 574 | def SocketClient(address): | 
 | 575 |     ''' | 
 | 576 |     Return a connection object connected to the socket given by `address` | 
 | 577 |     ''' | 
 | 578 |     family = address_type(address) | 
| Victor Stinner | 2b69506 | 2011-01-03 15:47:59 +0000 | [diff] [blame] | 579 |     with socket.socket( getattr(socket, family) ) as s: | 
| Charles-François Natali | e6eabd4 | 2011-11-19 09:59:43 +0100 | [diff] [blame] | 580 |         s.connect(address) | 
| Antoine Pitrou | df97cbe | 2012-04-07 22:38:52 +0200 | [diff] [blame] | 581 |         return Connection(s.detach()) | 
| Benjamin Peterson | e711caf | 2008-06-11 16:44:04 +0000 | [diff] [blame] | 582 |  | 
 | 583 | # | 
 | 584 | # Definitions for connections based on named pipes | 
 | 585 | # | 
 | 586 |  | 
 | 587 | if sys.platform == 'win32': | 
 | 588 |  | 
 | 589 |     class PipeListener(object): | 
 | 590 |         ''' | 
 | 591 |         Representation of a named pipe | 
 | 592 |         ''' | 
 | 593 |         def __init__(self, address, backlog=None): | 
 | 594 |             self._address = address | 
| Antoine Pitrou | bdb1cf1 | 2012-03-05 19:28:37 +0100 | [diff] [blame] | 595 |             self._handle_queue = [self._new_handle(first=True)] | 
 | 596 |  | 
| Benjamin Peterson | e711caf | 2008-06-11 16:44:04 +0000 | [diff] [blame] | 597 |             self._last_accepted = None | 
| Benjamin Peterson | e711caf | 2008-06-11 16:44:04 +0000 | [diff] [blame] | 598 |             sub_debug('listener created with address=%r', self._address) | 
| Benjamin Peterson | e711caf | 2008-06-11 16:44:04 +0000 | [diff] [blame] | 599 |             self.close = Finalize( | 
 | 600 |                 self, PipeListener._finalize_pipe_listener, | 
 | 601 |                 args=(self._handle_queue, self._address), exitpriority=0 | 
 | 602 |                 ) | 
 | 603 |  | 
| Antoine Pitrou | bdb1cf1 | 2012-03-05 19:28:37 +0100 | [diff] [blame] | 604 |         def _new_handle(self, first=False): | 
| Antoine Pitrou | 23bba4c | 2012-04-18 20:51:15 +0200 | [diff] [blame] | 605 |             flags = _winapi.PIPE_ACCESS_DUPLEX | _winapi.FILE_FLAG_OVERLAPPED | 
| Antoine Pitrou | bdb1cf1 | 2012-03-05 19:28:37 +0100 | [diff] [blame] | 606 |             if first: | 
| Antoine Pitrou | 23bba4c | 2012-04-18 20:51:15 +0200 | [diff] [blame] | 607 |                 flags |= _winapi.FILE_FLAG_FIRST_PIPE_INSTANCE | 
 | 608 |             return _winapi.CreateNamedPipe( | 
| Antoine Pitrou | bdb1cf1 | 2012-03-05 19:28:37 +0100 | [diff] [blame] | 609 |                 self._address, flags, | 
| Antoine Pitrou | 23bba4c | 2012-04-18 20:51:15 +0200 | [diff] [blame] | 610 |                 _winapi.PIPE_TYPE_MESSAGE | _winapi.PIPE_READMODE_MESSAGE | | 
 | 611 |                 _winapi.PIPE_WAIT, | 
 | 612 |                 _winapi.PIPE_UNLIMITED_INSTANCES, BUFSIZE, BUFSIZE, | 
 | 613 |                 _winapi.NMPWAIT_WAIT_FOREVER, _winapi.NULL | 
| Benjamin Peterson | e711caf | 2008-06-11 16:44:04 +0000 | [diff] [blame] | 614 |                 ) | 
| Antoine Pitrou | bdb1cf1 | 2012-03-05 19:28:37 +0100 | [diff] [blame] | 615 |  | 
 | 616 |         def accept(self): | 
 | 617 |             self._handle_queue.append(self._new_handle()) | 
| Benjamin Peterson | e711caf | 2008-06-11 16:44:04 +0000 | [diff] [blame] | 618 |             handle = self._handle_queue.pop(0) | 
 | 619 |             try: | 
| Richard Oudkerk | fdb8dcf | 2012-05-05 19:45:37 +0100 | [diff] [blame] | 620 |                 ov = _winapi.ConnectNamedPipe(handle, overlapped=True) | 
 | 621 |             except OSError as e: | 
 | 622 |                 if e.winerror != _winapi.ERROR_NO_DATA: | 
 | 623 |                     raise | 
 | 624 |                 # ERROR_NO_DATA can occur if a client has already connected, | 
 | 625 |                 # written data and then disconnected -- see Issue 14725. | 
 | 626 |             else: | 
 | 627 |                 try: | 
 | 628 |                     res = _winapi.WaitForMultipleObjects( | 
 | 629 |                         [ov.event], False, INFINITE) | 
 | 630 |                 except: | 
 | 631 |                     ov.cancel() | 
 | 632 |                     _winapi.CloseHandle(handle) | 
 | 633 |                     raise | 
 | 634 |                 finally: | 
 | 635 |                     _, err = ov.GetOverlappedResult(True) | 
 | 636 |                     assert err == 0 | 
| Antoine Pitrou | 87cf220 | 2011-05-09 17:04:27 +0200 | [diff] [blame] | 637 |             return PipeConnection(handle) | 
| Benjamin Peterson | e711caf | 2008-06-11 16:44:04 +0000 | [diff] [blame] | 638 |  | 
 | 639 |         @staticmethod | 
 | 640 |         def _finalize_pipe_listener(queue, address): | 
 | 641 |             sub_debug('closing listener with address=%r', address) | 
 | 642 |             for handle in queue: | 
| Antoine Pitrou | 23bba4c | 2012-04-18 20:51:15 +0200 | [diff] [blame] | 643 |                 _winapi.CloseHandle(handle) | 
| Benjamin Peterson | e711caf | 2008-06-11 16:44:04 +0000 | [diff] [blame] | 644 |  | 
 | 645 |     def PipeClient(address): | 
 | 646 |         ''' | 
 | 647 |         Return a connection object connected to the pipe given by `address` | 
 | 648 |         ''' | 
| Antoine Pitrou | 45d61a3 | 2009-11-13 22:35:18 +0000 | [diff] [blame] | 649 |         t = _init_timeout() | 
| Benjamin Peterson | e711caf | 2008-06-11 16:44:04 +0000 | [diff] [blame] | 650 |         while 1: | 
 | 651 |             try: | 
| Antoine Pitrou | 23bba4c | 2012-04-18 20:51:15 +0200 | [diff] [blame] | 652 |                 _winapi.WaitNamedPipe(address, 1000) | 
 | 653 |                 h = _winapi.CreateFile( | 
 | 654 |                     address, _winapi.GENERIC_READ | _winapi.GENERIC_WRITE, | 
 | 655 |                     0, _winapi.NULL, _winapi.OPEN_EXISTING, | 
 | 656 |                     _winapi.FILE_FLAG_OVERLAPPED, _winapi.NULL | 
| Benjamin Peterson | e711caf | 2008-06-11 16:44:04 +0000 | [diff] [blame] | 657 |                     ) | 
 | 658 |             except WindowsError as e: | 
| Antoine Pitrou | 23bba4c | 2012-04-18 20:51:15 +0200 | [diff] [blame] | 659 |                 if e.winerror not in (_winapi.ERROR_SEM_TIMEOUT, | 
 | 660 |                                       _winapi.ERROR_PIPE_BUSY) or _check_timeout(t): | 
| Benjamin Peterson | e711caf | 2008-06-11 16:44:04 +0000 | [diff] [blame] | 661 |                     raise | 
 | 662 |             else: | 
 | 663 |                 break | 
 | 664 |         else: | 
 | 665 |             raise | 
 | 666 |  | 
| Antoine Pitrou | 23bba4c | 2012-04-18 20:51:15 +0200 | [diff] [blame] | 667 |         _winapi.SetNamedPipeHandleState( | 
 | 668 |             h, _winapi.PIPE_READMODE_MESSAGE, None, None | 
| Benjamin Peterson | e711caf | 2008-06-11 16:44:04 +0000 | [diff] [blame] | 669 |             ) | 
| Antoine Pitrou | 87cf220 | 2011-05-09 17:04:27 +0200 | [diff] [blame] | 670 |         return PipeConnection(h) | 
| Benjamin Peterson | e711caf | 2008-06-11 16:44:04 +0000 | [diff] [blame] | 671 |  | 
 | 672 | # | 
 | 673 | # Authentication stuff | 
 | 674 | # | 
 | 675 |  | 
 | 676 | MESSAGE_LENGTH = 20 | 
 | 677 |  | 
| Benjamin Peterson | 1fcfe21 | 2008-06-25 12:54:22 +0000 | [diff] [blame] | 678 | CHALLENGE = b'#CHALLENGE#' | 
 | 679 | WELCOME = b'#WELCOME#' | 
 | 680 | FAILURE = b'#FAILURE#' | 
| Benjamin Peterson | e711caf | 2008-06-11 16:44:04 +0000 | [diff] [blame] | 681 |  | 
 | 682 | def deliver_challenge(connection, authkey): | 
 | 683 |     import hmac | 
 | 684 |     assert isinstance(authkey, bytes) | 
 | 685 |     message = os.urandom(MESSAGE_LENGTH) | 
 | 686 |     connection.send_bytes(CHALLENGE + message) | 
 | 687 |     digest = hmac.new(authkey, message).digest() | 
 | 688 |     response = connection.recv_bytes(256)        # reject large message | 
 | 689 |     if response == digest: | 
 | 690 |         connection.send_bytes(WELCOME) | 
 | 691 |     else: | 
 | 692 |         connection.send_bytes(FAILURE) | 
 | 693 |         raise AuthenticationError('digest received was wrong') | 
 | 694 |  | 
 | 695 | def answer_challenge(connection, authkey): | 
 | 696 |     import hmac | 
 | 697 |     assert isinstance(authkey, bytes) | 
 | 698 |     message = connection.recv_bytes(256)         # reject large message | 
 | 699 |     assert message[:len(CHALLENGE)] == CHALLENGE, 'message = %r' % message | 
 | 700 |     message = message[len(CHALLENGE):] | 
 | 701 |     digest = hmac.new(authkey, message).digest() | 
 | 702 |     connection.send_bytes(digest) | 
 | 703 |     response = connection.recv_bytes(256)        # reject large message | 
 | 704 |     if response != WELCOME: | 
 | 705 |         raise AuthenticationError('digest sent was rejected') | 
 | 706 |  | 
 | 707 | # | 
 | 708 | # Support for using xmlrpclib for serialization | 
 | 709 | # | 
 | 710 |  | 
 | 711 | class ConnectionWrapper(object): | 
 | 712 |     def __init__(self, conn, dumps, loads): | 
 | 713 |         self._conn = conn | 
 | 714 |         self._dumps = dumps | 
 | 715 |         self._loads = loads | 
 | 716 |         for attr in ('fileno', 'close', 'poll', 'recv_bytes', 'send_bytes'): | 
 | 717 |             obj = getattr(conn, attr) | 
 | 718 |             setattr(self, attr, obj) | 
 | 719 |     def send(self, obj): | 
 | 720 |         s = self._dumps(obj) | 
 | 721 |         self._conn.send_bytes(s) | 
 | 722 |     def recv(self): | 
 | 723 |         s = self._conn.recv_bytes() | 
 | 724 |         return self._loads(s) | 
 | 725 |  | 
 | 726 | def _xml_dumps(obj): | 
| Marc-André Lemburg | 8f36af7 | 2011-02-25 15:42:01 +0000 | [diff] [blame] | 727 |     return xmlrpclib.dumps((obj,), None, None, None, 1).encode('utf-8') | 
| Benjamin Peterson | e711caf | 2008-06-11 16:44:04 +0000 | [diff] [blame] | 728 |  | 
 | 729 | def _xml_loads(s): | 
| Marc-André Lemburg | 8f36af7 | 2011-02-25 15:42:01 +0000 | [diff] [blame] | 730 |     (obj,), method = xmlrpclib.loads(s.decode('utf-8')) | 
| Benjamin Peterson | e711caf | 2008-06-11 16:44:04 +0000 | [diff] [blame] | 731 |     return obj | 
 | 732 |  | 
 | 733 | class XmlListener(Listener): | 
 | 734 |     def accept(self): | 
 | 735 |         global xmlrpclib | 
 | 736 |         import xmlrpc.client as xmlrpclib | 
 | 737 |         obj = Listener.accept(self) | 
 | 738 |         return ConnectionWrapper(obj, _xml_dumps, _xml_loads) | 
 | 739 |  | 
 | 740 | def XmlClient(*args, **kwds): | 
 | 741 |     global xmlrpclib | 
 | 742 |     import xmlrpc.client as xmlrpclib | 
 | 743 |     return ConnectionWrapper(Client(*args, **kwds), _xml_dumps, _xml_loads) | 
| Antoine Pitrou | 87cf220 | 2011-05-09 17:04:27 +0200 | [diff] [blame] | 744 |  | 
| Antoine Pitrou | bdb1cf1 | 2012-03-05 19:28:37 +0100 | [diff] [blame] | 745 | # | 
 | 746 | # Wait | 
 | 747 | # | 
 | 748 |  | 
 | 749 | if sys.platform == 'win32': | 
 | 750 |  | 
 | 751 |     def _exhaustive_wait(handles, timeout): | 
 | 752 |         # Return ALL handles which are currently signalled.  (Only | 
 | 753 |         # returning the first signalled might create starvation issues.) | 
 | 754 |         L = list(handles) | 
 | 755 |         ready = [] | 
 | 756 |         while L: | 
| Antoine Pitrou | 23bba4c | 2012-04-18 20:51:15 +0200 | [diff] [blame] | 757 |             res = _winapi.WaitForMultipleObjects(L, False, timeout) | 
| Antoine Pitrou | bdb1cf1 | 2012-03-05 19:28:37 +0100 | [diff] [blame] | 758 |             if res == WAIT_TIMEOUT: | 
 | 759 |                 break | 
 | 760 |             elif WAIT_OBJECT_0 <= res < WAIT_OBJECT_0 + len(L): | 
 | 761 |                 res -= WAIT_OBJECT_0 | 
 | 762 |             elif WAIT_ABANDONED_0 <= res < WAIT_ABANDONED_0 + len(L): | 
 | 763 |                 res -= WAIT_ABANDONED_0 | 
 | 764 |             else: | 
 | 765 |                 raise RuntimeError('Should not get here') | 
 | 766 |             ready.append(L[res]) | 
 | 767 |             L = L[res+1:] | 
 | 768 |             timeout = 0 | 
 | 769 |         return ready | 
 | 770 |  | 
| Antoine Pitrou | 23bba4c | 2012-04-18 20:51:15 +0200 | [diff] [blame] | 771 |     _ready_errors = {_winapi.ERROR_BROKEN_PIPE, _winapi.ERROR_NETNAME_DELETED} | 
| Antoine Pitrou | bdb1cf1 | 2012-03-05 19:28:37 +0100 | [diff] [blame] | 772 |  | 
 | 773 |     def wait(object_list, timeout=None): | 
 | 774 |         ''' | 
 | 775 |         Wait till an object in object_list is ready/readable. | 
 | 776 |  | 
 | 777 |         Returns list of those objects in object_list which are ready/readable. | 
 | 778 |         ''' | 
 | 779 |         if timeout is None: | 
 | 780 |             timeout = INFINITE | 
 | 781 |         elif timeout < 0: | 
 | 782 |             timeout = 0 | 
 | 783 |         else: | 
 | 784 |             timeout = int(timeout * 1000 + 0.5) | 
 | 785 |  | 
 | 786 |         object_list = list(object_list) | 
 | 787 |         waithandle_to_obj = {} | 
 | 788 |         ov_list = [] | 
 | 789 |         ready_objects = set() | 
 | 790 |         ready_handles = set() | 
 | 791 |  | 
 | 792 |         try: | 
 | 793 |             for o in object_list: | 
 | 794 |                 try: | 
 | 795 |                     fileno = getattr(o, 'fileno') | 
 | 796 |                 except AttributeError: | 
 | 797 |                     waithandle_to_obj[o.__index__()] = o | 
 | 798 |                 else: | 
 | 799 |                     # start an overlapped read of length zero | 
 | 800 |                     try: | 
| Antoine Pitrou | 23bba4c | 2012-04-18 20:51:15 +0200 | [diff] [blame] | 801 |                         ov, err = _winapi.ReadFile(fileno(), 0, True) | 
| Antoine Pitrou | bdb1cf1 | 2012-03-05 19:28:37 +0100 | [diff] [blame] | 802 |                     except OSError as e: | 
 | 803 |                         err = e.winerror | 
 | 804 |                         if err not in _ready_errors: | 
 | 805 |                             raise | 
| Antoine Pitrou | 23bba4c | 2012-04-18 20:51:15 +0200 | [diff] [blame] | 806 |                     if err == _winapi.ERROR_IO_PENDING: | 
| Antoine Pitrou | bdb1cf1 | 2012-03-05 19:28:37 +0100 | [diff] [blame] | 807 |                         ov_list.append(ov) | 
 | 808 |                         waithandle_to_obj[ov.event] = o | 
 | 809 |                     else: | 
 | 810 |                         # If o.fileno() is an overlapped pipe handle and | 
 | 811 |                         # err == 0 then there is a zero length message | 
 | 812 |                         # in the pipe, but it HAS NOT been consumed. | 
 | 813 |                         ready_objects.add(o) | 
 | 814 |                         timeout = 0 | 
 | 815 |  | 
 | 816 |             ready_handles = _exhaustive_wait(waithandle_to_obj.keys(), timeout) | 
 | 817 |         finally: | 
 | 818 |             # request that overlapped reads stop | 
 | 819 |             for ov in ov_list: | 
 | 820 |                 ov.cancel() | 
 | 821 |  | 
 | 822 |             # wait for all overlapped reads to stop | 
 | 823 |             for ov in ov_list: | 
 | 824 |                 try: | 
 | 825 |                     _, err = ov.GetOverlappedResult(True) | 
 | 826 |                 except OSError as e: | 
 | 827 |                     err = e.winerror | 
 | 828 |                     if err not in _ready_errors: | 
 | 829 |                         raise | 
| Antoine Pitrou | 23bba4c | 2012-04-18 20:51:15 +0200 | [diff] [blame] | 830 |                 if err != _winapi.ERROR_OPERATION_ABORTED: | 
| Antoine Pitrou | bdb1cf1 | 2012-03-05 19:28:37 +0100 | [diff] [blame] | 831 |                     o = waithandle_to_obj[ov.event] | 
 | 832 |                     ready_objects.add(o) | 
 | 833 |                     if err == 0: | 
 | 834 |                         # If o.fileno() is an overlapped pipe handle then | 
 | 835 |                         # a zero length message HAS been consumed. | 
 | 836 |                         if hasattr(o, '_got_empty_message'): | 
 | 837 |                             o._got_empty_message = True | 
 | 838 |  | 
 | 839 |         ready_objects.update(waithandle_to_obj[h] for h in ready_handles) | 
 | 840 |         return [o for o in object_list if o in ready_objects] | 
 | 841 |  | 
 | 842 | else: | 
 | 843 |  | 
 | 844 |     def wait(object_list, timeout=None): | 
 | 845 |         ''' | 
 | 846 |         Wait till an object in object_list is ready/readable. | 
 | 847 |  | 
 | 848 |         Returns list of those objects in object_list which are ready/readable. | 
 | 849 |         ''' | 
 | 850 |         if timeout is not None: | 
 | 851 |             if timeout <= 0: | 
 | 852 |                 return select.select(object_list, [], [], 0)[0] | 
 | 853 |             else: | 
 | 854 |                 deadline = time.time() + timeout | 
 | 855 |         while True: | 
 | 856 |             try: | 
 | 857 |                 return select.select(object_list, [], [], timeout)[0] | 
 | 858 |             except OSError as e: | 
 | 859 |                 if e.errno != errno.EINTR: | 
 | 860 |                     raise | 
 | 861 |             if timeout is not None: | 
 | 862 |                 timeout = deadline - time.time() | 
| Antoine Pitrou | 5438ed1 | 2012-04-24 22:56:57 +0200 | [diff] [blame] | 863 |  | 
 | 864 | # | 
 | 865 | # Make connection and socket objects sharable if possible | 
 | 866 | # | 
 | 867 |  | 
 | 868 | if sys.platform == 'win32': | 
 | 869 |     from . import reduction | 
 | 870 |     ForkingPickler.register(socket.socket, reduction.reduce_socket) | 
 | 871 |     ForkingPickler.register(Connection, reduction.reduce_connection) | 
 | 872 |     ForkingPickler.register(PipeConnection, reduction.reduce_pipe_connection) | 
 | 873 | else: | 
 | 874 |     try: | 
 | 875 |         from . import reduction | 
 | 876 |     except ImportError: | 
 | 877 |         pass | 
 | 878 |     else: | 
 | 879 |         ForkingPickler.register(socket.socket, reduction.reduce_socket) | 
 | 880 |         ForkingPickler.register(Connection, reduction.reduce_connection) |