blob: 1093d9fb97f23c93f4345db07ce1b73887fe75e9 [file] [log] [blame]
Benjamin Petersone711caf2008-06-11 16:44:04 +00001#
2# A higher level module for using sockets (or Windows named pipes)
3#
4# multiprocessing/connection.py
5#
R. David Murray3fc969a2010-12-14 01:38:16 +00006# Copyright (c) 2006-2008, R Oudkerk
Richard Oudkerk3e268aa2012-04-30 12:13:55 +01007# Licensed to PSF under a Contributor Agreement.
Benjamin Petersone711caf2008-06-11 16:44:04 +00008#
9
Antoine Pitroubdb1cf12012-03-05 19:28:37 +010010__all__ = [ 'Client', 'Listener', 'Pipe', 'wait' ]
Benjamin Petersone711caf2008-06-11 16:44:04 +000011
Antoine Pitrou87cf2202011-05-09 17:04:27 +020012import io
Benjamin Petersone711caf2008-06-11 16:44:04 +000013import os
14import sys
Antoine Pitrou87cf2202011-05-09 17:04:27 +020015import select
Benjamin Petersone711caf2008-06-11 16:44:04 +000016import socket
Antoine Pitrou87cf2202011-05-09 17:04:27 +020017import struct
Georg Brandl6aa2d1f2008-08-12 08:35:52 +000018import errno
Benjamin Petersone711caf2008-06-11 16:44:04 +000019import time
20import tempfile
21import itertools
22
23import _multiprocessing
Antoine Pitrou87cf2202011-05-09 17:04:27 +020024from multiprocessing import current_process, AuthenticationError, BufferTooShort
Richard Oudkerk59d54042012-05-10 16:11:12 +010025from multiprocessing.util import get_temp_dir, Finalize, sub_debug, debug
Antoine Pitrou5438ed12012-04-24 22:56:57 +020026from multiprocessing.forking import ForkingPickler
Antoine Pitrou87cf2202011-05-09 17:04:27 +020027try:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +020028 import _winapi
29 from _winapi import WAIT_OBJECT_0, WAIT_TIMEOUT, INFINITE
Brett Cannoncd171c82013-07-04 17:43:24 -040030except ImportError:
Antoine Pitrou87cf2202011-05-09 17:04:27 +020031 if sys.platform == 'win32':
32 raise
Antoine Pitrou23bba4c2012-04-18 20:51:15 +020033 _winapi = None
Benjamin Petersone711caf2008-06-11 16:44:04 +000034
35#
36#
37#
38
39BUFSIZE = 8192
Antoine Pitrou45d61a32009-11-13 22:35:18 +000040# A very generous timeout when it comes to local connections...
41CONNECTION_TIMEOUT = 20.
Benjamin Petersone711caf2008-06-11 16:44:04 +000042
43_mmap_counter = itertools.count()
44
45default_family = 'AF_INET'
46families = ['AF_INET']
47
48if hasattr(socket, 'AF_UNIX'):
49 default_family = 'AF_UNIX'
50 families += ['AF_UNIX']
51
52if sys.platform == 'win32':
53 default_family = 'AF_PIPE'
54 families += ['AF_PIPE']
55
Antoine Pitrou45d61a32009-11-13 22:35:18 +000056
57def _init_timeout(timeout=CONNECTION_TIMEOUT):
58 return time.time() + timeout
59
60def _check_timeout(t):
61 return time.time() > t
62
Benjamin Petersone711caf2008-06-11 16:44:04 +000063#
64#
65#
66
67def arbitrary_address(family):
68 '''
69 Return an arbitrary free address for the given family
70 '''
71 if family == 'AF_INET':
72 return ('localhost', 0)
73 elif family == 'AF_UNIX':
74 return tempfile.mktemp(prefix='listener-', dir=get_temp_dir())
75 elif family == 'AF_PIPE':
76 return tempfile.mktemp(prefix=r'\\.\pipe\pyc-%d-%d-' %
77 (os.getpid(), next(_mmap_counter)))
78 else:
79 raise ValueError('unrecognized family')
80
Antoine Pitrou709176f2012-04-01 17:19:09 +020081def _validate_family(family):
82 '''
83 Checks if the family is valid for the current environment.
84 '''
85 if sys.platform != 'win32' and family == 'AF_PIPE':
86 raise ValueError('Family %s is not recognized.' % family)
87
Antoine Pitrou6d20cba2012-04-03 20:12:23 +020088 if sys.platform == 'win32' and family == 'AF_UNIX':
89 # double check
90 if not hasattr(socket, family):
91 raise ValueError('Family %s is not recognized.' % family)
Benjamin Petersone711caf2008-06-11 16:44:04 +000092
93def address_type(address):
94 '''
95 Return the types of the address
96
97 This can be 'AF_INET', 'AF_UNIX', or 'AF_PIPE'
98 '''
99 if type(address) == tuple:
100 return 'AF_INET'
101 elif type(address) is str and address.startswith('\\\\'):
102 return 'AF_PIPE'
103 elif type(address) is str:
104 return 'AF_UNIX'
105 else:
106 raise ValueError('address type of %r unrecognized' % address)
107
108#
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200109# Connection classes
110#
111
112class _ConnectionBase:
113 _handle = None
114
115 def __init__(self, handle, readable=True, writable=True):
116 handle = handle.__index__()
117 if handle < 0:
118 raise ValueError("invalid handle")
119 if not readable and not writable:
120 raise ValueError(
121 "at least one of `readable` and `writable` must be True")
122 self._handle = handle
123 self._readable = readable
124 self._writable = writable
125
Antoine Pitrou60001202011-07-09 01:03:46 +0200126 # XXX should we use util.Finalize instead of a __del__?
127
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200128 def __del__(self):
129 if self._handle is not None:
130 self._close()
131
132 def _check_closed(self):
133 if self._handle is None:
Andrew Svetlovf7a17b42012-12-25 16:47:37 +0200134 raise OSError("handle is closed")
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200135
136 def _check_readable(self):
137 if not self._readable:
Andrew Svetlovf7a17b42012-12-25 16:47:37 +0200138 raise OSError("connection is write-only")
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200139
140 def _check_writable(self):
141 if not self._writable:
Andrew Svetlovf7a17b42012-12-25 16:47:37 +0200142 raise OSError("connection is read-only")
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200143
144 def _bad_message_length(self):
145 if self._writable:
146 self._readable = False
147 else:
148 self.close()
Andrew Svetlovf7a17b42012-12-25 16:47:37 +0200149 raise OSError("bad message length")
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200150
151 @property
152 def closed(self):
153 """True if the connection is closed"""
154 return self._handle is None
155
156 @property
157 def readable(self):
158 """True if the connection is readable"""
159 return self._readable
160
161 @property
162 def writable(self):
163 """True if the connection is writable"""
164 return self._writable
165
166 def fileno(self):
167 """File descriptor or handle of the connection"""
168 self._check_closed()
169 return self._handle
170
171 def close(self):
172 """Close the connection"""
173 if self._handle is not None:
174 try:
175 self._close()
176 finally:
177 self._handle = None
178
179 def send_bytes(self, buf, offset=0, size=None):
180 """Send the bytes data from a bytes-like object"""
181 self._check_closed()
182 self._check_writable()
183 m = memoryview(buf)
184 # HACK for byte-indexing of non-bytewise buffers (e.g. array.array)
185 if m.itemsize > 1:
186 m = memoryview(bytes(m))
187 n = len(m)
188 if offset < 0:
189 raise ValueError("offset is negative")
190 if n < offset:
191 raise ValueError("buffer length < offset")
192 if size is None:
193 size = n - offset
194 elif size < 0:
195 raise ValueError("size is negative")
196 elif offset + size > n:
197 raise ValueError("buffer length < offset + size")
198 self._send_bytes(m[offset:offset + size])
199
200 def send(self, obj):
201 """Send a (picklable) object"""
202 self._check_closed()
203 self._check_writable()
Charles-François Natalia6550752013-03-24 15:21:49 +0100204 self._send_bytes(ForkingPickler.dumps(obj))
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200205
206 def recv_bytes(self, maxlength=None):
207 """
208 Receive bytes data as a bytes object.
209 """
210 self._check_closed()
211 self._check_readable()
212 if maxlength is not None and maxlength < 0:
213 raise ValueError("negative maxlength")
214 buf = self._recv_bytes(maxlength)
215 if buf is None:
216 self._bad_message_length()
217 return buf.getvalue()
218
219 def recv_bytes_into(self, buf, offset=0):
220 """
221 Receive bytes data into a writeable buffer-like object.
222 Return the number of bytes read.
223 """
224 self._check_closed()
225 self._check_readable()
226 with memoryview(buf) as m:
227 # Get bytesize of arbitrary buffer
228 itemsize = m.itemsize
229 bytesize = itemsize * len(m)
230 if offset < 0:
231 raise ValueError("negative offset")
232 elif offset > bytesize:
233 raise ValueError("offset too large")
234 result = self._recv_bytes()
235 size = result.tell()
236 if bytesize < offset + size:
237 raise BufferTooShort(result.getvalue())
238 # Message can fit in dest
239 result.seek(0)
240 result.readinto(m[offset // itemsize :
241 (offset + size) // itemsize])
242 return size
243
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100244 def recv(self):
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200245 """Receive a (picklable) object"""
246 self._check_closed()
247 self._check_readable()
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100248 buf = self._recv_bytes()
Charles-François Natalia6550752013-03-24 15:21:49 +0100249 return ForkingPickler.loads(buf.getbuffer())
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200250
251 def poll(self, timeout=0.0):
252 """Whether there is any input available to be read"""
253 self._check_closed()
254 self._check_readable()
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200255 return self._poll(timeout)
256
Richard Oudkerkd69cfe82012-06-18 17:47:52 +0100257 def __enter__(self):
258 return self
259
260 def __exit__(self, exc_type, exc_value, exc_tb):
261 self.close()
262
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200263
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200264if _winapi:
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200265
266 class PipeConnection(_ConnectionBase):
267 """
268 Connection class based on a Windows named pipe.
Antoine Pitroudd696492011-06-08 17:21:55 +0200269 Overlapped I/O is used, so the handles must have been created
270 with FILE_FLAG_OVERLAPPED.
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200271 """
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100272 _got_empty_message = False
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200273
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200274 def _close(self, _CloseHandle=_winapi.CloseHandle):
Antoine Pitrou71a28a92011-07-09 01:03:00 +0200275 _CloseHandle(self._handle)
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200276
277 def _send_bytes(self, buf):
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200278 ov, err = _winapi.WriteFile(self._handle, buf, overlapped=True)
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100279 try:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200280 if err == _winapi.ERROR_IO_PENDING:
281 waitres = _winapi.WaitForMultipleObjects(
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100282 [ov.event], False, INFINITE)
283 assert waitres == WAIT_OBJECT_0
284 except:
285 ov.cancel()
286 raise
287 finally:
288 nwritten, err = ov.GetOverlappedResult(True)
289 assert err == 0
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200290 assert nwritten == len(buf)
291
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100292 def _recv_bytes(self, maxsize=None):
293 if self._got_empty_message:
294 self._got_empty_message = False
295 return io.BytesIO()
Antoine Pitroudd696492011-06-08 17:21:55 +0200296 else:
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100297 bsize = 128 if maxsize is None else min(maxsize, 128)
Antoine Pitroudd696492011-06-08 17:21:55 +0200298 try:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200299 ov, err = _winapi.ReadFile(self._handle, bsize,
300 overlapped=True)
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100301 try:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200302 if err == _winapi.ERROR_IO_PENDING:
303 waitres = _winapi.WaitForMultipleObjects(
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100304 [ov.event], False, INFINITE)
305 assert waitres == WAIT_OBJECT_0
306 except:
307 ov.cancel()
308 raise
309 finally:
310 nread, err = ov.GetOverlappedResult(True)
311 if err == 0:
312 f = io.BytesIO()
313 f.write(ov.getbuffer())
314 return f
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200315 elif err == _winapi.ERROR_MORE_DATA:
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100316 return self._get_more_data(ov, maxsize)
Andrew Svetlovf7a17b42012-12-25 16:47:37 +0200317 except OSError as e:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200318 if e.winerror == _winapi.ERROR_BROKEN_PIPE:
Antoine Pitroudd696492011-06-08 17:21:55 +0200319 raise EOFError
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100320 else:
321 raise
322 raise RuntimeError("shouldn't get here; expected KeyboardInterrupt")
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200323
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100324 def _poll(self, timeout):
325 if (self._got_empty_message or
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200326 _winapi.PeekNamedPipe(self._handle)[0] != 0):
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200327 return True
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100328 return bool(wait([self], timeout))
329
330 def _get_more_data(self, ov, maxsize):
331 buf = ov.getbuffer()
332 f = io.BytesIO()
333 f.write(buf)
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200334 left = _winapi.PeekNamedPipe(self._handle)[1]
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100335 assert left > 0
336 if maxsize is not None and len(buf) + left > maxsize:
337 self._bad_message_length()
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200338 ov, err = _winapi.ReadFile(self._handle, left, overlapped=True)
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100339 rbytes, err = ov.GetOverlappedResult(True)
340 assert err == 0
341 assert rbytes == left
342 f.write(ov.getbuffer())
343 return f
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200344
345
346class Connection(_ConnectionBase):
347 """
348 Connection class based on an arbitrary file descriptor (Unix only), or
349 a socket handle (Windows).
350 """
351
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200352 if _winapi:
353 def _close(self, _close=_multiprocessing.closesocket):
Antoine Pitrou71a28a92011-07-09 01:03:00 +0200354 _close(self._handle)
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200355 _write = _multiprocessing.send
356 _read = _multiprocessing.recv
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200357 else:
Antoine Pitrou71a28a92011-07-09 01:03:00 +0200358 def _close(self, _close=os.close):
359 _close(self._handle)
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200360 _write = os.write
361 _read = os.read
362
363 def _send(self, buf, write=_write):
364 remaining = len(buf)
365 while True:
Richard Oudkerkcca8c532013-07-01 18:59:26 +0100366 try:
367 n = write(self._handle, buf)
368 except InterruptedError:
369 continue
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200370 remaining -= n
371 if remaining == 0:
372 break
373 buf = buf[n:]
374
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100375 def _recv(self, size, read=_read):
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200376 buf = io.BytesIO()
Antoine Pitroudd696492011-06-08 17:21:55 +0200377 handle = self._handle
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200378 remaining = size
379 while remaining > 0:
Richard Oudkerkcca8c532013-07-01 18:59:26 +0100380 try:
381 chunk = read(handle, remaining)
382 except InterruptedError:
383 continue
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200384 n = len(chunk)
385 if n == 0:
386 if remaining == size:
387 raise EOFError
388 else:
Andrew Svetlovf7a17b42012-12-25 16:47:37 +0200389 raise OSError("got end of file during message")
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200390 buf.write(chunk)
391 remaining -= n
392 return buf
393
394 def _send_bytes(self, buf):
395 # For wire compatibility with 3.2 and lower
396 n = len(buf)
Charles-François Natali225aa4f2011-09-20 19:27:39 +0200397 self._send(struct.pack("!i", n))
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200398 # The condition is necessary to avoid "broken pipe" errors
399 # when sending a 0-length buffer if the other end closed the pipe.
400 if n > 0:
401 self._send(buf)
402
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100403 def _recv_bytes(self, maxsize=None):
404 buf = self._recv(4)
Charles-François Natali225aa4f2011-09-20 19:27:39 +0200405 size, = struct.unpack("!i", buf.getvalue())
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200406 if maxsize is not None and size > maxsize:
407 return None
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100408 return self._recv(size)
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200409
410 def _poll(self, timeout):
Richard Oudkerked9e06c2013-01-13 22:46:48 +0000411 r = wait([self], timeout)
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200412 return bool(r)
413
414
415#
Benjamin Petersone711caf2008-06-11 16:44:04 +0000416# Public functions
417#
418
419class Listener(object):
420 '''
421 Returns a listener object.
422
423 This is a wrapper for a bound socket which is 'listening' for
424 connections, or for a Windows named pipe.
425 '''
426 def __init__(self, address=None, family=None, backlog=1, authkey=None):
427 family = family or (address and address_type(address)) \
428 or default_family
429 address = address or arbitrary_address(family)
430
Antoine Pitrou709176f2012-04-01 17:19:09 +0200431 _validate_family(family)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000432 if family == 'AF_PIPE':
433 self._listener = PipeListener(address, backlog)
434 else:
435 self._listener = SocketListener(address, family, backlog)
436
437 if authkey is not None and not isinstance(authkey, bytes):
438 raise TypeError('authkey should be a byte string')
439
440 self._authkey = authkey
441
442 def accept(self):
443 '''
444 Accept a connection on the bound socket or named pipe of `self`.
445
446 Returns a `Connection` object.
447 '''
Richard Oudkerkd69cfe82012-06-18 17:47:52 +0100448 if self._listener is None:
Andrew Svetlovf7a17b42012-12-25 16:47:37 +0200449 raise OSError('listener is closed')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000450 c = self._listener.accept()
451 if self._authkey:
452 deliver_challenge(c, self._authkey)
453 answer_challenge(c, self._authkey)
454 return c
455
456 def close(self):
457 '''
458 Close the bound socket or named pipe of `self`.
459 '''
Richard Oudkerkd69cfe82012-06-18 17:47:52 +0100460 if self._listener is not None:
461 self._listener.close()
462 self._listener = None
Benjamin Petersone711caf2008-06-11 16:44:04 +0000463
464 address = property(lambda self: self._listener._address)
465 last_accepted = property(lambda self: self._listener._last_accepted)
466
Richard Oudkerkd69cfe82012-06-18 17:47:52 +0100467 def __enter__(self):
468 return self
469
470 def __exit__(self, exc_type, exc_value, exc_tb):
471 self.close()
472
Benjamin Petersone711caf2008-06-11 16:44:04 +0000473
474def Client(address, family=None, authkey=None):
475 '''
476 Returns a connection to the address of a `Listener`
477 '''
478 family = family or address_type(address)
Antoine Pitrou709176f2012-04-01 17:19:09 +0200479 _validate_family(family)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000480 if family == 'AF_PIPE':
481 c = PipeClient(address)
482 else:
483 c = SocketClient(address)
484
485 if authkey is not None and not isinstance(authkey, bytes):
486 raise TypeError('authkey should be a byte string')
487
488 if authkey is not None:
489 answer_challenge(c, authkey)
490 deliver_challenge(c, authkey)
491
492 return c
493
494
495if sys.platform != 'win32':
496
497 def Pipe(duplex=True):
498 '''
499 Returns pair of connection objects at either end of a pipe
500 '''
501 if duplex:
502 s1, s2 = socket.socketpair()
Richard Oudkerkb15e6222012-07-27 14:19:00 +0100503 s1.setblocking(True)
504 s2.setblocking(True)
Antoine Pitrou5aa878c2011-05-09 21:00:28 +0200505 c1 = Connection(s1.detach())
506 c2 = Connection(s2.detach())
Benjamin Petersone711caf2008-06-11 16:44:04 +0000507 else:
508 fd1, fd2 = os.pipe()
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200509 c1 = Connection(fd1, writable=False)
510 c2 = Connection(fd2, readable=False)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000511
512 return c1, c2
513
514else:
515
Benjamin Petersone711caf2008-06-11 16:44:04 +0000516 def Pipe(duplex=True):
517 '''
518 Returns pair of connection objects at either end of a pipe
519 '''
520 address = arbitrary_address('AF_PIPE')
521 if duplex:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200522 openmode = _winapi.PIPE_ACCESS_DUPLEX
523 access = _winapi.GENERIC_READ | _winapi.GENERIC_WRITE
Benjamin Petersone711caf2008-06-11 16:44:04 +0000524 obsize, ibsize = BUFSIZE, BUFSIZE
525 else:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200526 openmode = _winapi.PIPE_ACCESS_INBOUND
527 access = _winapi.GENERIC_WRITE
Benjamin Petersone711caf2008-06-11 16:44:04 +0000528 obsize, ibsize = 0, BUFSIZE
529
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200530 h1 = _winapi.CreateNamedPipe(
531 address, openmode | _winapi.FILE_FLAG_OVERLAPPED |
532 _winapi.FILE_FLAG_FIRST_PIPE_INSTANCE,
533 _winapi.PIPE_TYPE_MESSAGE | _winapi.PIPE_READMODE_MESSAGE |
534 _winapi.PIPE_WAIT,
535 1, obsize, ibsize, _winapi.NMPWAIT_WAIT_FOREVER, _winapi.NULL
Benjamin Petersone711caf2008-06-11 16:44:04 +0000536 )
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200537 h2 = _winapi.CreateFile(
538 address, access, 0, _winapi.NULL, _winapi.OPEN_EXISTING,
539 _winapi.FILE_FLAG_OVERLAPPED, _winapi.NULL
Benjamin Petersone711caf2008-06-11 16:44:04 +0000540 )
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200541 _winapi.SetNamedPipeHandleState(
542 h2, _winapi.PIPE_READMODE_MESSAGE, None, None
Benjamin Petersone711caf2008-06-11 16:44:04 +0000543 )
544
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200545 overlapped = _winapi.ConnectNamedPipe(h1, overlapped=True)
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100546 _, err = overlapped.GetOverlappedResult(True)
547 assert err == 0
Benjamin Petersone711caf2008-06-11 16:44:04 +0000548
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200549 c1 = PipeConnection(h1, writable=duplex)
550 c2 = PipeConnection(h2, readable=duplex)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000551
552 return c1, c2
553
554#
555# Definitions for connections based on sockets
556#
557
558class SocketListener(object):
559 '''
Georg Brandl734e2682008-08-12 08:18:18 +0000560 Representation of a socket which is bound to an address and listening
Benjamin Petersone711caf2008-06-11 16:44:04 +0000561 '''
562 def __init__(self, address, family, backlog=1):
563 self._socket = socket.socket(getattr(socket, family))
Charles-François Nataliedc67fe2012-02-04 15:12:08 +0100564 try:
Charles-François Natalied4a8fc2012-02-08 21:15:58 +0100565 # SO_REUSEADDR has different semantics on Windows (issue #2550).
566 if os.name == 'posix':
567 self._socket.setsockopt(socket.SOL_SOCKET,
568 socket.SO_REUSEADDR, 1)
Richard Oudkerkb15e6222012-07-27 14:19:00 +0100569 self._socket.setblocking(True)
Charles-François Nataliedc67fe2012-02-04 15:12:08 +0100570 self._socket.bind(address)
571 self._socket.listen(backlog)
572 self._address = self._socket.getsockname()
573 except OSError:
574 self._socket.close()
575 raise
Benjamin Petersone711caf2008-06-11 16:44:04 +0000576 self._family = family
577 self._last_accepted = None
578
Benjamin Petersone711caf2008-06-11 16:44:04 +0000579 if family == 'AF_UNIX':
580 self._unlink = Finalize(
Georg Brandl2ee470f2008-07-16 12:55:28 +0000581 self, os.unlink, args=(address,), exitpriority=0
Benjamin Petersone711caf2008-06-11 16:44:04 +0000582 )
583 else:
584 self._unlink = None
585
586 def accept(self):
Richard Oudkerkcca8c532013-07-01 18:59:26 +0100587 while True:
588 try:
589 s, self._last_accepted = self._socket.accept()
590 except InterruptedError:
591 pass
592 else:
593 break
Richard Oudkerkb15e6222012-07-27 14:19:00 +0100594 s.setblocking(True)
Antoine Pitroudf97cbe2012-04-07 22:38:52 +0200595 return Connection(s.detach())
Benjamin Petersone711caf2008-06-11 16:44:04 +0000596
597 def close(self):
598 self._socket.close()
599 if self._unlink is not None:
600 self._unlink()
601
602
603def SocketClient(address):
604 '''
605 Return a connection object connected to the socket given by `address`
606 '''
607 family = address_type(address)
Victor Stinner2b695062011-01-03 15:47:59 +0000608 with socket.socket( getattr(socket, family) ) as s:
Richard Oudkerkb15e6222012-07-27 14:19:00 +0100609 s.setblocking(True)
Charles-François Natalie6eabd42011-11-19 09:59:43 +0100610 s.connect(address)
Antoine Pitroudf97cbe2012-04-07 22:38:52 +0200611 return Connection(s.detach())
Benjamin Petersone711caf2008-06-11 16:44:04 +0000612
613#
614# Definitions for connections based on named pipes
615#
616
617if sys.platform == 'win32':
618
619 class PipeListener(object):
620 '''
621 Representation of a named pipe
622 '''
623 def __init__(self, address, backlog=None):
624 self._address = address
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100625 self._handle_queue = [self._new_handle(first=True)]
626
Benjamin Petersone711caf2008-06-11 16:44:04 +0000627 self._last_accepted = None
Benjamin Petersone711caf2008-06-11 16:44:04 +0000628 sub_debug('listener created with address=%r', self._address)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000629 self.close = Finalize(
630 self, PipeListener._finalize_pipe_listener,
631 args=(self._handle_queue, self._address), exitpriority=0
632 )
633
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100634 def _new_handle(self, first=False):
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200635 flags = _winapi.PIPE_ACCESS_DUPLEX | _winapi.FILE_FLAG_OVERLAPPED
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100636 if first:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200637 flags |= _winapi.FILE_FLAG_FIRST_PIPE_INSTANCE
638 return _winapi.CreateNamedPipe(
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100639 self._address, flags,
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200640 _winapi.PIPE_TYPE_MESSAGE | _winapi.PIPE_READMODE_MESSAGE |
641 _winapi.PIPE_WAIT,
642 _winapi.PIPE_UNLIMITED_INSTANCES, BUFSIZE, BUFSIZE,
643 _winapi.NMPWAIT_WAIT_FOREVER, _winapi.NULL
Benjamin Petersone711caf2008-06-11 16:44:04 +0000644 )
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100645
646 def accept(self):
647 self._handle_queue.append(self._new_handle())
Benjamin Petersone711caf2008-06-11 16:44:04 +0000648 handle = self._handle_queue.pop(0)
649 try:
Richard Oudkerkfdb8dcf2012-05-05 19:45:37 +0100650 ov = _winapi.ConnectNamedPipe(handle, overlapped=True)
651 except OSError as e:
652 if e.winerror != _winapi.ERROR_NO_DATA:
653 raise
654 # ERROR_NO_DATA can occur if a client has already connected,
655 # written data and then disconnected -- see Issue 14725.
656 else:
657 try:
658 res = _winapi.WaitForMultipleObjects(
659 [ov.event], False, INFINITE)
660 except:
661 ov.cancel()
662 _winapi.CloseHandle(handle)
663 raise
664 finally:
665 _, err = ov.GetOverlappedResult(True)
666 assert err == 0
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200667 return PipeConnection(handle)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000668
669 @staticmethod
670 def _finalize_pipe_listener(queue, address):
671 sub_debug('closing listener with address=%r', address)
672 for handle in queue:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200673 _winapi.CloseHandle(handle)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000674
675 def PipeClient(address):
676 '''
677 Return a connection object connected to the pipe given by `address`
678 '''
Antoine Pitrou45d61a32009-11-13 22:35:18 +0000679 t = _init_timeout()
Benjamin Petersone711caf2008-06-11 16:44:04 +0000680 while 1:
681 try:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200682 _winapi.WaitNamedPipe(address, 1000)
683 h = _winapi.CreateFile(
684 address, _winapi.GENERIC_READ | _winapi.GENERIC_WRITE,
685 0, _winapi.NULL, _winapi.OPEN_EXISTING,
686 _winapi.FILE_FLAG_OVERLAPPED, _winapi.NULL
Benjamin Petersone711caf2008-06-11 16:44:04 +0000687 )
Andrew Svetlov2606a6f2012-12-19 14:33:35 +0200688 except OSError as e:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200689 if e.winerror not in (_winapi.ERROR_SEM_TIMEOUT,
690 _winapi.ERROR_PIPE_BUSY) or _check_timeout(t):
Benjamin Petersone711caf2008-06-11 16:44:04 +0000691 raise
692 else:
693 break
694 else:
695 raise
696
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200697 _winapi.SetNamedPipeHandleState(
698 h, _winapi.PIPE_READMODE_MESSAGE, None, None
Benjamin Petersone711caf2008-06-11 16:44:04 +0000699 )
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200700 return PipeConnection(h)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000701
702#
703# Authentication stuff
704#
705
706MESSAGE_LENGTH = 20
707
Benjamin Peterson1fcfe212008-06-25 12:54:22 +0000708CHALLENGE = b'#CHALLENGE#'
709WELCOME = b'#WELCOME#'
710FAILURE = b'#FAILURE#'
Benjamin Petersone711caf2008-06-11 16:44:04 +0000711
712def deliver_challenge(connection, authkey):
713 import hmac
714 assert isinstance(authkey, bytes)
715 message = os.urandom(MESSAGE_LENGTH)
716 connection.send_bytes(CHALLENGE + message)
717 digest = hmac.new(authkey, message).digest()
718 response = connection.recv_bytes(256) # reject large message
719 if response == digest:
720 connection.send_bytes(WELCOME)
721 else:
722 connection.send_bytes(FAILURE)
723 raise AuthenticationError('digest received was wrong')
724
725def answer_challenge(connection, authkey):
726 import hmac
727 assert isinstance(authkey, bytes)
728 message = connection.recv_bytes(256) # reject large message
729 assert message[:len(CHALLENGE)] == CHALLENGE, 'message = %r' % message
730 message = message[len(CHALLENGE):]
731 digest = hmac.new(authkey, message).digest()
732 connection.send_bytes(digest)
733 response = connection.recv_bytes(256) # reject large message
734 if response != WELCOME:
735 raise AuthenticationError('digest sent was rejected')
736
737#
738# Support for using xmlrpclib for serialization
739#
740
741class ConnectionWrapper(object):
742 def __init__(self, conn, dumps, loads):
743 self._conn = conn
744 self._dumps = dumps
745 self._loads = loads
746 for attr in ('fileno', 'close', 'poll', 'recv_bytes', 'send_bytes'):
747 obj = getattr(conn, attr)
748 setattr(self, attr, obj)
749 def send(self, obj):
750 s = self._dumps(obj)
751 self._conn.send_bytes(s)
752 def recv(self):
753 s = self._conn.recv_bytes()
754 return self._loads(s)
755
756def _xml_dumps(obj):
Marc-André Lemburg8f36af72011-02-25 15:42:01 +0000757 return xmlrpclib.dumps((obj,), None, None, None, 1).encode('utf-8')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000758
759def _xml_loads(s):
Marc-André Lemburg8f36af72011-02-25 15:42:01 +0000760 (obj,), method = xmlrpclib.loads(s.decode('utf-8'))
Benjamin Petersone711caf2008-06-11 16:44:04 +0000761 return obj
762
763class XmlListener(Listener):
764 def accept(self):
765 global xmlrpclib
766 import xmlrpc.client as xmlrpclib
767 obj = Listener.accept(self)
768 return ConnectionWrapper(obj, _xml_dumps, _xml_loads)
769
770def XmlClient(*args, **kwds):
771 global xmlrpclib
772 import xmlrpc.client as xmlrpclib
773 return ConnectionWrapper(Client(*args, **kwds), _xml_dumps, _xml_loads)
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200774
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100775#
776# Wait
777#
778
779if sys.platform == 'win32':
780
781 def _exhaustive_wait(handles, timeout):
782 # Return ALL handles which are currently signalled. (Only
783 # returning the first signalled might create starvation issues.)
784 L = list(handles)
785 ready = []
786 while L:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200787 res = _winapi.WaitForMultipleObjects(L, False, timeout)
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100788 if res == WAIT_TIMEOUT:
789 break
790 elif WAIT_OBJECT_0 <= res < WAIT_OBJECT_0 + len(L):
791 res -= WAIT_OBJECT_0
792 elif WAIT_ABANDONED_0 <= res < WAIT_ABANDONED_0 + len(L):
793 res -= WAIT_ABANDONED_0
794 else:
795 raise RuntimeError('Should not get here')
796 ready.append(L[res])
797 L = L[res+1:]
798 timeout = 0
799 return ready
800
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200801 _ready_errors = {_winapi.ERROR_BROKEN_PIPE, _winapi.ERROR_NETNAME_DELETED}
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100802
803 def wait(object_list, timeout=None):
804 '''
805 Wait till an object in object_list is ready/readable.
806
807 Returns list of those objects in object_list which are ready/readable.
808 '''
809 if timeout is None:
810 timeout = INFINITE
811 elif timeout < 0:
812 timeout = 0
813 else:
814 timeout = int(timeout * 1000 + 0.5)
815
816 object_list = list(object_list)
817 waithandle_to_obj = {}
818 ov_list = []
819 ready_objects = set()
820 ready_handles = set()
821
822 try:
823 for o in object_list:
824 try:
825 fileno = getattr(o, 'fileno')
826 except AttributeError:
827 waithandle_to_obj[o.__index__()] = o
828 else:
829 # start an overlapped read of length zero
830 try:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200831 ov, err = _winapi.ReadFile(fileno(), 0, True)
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100832 except OSError as e:
833 err = e.winerror
834 if err not in _ready_errors:
835 raise
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200836 if err == _winapi.ERROR_IO_PENDING:
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100837 ov_list.append(ov)
838 waithandle_to_obj[ov.event] = o
839 else:
840 # If o.fileno() is an overlapped pipe handle and
841 # err == 0 then there is a zero length message
842 # in the pipe, but it HAS NOT been consumed.
843 ready_objects.add(o)
844 timeout = 0
845
846 ready_handles = _exhaustive_wait(waithandle_to_obj.keys(), timeout)
847 finally:
848 # request that overlapped reads stop
849 for ov in ov_list:
850 ov.cancel()
851
852 # wait for all overlapped reads to stop
853 for ov in ov_list:
854 try:
855 _, err = ov.GetOverlappedResult(True)
856 except OSError as e:
857 err = e.winerror
858 if err not in _ready_errors:
859 raise
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200860 if err != _winapi.ERROR_OPERATION_ABORTED:
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100861 o = waithandle_to_obj[ov.event]
862 ready_objects.add(o)
863 if err == 0:
864 # If o.fileno() is an overlapped pipe handle then
865 # a zero length message HAS been consumed.
866 if hasattr(o, '_got_empty_message'):
867 o._got_empty_message = True
868
869 ready_objects.update(waithandle_to_obj[h] for h in ready_handles)
870 return [o for o in object_list if o in ready_objects]
871
872else:
873
Giampaolo Rodola'0c8ad612013-01-14 02:24:05 +0100874 if hasattr(select, 'poll'):
875 def _poll(fds, timeout):
876 if timeout is not None:
Giampaolo Rodola'30830712013-04-17 13:12:27 +0200877 timeout = int(timeout * 1000) # timeout is in milliseconds
Giampaolo Rodola'0c8ad612013-01-14 02:24:05 +0100878 fd_map = {}
879 pollster = select.poll()
880 for fd in fds:
881 pollster.register(fd, select.POLLIN)
882 if hasattr(fd, 'fileno'):
883 fd_map[fd.fileno()] = fd
884 else:
885 fd_map[fd] = fd
886 ls = []
887 for fd, event in pollster.poll(timeout):
888 if event & select.POLLNVAL:
889 raise ValueError('invalid file descriptor %i' % fd)
890 ls.append(fd_map[fd])
891 return ls
892 else:
893 def _poll(fds, timeout):
894 return select.select(fds, [], [], timeout)[0]
895
896
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100897 def wait(object_list, timeout=None):
898 '''
899 Wait till an object in object_list is ready/readable.
900
901 Returns list of those objects in object_list which are ready/readable.
902 '''
903 if timeout is not None:
904 if timeout <= 0:
Giampaolo Rodola'0c8ad612013-01-14 02:24:05 +0100905 return _poll(object_list, 0)
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100906 else:
907 deadline = time.time() + timeout
908 while True:
909 try:
Giampaolo Rodola'0c8ad612013-01-14 02:24:05 +0100910 return _poll(object_list, timeout)
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100911 except OSError as e:
912 if e.errno != errno.EINTR:
913 raise
914 if timeout is not None:
915 timeout = deadline - time.time()
Antoine Pitrou5438ed12012-04-24 22:56:57 +0200916
917#
918# Make connection and socket objects sharable if possible
919#
920
921if sys.platform == 'win32':
922 from . import reduction
923 ForkingPickler.register(socket.socket, reduction.reduce_socket)
924 ForkingPickler.register(Connection, reduction.reduce_connection)
925 ForkingPickler.register(PipeConnection, reduction.reduce_pipe_connection)
926else:
927 try:
928 from . import reduction
929 except ImportError:
930 pass
931 else:
932 ForkingPickler.register(socket.socket, reduction.reduce_socket)
933 ForkingPickler.register(Connection, reduction.reduce_connection)