blob: 22589d0422cf7aace53b30cc8f3c31348bbbd9c9 [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 pickle
16import select
Benjamin Petersone711caf2008-06-11 16:44:04 +000017import socket
Antoine Pitrou87cf2202011-05-09 17:04:27 +020018import struct
Georg Brandl6aa2d1f2008-08-12 08:35:52 +000019import errno
Benjamin Petersone711caf2008-06-11 16:44:04 +000020import time
21import tempfile
22import itertools
23
24import _multiprocessing
Antoine Pitrou87cf2202011-05-09 17:04:27 +020025from multiprocessing import current_process, AuthenticationError, BufferTooShort
Richard Oudkerk59d54042012-05-10 16:11:12 +010026from multiprocessing.util import get_temp_dir, Finalize, sub_debug, debug
Antoine Pitrou5438ed12012-04-24 22:56:57 +020027from multiprocessing.forking import ForkingPickler
Antoine Pitrou87cf2202011-05-09 17:04:27 +020028try:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +020029 import _winapi
30 from _winapi import WAIT_OBJECT_0, WAIT_TIMEOUT, INFINITE
Antoine Pitrou87cf2202011-05-09 17:04:27 +020031except ImportError:
32 if sys.platform == 'win32':
33 raise
Antoine Pitrou23bba4c2012-04-18 20:51:15 +020034 _winapi = None
Benjamin Petersone711caf2008-06-11 16:44:04 +000035
36#
37#
38#
39
40BUFSIZE = 8192
Antoine Pitrou45d61a32009-11-13 22:35:18 +000041# A very generous timeout when it comes to local connections...
42CONNECTION_TIMEOUT = 20.
Benjamin Petersone711caf2008-06-11 16:44:04 +000043
44_mmap_counter = itertools.count()
45
46default_family = 'AF_INET'
47families = ['AF_INET']
48
49if hasattr(socket, 'AF_UNIX'):
50 default_family = 'AF_UNIX'
51 families += ['AF_UNIX']
52
53if sys.platform == 'win32':
54 default_family = 'AF_PIPE'
55 families += ['AF_PIPE']
56
Antoine Pitrou45d61a32009-11-13 22:35:18 +000057
58def _init_timeout(timeout=CONNECTION_TIMEOUT):
59 return time.time() + timeout
60
61def _check_timeout(t):
62 return time.time() > t
63
Benjamin Petersone711caf2008-06-11 16:44:04 +000064#
65#
66#
67
68def 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 Pitrou709176f2012-04-01 17:19:09 +020082def _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 Pitrou6d20cba2012-04-03 20:12:23 +020089 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 Petersone711caf2008-06-11 16:44:04 +000093
94def 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 Pitrou87cf2202011-05-09 17:04:27 +0200110# Connection classes
111#
112
113class _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 Pitrou60001202011-07-09 01:03:46 +0200127 # XXX should we use util.Finalize instead of a __del__?
128
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200129 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 Pitrou5438ed12012-04-24 22:56:57 +0200205 buf = io.BytesIO()
206 ForkingPickler(buf, pickle.HIGHEST_PROTOCOL).dump(obj)
207 self._send_bytes(buf.getbuffer())
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200208
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 Pitroubdb1cf12012-03-05 19:28:37 +0100247 def recv(self):
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200248 """Receive a (picklable) object"""
249 self._check_closed()
250 self._check_readable()
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100251 buf = self._recv_bytes()
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200252 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 Pitrou87cf2202011-05-09 17:04:27 +0200258 return self._poll(timeout)
259
Richard Oudkerkd69cfe82012-06-18 17:47:52 +0100260 def __enter__(self):
261 return self
262
263 def __exit__(self, exc_type, exc_value, exc_tb):
264 self.close()
265
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200266
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200267if _winapi:
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200268
269 class PipeConnection(_ConnectionBase):
270 """
271 Connection class based on a Windows named pipe.
Antoine Pitroudd696492011-06-08 17:21:55 +0200272 Overlapped I/O is used, so the handles must have been created
273 with FILE_FLAG_OVERLAPPED.
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200274 """
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100275 _got_empty_message = False
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200276
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200277 def _close(self, _CloseHandle=_winapi.CloseHandle):
Antoine Pitrou71a28a92011-07-09 01:03:00 +0200278 _CloseHandle(self._handle)
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200279
280 def _send_bytes(self, buf):
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200281 ov, err = _winapi.WriteFile(self._handle, buf, overlapped=True)
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100282 try:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200283 if err == _winapi.ERROR_IO_PENDING:
284 waitres = _winapi.WaitForMultipleObjects(
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100285 [ov.event], False, INFINITE)
286 assert waitres == WAIT_OBJECT_0
287 except:
288 ov.cancel()
289 raise
290 finally:
291 nwritten, err = ov.GetOverlappedResult(True)
292 assert err == 0
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200293 assert nwritten == len(buf)
294
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100295 def _recv_bytes(self, maxsize=None):
296 if self._got_empty_message:
297 self._got_empty_message = False
298 return io.BytesIO()
Antoine Pitroudd696492011-06-08 17:21:55 +0200299 else:
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100300 bsize = 128 if maxsize is None else min(maxsize, 128)
Antoine Pitroudd696492011-06-08 17:21:55 +0200301 try:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200302 ov, err = _winapi.ReadFile(self._handle, bsize,
303 overlapped=True)
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100304 try:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200305 if err == _winapi.ERROR_IO_PENDING:
306 waitres = _winapi.WaitForMultipleObjects(
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100307 [ov.event], False, INFINITE)
308 assert waitres == WAIT_OBJECT_0
309 except:
310 ov.cancel()
311 raise
312 finally:
313 nread, err = ov.GetOverlappedResult(True)
314 if err == 0:
315 f = io.BytesIO()
316 f.write(ov.getbuffer())
317 return f
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200318 elif err == _winapi.ERROR_MORE_DATA:
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100319 return self._get_more_data(ov, maxsize)
Antoine Pitroudd696492011-06-08 17:21:55 +0200320 except IOError as e:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200321 if e.winerror == _winapi.ERROR_BROKEN_PIPE:
Antoine Pitroudd696492011-06-08 17:21:55 +0200322 raise EOFError
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100323 else:
324 raise
325 raise RuntimeError("shouldn't get here; expected KeyboardInterrupt")
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200326
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100327 def _poll(self, timeout):
328 if (self._got_empty_message or
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200329 _winapi.PeekNamedPipe(self._handle)[0] != 0):
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200330 return True
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100331 return bool(wait([self], timeout))
332
333 def _get_more_data(self, ov, maxsize):
334 buf = ov.getbuffer()
335 f = io.BytesIO()
336 f.write(buf)
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200337 left = _winapi.PeekNamedPipe(self._handle)[1]
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100338 assert left > 0
339 if maxsize is not None and len(buf) + left > maxsize:
340 self._bad_message_length()
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200341 ov, err = _winapi.ReadFile(self._handle, left, overlapped=True)
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100342 rbytes, err = ov.GetOverlappedResult(True)
343 assert err == 0
344 assert rbytes == left
345 f.write(ov.getbuffer())
346 return f
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200347
348
349class Connection(_ConnectionBase):
350 """
351 Connection class based on an arbitrary file descriptor (Unix only), or
352 a socket handle (Windows).
353 """
354
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200355 if _winapi:
356 def _close(self, _close=_multiprocessing.closesocket):
Antoine Pitrou71a28a92011-07-09 01:03:00 +0200357 _close(self._handle)
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200358 _write = _multiprocessing.send
359 _read = _multiprocessing.recv
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200360 else:
Antoine Pitrou71a28a92011-07-09 01:03:00 +0200361 def _close(self, _close=os.close):
362 _close(self._handle)
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200363 _write = os.write
364 _read = os.read
365
366 def _send(self, buf, write=_write):
367 remaining = len(buf)
368 while True:
Richard Oudkerkcca8c532013-07-01 18:59:26 +0100369 try:
370 n = write(self._handle, buf)
371 except InterruptedError:
372 continue
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200373 remaining -= n
374 if remaining == 0:
375 break
376 buf = buf[n:]
377
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100378 def _recv(self, size, read=_read):
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200379 buf = io.BytesIO()
Antoine Pitroudd696492011-06-08 17:21:55 +0200380 handle = self._handle
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200381 remaining = size
382 while remaining > 0:
Richard Oudkerkcca8c532013-07-01 18:59:26 +0100383 try:
384 chunk = read(handle, remaining)
385 except InterruptedError:
386 continue
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200387 n = len(chunk)
388 if n == 0:
389 if remaining == size:
390 raise EOFError
391 else:
392 raise IOError("got end of file during message")
393 buf.write(chunk)
394 remaining -= n
395 return buf
396
397 def _send_bytes(self, buf):
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200398 n = len(buf)
Antoine Pitroub7d6d2a2014-02-08 23:03:56 +0100399 # For wire compatibility with 3.2 and lower
400 header = struct.pack("!i", n)
401 if n > 16384:
402 # The payload is large so Nagle's algorithm won't be triggered
403 # and we'd better avoid the cost of concatenation.
404 chunks = [header, buf]
405 elif n > 0:
406 # Issue # 20540: concatenate before sending, to avoid delays due
407 # to Nagle's algorithm on a TCP socket.
408 chunks = [header + buf]
409 else:
410 # This code path is necessary to avoid "broken pipe" errors
411 # when sending a 0-length buffer if the other end closed the pipe.
412 chunks = [header]
413 for chunk in chunks:
414 self._send(chunk)
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200415
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100416 def _recv_bytes(self, maxsize=None):
417 buf = self._recv(4)
Charles-François Natali225aa4f2011-09-20 19:27:39 +0200418 size, = struct.unpack("!i", buf.getvalue())
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200419 if maxsize is not None and size > maxsize:
420 return None
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100421 return self._recv(size)
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200422
423 def _poll(self, timeout):
Richard Oudkerked9e06c2013-01-13 22:46:48 +0000424 r = wait([self], timeout)
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200425 return bool(r)
426
427
428#
Benjamin Petersone711caf2008-06-11 16:44:04 +0000429# Public functions
430#
431
432class Listener(object):
433 '''
434 Returns a listener object.
435
436 This is a wrapper for a bound socket which is 'listening' for
437 connections, or for a Windows named pipe.
438 '''
439 def __init__(self, address=None, family=None, backlog=1, authkey=None):
440 family = family or (address and address_type(address)) \
441 or default_family
442 address = address or arbitrary_address(family)
443
Antoine Pitrou709176f2012-04-01 17:19:09 +0200444 _validate_family(family)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000445 if family == 'AF_PIPE':
446 self._listener = PipeListener(address, backlog)
447 else:
448 self._listener = SocketListener(address, family, backlog)
449
450 if authkey is not None and not isinstance(authkey, bytes):
451 raise TypeError('authkey should be a byte string')
452
453 self._authkey = authkey
454
455 def accept(self):
456 '''
457 Accept a connection on the bound socket or named pipe of `self`.
458
459 Returns a `Connection` object.
460 '''
Richard Oudkerkd69cfe82012-06-18 17:47:52 +0100461 if self._listener is None:
462 raise IOError('listener is closed')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000463 c = self._listener.accept()
464 if self._authkey:
465 deliver_challenge(c, self._authkey)
466 answer_challenge(c, self._authkey)
467 return c
468
469 def close(self):
470 '''
471 Close the bound socket or named pipe of `self`.
472 '''
Richard Oudkerkd69cfe82012-06-18 17:47:52 +0100473 if self._listener is not None:
474 self._listener.close()
475 self._listener = None
Benjamin Petersone711caf2008-06-11 16:44:04 +0000476
477 address = property(lambda self: self._listener._address)
478 last_accepted = property(lambda self: self._listener._last_accepted)
479
Richard Oudkerkd69cfe82012-06-18 17:47:52 +0100480 def __enter__(self):
481 return self
482
483 def __exit__(self, exc_type, exc_value, exc_tb):
484 self.close()
485
Benjamin Petersone711caf2008-06-11 16:44:04 +0000486
487def Client(address, family=None, authkey=None):
488 '''
489 Returns a connection to the address of a `Listener`
490 '''
491 family = family or address_type(address)
Antoine Pitrou709176f2012-04-01 17:19:09 +0200492 _validate_family(family)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000493 if family == 'AF_PIPE':
494 c = PipeClient(address)
495 else:
496 c = SocketClient(address)
497
498 if authkey is not None and not isinstance(authkey, bytes):
499 raise TypeError('authkey should be a byte string')
500
501 if authkey is not None:
502 answer_challenge(c, authkey)
503 deliver_challenge(c, authkey)
504
505 return c
506
507
508if sys.platform != 'win32':
509
510 def Pipe(duplex=True):
511 '''
512 Returns pair of connection objects at either end of a pipe
513 '''
514 if duplex:
515 s1, s2 = socket.socketpair()
Richard Oudkerkb15e6222012-07-27 14:19:00 +0100516 s1.setblocking(True)
517 s2.setblocking(True)
Antoine Pitrou5aa878c2011-05-09 21:00:28 +0200518 c1 = Connection(s1.detach())
519 c2 = Connection(s2.detach())
Benjamin Petersone711caf2008-06-11 16:44:04 +0000520 else:
521 fd1, fd2 = os.pipe()
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200522 c1 = Connection(fd1, writable=False)
523 c2 = Connection(fd2, readable=False)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000524
525 return c1, c2
526
527else:
528
Benjamin Petersone711caf2008-06-11 16:44:04 +0000529 def Pipe(duplex=True):
530 '''
531 Returns pair of connection objects at either end of a pipe
532 '''
533 address = arbitrary_address('AF_PIPE')
534 if duplex:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200535 openmode = _winapi.PIPE_ACCESS_DUPLEX
536 access = _winapi.GENERIC_READ | _winapi.GENERIC_WRITE
Benjamin Petersone711caf2008-06-11 16:44:04 +0000537 obsize, ibsize = BUFSIZE, BUFSIZE
538 else:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200539 openmode = _winapi.PIPE_ACCESS_INBOUND
540 access = _winapi.GENERIC_WRITE
Benjamin Petersone711caf2008-06-11 16:44:04 +0000541 obsize, ibsize = 0, BUFSIZE
542
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200543 h1 = _winapi.CreateNamedPipe(
544 address, openmode | _winapi.FILE_FLAG_OVERLAPPED |
545 _winapi.FILE_FLAG_FIRST_PIPE_INSTANCE,
546 _winapi.PIPE_TYPE_MESSAGE | _winapi.PIPE_READMODE_MESSAGE |
547 _winapi.PIPE_WAIT,
548 1, obsize, ibsize, _winapi.NMPWAIT_WAIT_FOREVER, _winapi.NULL
Benjamin Petersone711caf2008-06-11 16:44:04 +0000549 )
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200550 h2 = _winapi.CreateFile(
551 address, access, 0, _winapi.NULL, _winapi.OPEN_EXISTING,
552 _winapi.FILE_FLAG_OVERLAPPED, _winapi.NULL
Benjamin Petersone711caf2008-06-11 16:44:04 +0000553 )
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200554 _winapi.SetNamedPipeHandleState(
555 h2, _winapi.PIPE_READMODE_MESSAGE, None, None
Benjamin Petersone711caf2008-06-11 16:44:04 +0000556 )
557
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200558 overlapped = _winapi.ConnectNamedPipe(h1, overlapped=True)
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100559 _, err = overlapped.GetOverlappedResult(True)
560 assert err == 0
Benjamin Petersone711caf2008-06-11 16:44:04 +0000561
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200562 c1 = PipeConnection(h1, writable=duplex)
563 c2 = PipeConnection(h2, readable=duplex)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000564
565 return c1, c2
566
567#
568# Definitions for connections based on sockets
569#
570
571class SocketListener(object):
572 '''
Georg Brandl734e2682008-08-12 08:18:18 +0000573 Representation of a socket which is bound to an address and listening
Benjamin Petersone711caf2008-06-11 16:44:04 +0000574 '''
575 def __init__(self, address, family, backlog=1):
576 self._socket = socket.socket(getattr(socket, family))
Charles-François Nataliedc67fe2012-02-04 15:12:08 +0100577 try:
Charles-François Natalied4a8fc2012-02-08 21:15:58 +0100578 # SO_REUSEADDR has different semantics on Windows (issue #2550).
579 if os.name == 'posix':
580 self._socket.setsockopt(socket.SOL_SOCKET,
581 socket.SO_REUSEADDR, 1)
Richard Oudkerkb15e6222012-07-27 14:19:00 +0100582 self._socket.setblocking(True)
Charles-François Nataliedc67fe2012-02-04 15:12:08 +0100583 self._socket.bind(address)
584 self._socket.listen(backlog)
585 self._address = self._socket.getsockname()
586 except OSError:
587 self._socket.close()
588 raise
Benjamin Petersone711caf2008-06-11 16:44:04 +0000589 self._family = family
590 self._last_accepted = None
591
Benjamin Petersone711caf2008-06-11 16:44:04 +0000592 if family == 'AF_UNIX':
593 self._unlink = Finalize(
Georg Brandl2ee470f2008-07-16 12:55:28 +0000594 self, os.unlink, args=(address,), exitpriority=0
Benjamin Petersone711caf2008-06-11 16:44:04 +0000595 )
596 else:
597 self._unlink = None
598
599 def accept(self):
Richard Oudkerkcca8c532013-07-01 18:59:26 +0100600 while True:
601 try:
602 s, self._last_accepted = self._socket.accept()
603 except InterruptedError:
604 pass
605 else:
606 break
Richard Oudkerkb15e6222012-07-27 14:19:00 +0100607 s.setblocking(True)
Antoine Pitroudf97cbe2012-04-07 22:38:52 +0200608 return Connection(s.detach())
Benjamin Petersone711caf2008-06-11 16:44:04 +0000609
610 def close(self):
611 self._socket.close()
612 if self._unlink is not None:
613 self._unlink()
614
615
616def SocketClient(address):
617 '''
618 Return a connection object connected to the socket given by `address`
619 '''
620 family = address_type(address)
Victor Stinner2b695062011-01-03 15:47:59 +0000621 with socket.socket( getattr(socket, family) ) as s:
Richard Oudkerkb15e6222012-07-27 14:19:00 +0100622 s.setblocking(True)
Charles-François Natalie6eabd42011-11-19 09:59:43 +0100623 s.connect(address)
Antoine Pitroudf97cbe2012-04-07 22:38:52 +0200624 return Connection(s.detach())
Benjamin Petersone711caf2008-06-11 16:44:04 +0000625
626#
627# Definitions for connections based on named pipes
628#
629
630if sys.platform == 'win32':
631
632 class PipeListener(object):
633 '''
634 Representation of a named pipe
635 '''
636 def __init__(self, address, backlog=None):
637 self._address = address
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100638 self._handle_queue = [self._new_handle(first=True)]
639
Benjamin Petersone711caf2008-06-11 16:44:04 +0000640 self._last_accepted = None
Benjamin Petersone711caf2008-06-11 16:44:04 +0000641 sub_debug('listener created with address=%r', self._address)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000642 self.close = Finalize(
643 self, PipeListener._finalize_pipe_listener,
644 args=(self._handle_queue, self._address), exitpriority=0
645 )
646
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100647 def _new_handle(self, first=False):
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200648 flags = _winapi.PIPE_ACCESS_DUPLEX | _winapi.FILE_FLAG_OVERLAPPED
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100649 if first:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200650 flags |= _winapi.FILE_FLAG_FIRST_PIPE_INSTANCE
651 return _winapi.CreateNamedPipe(
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100652 self._address, flags,
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200653 _winapi.PIPE_TYPE_MESSAGE | _winapi.PIPE_READMODE_MESSAGE |
654 _winapi.PIPE_WAIT,
655 _winapi.PIPE_UNLIMITED_INSTANCES, BUFSIZE, BUFSIZE,
656 _winapi.NMPWAIT_WAIT_FOREVER, _winapi.NULL
Benjamin Petersone711caf2008-06-11 16:44:04 +0000657 )
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100658
659 def accept(self):
660 self._handle_queue.append(self._new_handle())
Benjamin Petersone711caf2008-06-11 16:44:04 +0000661 handle = self._handle_queue.pop(0)
662 try:
Richard Oudkerkfdb8dcf2012-05-05 19:45:37 +0100663 ov = _winapi.ConnectNamedPipe(handle, overlapped=True)
664 except OSError as e:
665 if e.winerror != _winapi.ERROR_NO_DATA:
666 raise
667 # ERROR_NO_DATA can occur if a client has already connected,
668 # written data and then disconnected -- see Issue 14725.
669 else:
670 try:
671 res = _winapi.WaitForMultipleObjects(
672 [ov.event], False, INFINITE)
673 except:
674 ov.cancel()
675 _winapi.CloseHandle(handle)
676 raise
677 finally:
678 _, err = ov.GetOverlappedResult(True)
679 assert err == 0
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200680 return PipeConnection(handle)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000681
682 @staticmethod
683 def _finalize_pipe_listener(queue, address):
684 sub_debug('closing listener with address=%r', address)
685 for handle in queue:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200686 _winapi.CloseHandle(handle)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000687
688 def PipeClient(address):
689 '''
690 Return a connection object connected to the pipe given by `address`
691 '''
Antoine Pitrou45d61a32009-11-13 22:35:18 +0000692 t = _init_timeout()
Benjamin Petersone711caf2008-06-11 16:44:04 +0000693 while 1:
694 try:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200695 _winapi.WaitNamedPipe(address, 1000)
696 h = _winapi.CreateFile(
697 address, _winapi.GENERIC_READ | _winapi.GENERIC_WRITE,
698 0, _winapi.NULL, _winapi.OPEN_EXISTING,
699 _winapi.FILE_FLAG_OVERLAPPED, _winapi.NULL
Benjamin Petersone711caf2008-06-11 16:44:04 +0000700 )
701 except WindowsError as e:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200702 if e.winerror not in (_winapi.ERROR_SEM_TIMEOUT,
703 _winapi.ERROR_PIPE_BUSY) or _check_timeout(t):
Benjamin Petersone711caf2008-06-11 16:44:04 +0000704 raise
705 else:
706 break
707 else:
708 raise
709
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200710 _winapi.SetNamedPipeHandleState(
711 h, _winapi.PIPE_READMODE_MESSAGE, None, None
Benjamin Petersone711caf2008-06-11 16:44:04 +0000712 )
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200713 return PipeConnection(h)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000714
715#
716# Authentication stuff
717#
718
719MESSAGE_LENGTH = 20
720
Benjamin Peterson1fcfe212008-06-25 12:54:22 +0000721CHALLENGE = b'#CHALLENGE#'
722WELCOME = b'#WELCOME#'
723FAILURE = b'#FAILURE#'
Benjamin Petersone711caf2008-06-11 16:44:04 +0000724
725def deliver_challenge(connection, authkey):
726 import hmac
727 assert isinstance(authkey, bytes)
728 message = os.urandom(MESSAGE_LENGTH)
729 connection.send_bytes(CHALLENGE + message)
730 digest = hmac.new(authkey, message).digest()
731 response = connection.recv_bytes(256) # reject large message
732 if response == digest:
733 connection.send_bytes(WELCOME)
734 else:
735 connection.send_bytes(FAILURE)
736 raise AuthenticationError('digest received was wrong')
737
738def answer_challenge(connection, authkey):
739 import hmac
740 assert isinstance(authkey, bytes)
741 message = connection.recv_bytes(256) # reject large message
742 assert message[:len(CHALLENGE)] == CHALLENGE, 'message = %r' % message
743 message = message[len(CHALLENGE):]
744 digest = hmac.new(authkey, message).digest()
745 connection.send_bytes(digest)
746 response = connection.recv_bytes(256) # reject large message
747 if response != WELCOME:
748 raise AuthenticationError('digest sent was rejected')
749
750#
751# Support for using xmlrpclib for serialization
752#
753
754class ConnectionWrapper(object):
755 def __init__(self, conn, dumps, loads):
756 self._conn = conn
757 self._dumps = dumps
758 self._loads = loads
759 for attr in ('fileno', 'close', 'poll', 'recv_bytes', 'send_bytes'):
760 obj = getattr(conn, attr)
761 setattr(self, attr, obj)
762 def send(self, obj):
763 s = self._dumps(obj)
764 self._conn.send_bytes(s)
765 def recv(self):
766 s = self._conn.recv_bytes()
767 return self._loads(s)
768
769def _xml_dumps(obj):
Marc-André Lemburg8f36af72011-02-25 15:42:01 +0000770 return xmlrpclib.dumps((obj,), None, None, None, 1).encode('utf-8')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000771
772def _xml_loads(s):
Marc-André Lemburg8f36af72011-02-25 15:42:01 +0000773 (obj,), method = xmlrpclib.loads(s.decode('utf-8'))
Benjamin Petersone711caf2008-06-11 16:44:04 +0000774 return obj
775
776class XmlListener(Listener):
777 def accept(self):
778 global xmlrpclib
779 import xmlrpc.client as xmlrpclib
780 obj = Listener.accept(self)
781 return ConnectionWrapper(obj, _xml_dumps, _xml_loads)
782
783def XmlClient(*args, **kwds):
784 global xmlrpclib
785 import xmlrpc.client as xmlrpclib
786 return ConnectionWrapper(Client(*args, **kwds), _xml_dumps, _xml_loads)
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200787
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100788#
789# Wait
790#
791
792if sys.platform == 'win32':
793
794 def _exhaustive_wait(handles, timeout):
795 # Return ALL handles which are currently signalled. (Only
796 # returning the first signalled might create starvation issues.)
797 L = list(handles)
798 ready = []
799 while L:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200800 res = _winapi.WaitForMultipleObjects(L, False, timeout)
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100801 if res == WAIT_TIMEOUT:
802 break
803 elif WAIT_OBJECT_0 <= res < WAIT_OBJECT_0 + len(L):
804 res -= WAIT_OBJECT_0
805 elif WAIT_ABANDONED_0 <= res < WAIT_ABANDONED_0 + len(L):
806 res -= WAIT_ABANDONED_0
807 else:
808 raise RuntimeError('Should not get here')
809 ready.append(L[res])
810 L = L[res+1:]
811 timeout = 0
812 return ready
813
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200814 _ready_errors = {_winapi.ERROR_BROKEN_PIPE, _winapi.ERROR_NETNAME_DELETED}
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100815
816 def wait(object_list, timeout=None):
817 '''
818 Wait till an object in object_list is ready/readable.
819
820 Returns list of those objects in object_list which are ready/readable.
821 '''
822 if timeout is None:
823 timeout = INFINITE
824 elif timeout < 0:
825 timeout = 0
826 else:
827 timeout = int(timeout * 1000 + 0.5)
828
829 object_list = list(object_list)
830 waithandle_to_obj = {}
831 ov_list = []
832 ready_objects = set()
833 ready_handles = set()
834
835 try:
836 for o in object_list:
837 try:
838 fileno = getattr(o, 'fileno')
839 except AttributeError:
840 waithandle_to_obj[o.__index__()] = o
841 else:
842 # start an overlapped read of length zero
843 try:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200844 ov, err = _winapi.ReadFile(fileno(), 0, True)
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100845 except OSError as e:
846 err = e.winerror
847 if err not in _ready_errors:
848 raise
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200849 if err == _winapi.ERROR_IO_PENDING:
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100850 ov_list.append(ov)
851 waithandle_to_obj[ov.event] = o
852 else:
853 # If o.fileno() is an overlapped pipe handle and
854 # err == 0 then there is a zero length message
855 # in the pipe, but it HAS NOT been consumed.
856 ready_objects.add(o)
857 timeout = 0
858
859 ready_handles = _exhaustive_wait(waithandle_to_obj.keys(), timeout)
860 finally:
861 # request that overlapped reads stop
862 for ov in ov_list:
863 ov.cancel()
864
865 # wait for all overlapped reads to stop
866 for ov in ov_list:
867 try:
868 _, err = ov.GetOverlappedResult(True)
869 except OSError as e:
870 err = e.winerror
871 if err not in _ready_errors:
872 raise
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200873 if err != _winapi.ERROR_OPERATION_ABORTED:
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100874 o = waithandle_to_obj[ov.event]
875 ready_objects.add(o)
876 if err == 0:
877 # If o.fileno() is an overlapped pipe handle then
878 # a zero length message HAS been consumed.
879 if hasattr(o, '_got_empty_message'):
880 o._got_empty_message = True
881
882 ready_objects.update(waithandle_to_obj[h] for h in ready_handles)
883 return [o for o in object_list if o in ready_objects]
884
885else:
886
Giampaolo Rodola'67da8942013-01-14 02:24:25 +0100887 if hasattr(select, 'poll'):
888 def _poll(fds, timeout):
889 if timeout is not None:
Giampaolo Rodola'b38897f2013-04-17 13:08:59 +0200890 timeout = int(timeout * 1000) # timeout is in milliseconds
Giampaolo Rodola'67da8942013-01-14 02:24:25 +0100891 fd_map = {}
892 pollster = select.poll()
893 for fd in fds:
894 pollster.register(fd, select.POLLIN)
895 if hasattr(fd, 'fileno'):
896 fd_map[fd.fileno()] = fd
897 else:
898 fd_map[fd] = fd
899 ls = []
900 for fd, event in pollster.poll(timeout):
901 if event & select.POLLNVAL:
902 raise ValueError('invalid file descriptor %i' % fd)
903 ls.append(fd_map[fd])
904 return ls
905 else:
906 def _poll(fds, timeout):
907 return select.select(fds, [], [], timeout)[0]
908
909
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100910 def wait(object_list, timeout=None):
911 '''
912 Wait till an object in object_list is ready/readable.
913
914 Returns list of those objects in object_list which are ready/readable.
915 '''
916 if timeout is not None:
917 if timeout <= 0:
Giampaolo Rodola'67da8942013-01-14 02:24:25 +0100918 return _poll(object_list, 0)
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100919 else:
920 deadline = time.time() + timeout
921 while True:
922 try:
Giampaolo Rodola'67da8942013-01-14 02:24:25 +0100923 return _poll(object_list, timeout)
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100924 except OSError as e:
925 if e.errno != errno.EINTR:
926 raise
927 if timeout is not None:
928 timeout = deadline - time.time()
Antoine Pitrou5438ed12012-04-24 22:56:57 +0200929
930#
931# Make connection and socket objects sharable if possible
932#
933
934if sys.platform == 'win32':
935 from . import reduction
936 ForkingPickler.register(socket.socket, reduction.reduce_socket)
937 ForkingPickler.register(Connection, reduction.reduce_connection)
938 ForkingPickler.register(PipeConnection, reduction.reduce_pipe_connection)
939else:
940 try:
941 from . import reduction
942 except ImportError:
943 pass
944 else:
945 ForkingPickler.register(socket.socket, reduction.reduce_socket)
946 ForkingPickler.register(Connection, reduction.reduce_connection)