blob: 27fda9f22c68d18b3bf3cca6b09443e42480f2a8 [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
15import socket
Antoine Pitrou87cf2202011-05-09 17:04:27 +020016import struct
Georg Brandl6aa2d1f2008-08-12 08:35:52 +000017import errno
Benjamin Petersone711caf2008-06-11 16:44:04 +000018import time
19import tempfile
20import itertools
21
22import _multiprocessing
Richard Oudkerk84ed9a62013-08-14 15:35:41 +010023
24from . import reduction
25from . import util
26
27from . import AuthenticationError, BufferTooShort
28from .reduction import ForkingPickler
29
Antoine Pitrou87cf2202011-05-09 17:04:27 +020030try:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +020031 import _winapi
32 from _winapi import WAIT_OBJECT_0, WAIT_TIMEOUT, INFINITE
Brett Cannoncd171c82013-07-04 17:43:24 -040033except ImportError:
Antoine Pitrou87cf2202011-05-09 17:04:27 +020034 if sys.platform == 'win32':
35 raise
Antoine Pitrou23bba4c2012-04-18 20:51:15 +020036 _winapi = None
Benjamin Petersone711caf2008-06-11 16:44:04 +000037
38#
39#
40#
41
42BUFSIZE = 8192
Antoine Pitrou45d61a32009-11-13 22:35:18 +000043# A very generous timeout when it comes to local connections...
44CONNECTION_TIMEOUT = 20.
Benjamin Petersone711caf2008-06-11 16:44:04 +000045
46_mmap_counter = itertools.count()
47
48default_family = 'AF_INET'
49families = ['AF_INET']
50
51if hasattr(socket, 'AF_UNIX'):
52 default_family = 'AF_UNIX'
53 families += ['AF_UNIX']
54
55if sys.platform == 'win32':
56 default_family = 'AF_PIPE'
57 families += ['AF_PIPE']
58
Antoine Pitrou45d61a32009-11-13 22:35:18 +000059
60def _init_timeout(timeout=CONNECTION_TIMEOUT):
61 return time.time() + timeout
62
63def _check_timeout(t):
64 return time.time() > t
65
Benjamin Petersone711caf2008-06-11 16:44:04 +000066#
67#
68#
69
70def arbitrary_address(family):
71 '''
72 Return an arbitrary free address for the given family
73 '''
74 if family == 'AF_INET':
75 return ('localhost', 0)
76 elif family == 'AF_UNIX':
Richard Oudkerk84ed9a62013-08-14 15:35:41 +010077 return tempfile.mktemp(prefix='listener-', dir=util.get_temp_dir())
Benjamin Petersone711caf2008-06-11 16:44:04 +000078 elif family == 'AF_PIPE':
79 return tempfile.mktemp(prefix=r'\\.\pipe\pyc-%d-%d-' %
80 (os.getpid(), next(_mmap_counter)))
81 else:
82 raise ValueError('unrecognized family')
83
Antoine Pitrou709176f2012-04-01 17:19:09 +020084def _validate_family(family):
85 '''
86 Checks if the family is valid for the current environment.
87 '''
88 if sys.platform != 'win32' and family == 'AF_PIPE':
89 raise ValueError('Family %s is not recognized.' % family)
90
Antoine Pitrou6d20cba2012-04-03 20:12:23 +020091 if sys.platform == 'win32' and family == 'AF_UNIX':
92 # double check
93 if not hasattr(socket, family):
94 raise ValueError('Family %s is not recognized.' % family)
Benjamin Petersone711caf2008-06-11 16:44:04 +000095
96def address_type(address):
97 '''
98 Return the types of the address
99
100 This can be 'AF_INET', 'AF_UNIX', or 'AF_PIPE'
101 '''
102 if type(address) == tuple:
103 return 'AF_INET'
104 elif type(address) is str and address.startswith('\\\\'):
105 return 'AF_PIPE'
106 elif type(address) is str:
107 return 'AF_UNIX'
108 else:
109 raise ValueError('address type of %r unrecognized' % address)
110
111#
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200112# Connection classes
113#
114
115class _ConnectionBase:
116 _handle = None
117
118 def __init__(self, handle, readable=True, writable=True):
119 handle = handle.__index__()
120 if handle < 0:
121 raise ValueError("invalid handle")
122 if not readable and not writable:
123 raise ValueError(
124 "at least one of `readable` and `writable` must be True")
125 self._handle = handle
126 self._readable = readable
127 self._writable = writable
128
Antoine Pitrou60001202011-07-09 01:03:46 +0200129 # XXX should we use util.Finalize instead of a __del__?
130
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200131 def __del__(self):
132 if self._handle is not None:
133 self._close()
134
135 def _check_closed(self):
136 if self._handle is None:
Andrew Svetlovf7a17b42012-12-25 16:47:37 +0200137 raise OSError("handle is closed")
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200138
139 def _check_readable(self):
140 if not self._readable:
Andrew Svetlovf7a17b42012-12-25 16:47:37 +0200141 raise OSError("connection is write-only")
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200142
143 def _check_writable(self):
144 if not self._writable:
Andrew Svetlovf7a17b42012-12-25 16:47:37 +0200145 raise OSError("connection is read-only")
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200146
147 def _bad_message_length(self):
148 if self._writable:
149 self._readable = False
150 else:
151 self.close()
Andrew Svetlovf7a17b42012-12-25 16:47:37 +0200152 raise OSError("bad message length")
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200153
154 @property
155 def closed(self):
156 """True if the connection is closed"""
157 return self._handle is None
158
159 @property
160 def readable(self):
161 """True if the connection is readable"""
162 return self._readable
163
164 @property
165 def writable(self):
166 """True if the connection is writable"""
167 return self._writable
168
169 def fileno(self):
170 """File descriptor or handle of the connection"""
171 self._check_closed()
172 return self._handle
173
174 def close(self):
175 """Close the connection"""
176 if self._handle is not None:
177 try:
178 self._close()
179 finally:
180 self._handle = None
181
182 def send_bytes(self, buf, offset=0, size=None):
183 """Send the bytes data from a bytes-like object"""
184 self._check_closed()
185 self._check_writable()
186 m = memoryview(buf)
187 # HACK for byte-indexing of non-bytewise buffers (e.g. array.array)
188 if m.itemsize > 1:
189 m = memoryview(bytes(m))
190 n = len(m)
191 if offset < 0:
192 raise ValueError("offset is negative")
193 if n < offset:
194 raise ValueError("buffer length < offset")
195 if size is None:
196 size = n - offset
197 elif size < 0:
198 raise ValueError("size is negative")
199 elif offset + size > n:
200 raise ValueError("buffer length < offset + size")
201 self._send_bytes(m[offset:offset + size])
202
203 def send(self, obj):
204 """Send a (picklable) object"""
205 self._check_closed()
206 self._check_writable()
Charles-François Natalia6550752013-03-24 15:21:49 +0100207 self._send_bytes(ForkingPickler.dumps(obj))
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()
Charles-François Natalia6550752013-03-24 15:21:49 +0100252 return ForkingPickler.loads(buf.getbuffer())
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200253
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)
Andrew Svetlovf7a17b42012-12-25 16:47:37 +0200320 except OSError 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:
Andrew Svetlovf7a17b42012-12-25 16:47:37 +0200392 raise OSError("got end of file during message")
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200393 buf.write(chunk)
394 remaining -= n
395 return buf
396
397 def _send_bytes(self, buf):
398 # For wire compatibility with 3.2 and lower
399 n = len(buf)
Charles-François Natali225aa4f2011-09-20 19:27:39 +0200400 self._send(struct.pack("!i", n))
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200401 # The condition is necessary to avoid "broken pipe" errors
402 # when sending a 0-length buffer if the other end closed the pipe.
403 if n > 0:
404 self._send(buf)
405
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100406 def _recv_bytes(self, maxsize=None):
407 buf = self._recv(4)
Charles-François Natali225aa4f2011-09-20 19:27:39 +0200408 size, = struct.unpack("!i", buf.getvalue())
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200409 if maxsize is not None and size > maxsize:
410 return None
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100411 return self._recv(size)
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200412
413 def _poll(self, timeout):
Richard Oudkerked9e06c2013-01-13 22:46:48 +0000414 r = wait([self], timeout)
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200415 return bool(r)
416
417
418#
Benjamin Petersone711caf2008-06-11 16:44:04 +0000419# Public functions
420#
421
422class Listener(object):
423 '''
424 Returns a listener object.
425
426 This is a wrapper for a bound socket which is 'listening' for
427 connections, or for a Windows named pipe.
428 '''
429 def __init__(self, address=None, family=None, backlog=1, authkey=None):
430 family = family or (address and address_type(address)) \
431 or default_family
432 address = address or arbitrary_address(family)
433
Antoine Pitrou709176f2012-04-01 17:19:09 +0200434 _validate_family(family)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000435 if family == 'AF_PIPE':
436 self._listener = PipeListener(address, backlog)
437 else:
438 self._listener = SocketListener(address, family, backlog)
439
440 if authkey is not None and not isinstance(authkey, bytes):
441 raise TypeError('authkey should be a byte string')
442
443 self._authkey = authkey
444
445 def accept(self):
446 '''
447 Accept a connection on the bound socket or named pipe of `self`.
448
449 Returns a `Connection` object.
450 '''
Richard Oudkerkd69cfe82012-06-18 17:47:52 +0100451 if self._listener is None:
Andrew Svetlovf7a17b42012-12-25 16:47:37 +0200452 raise OSError('listener is closed')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000453 c = self._listener.accept()
454 if self._authkey:
455 deliver_challenge(c, self._authkey)
456 answer_challenge(c, self._authkey)
457 return c
458
459 def close(self):
460 '''
461 Close the bound socket or named pipe of `self`.
462 '''
Richard Oudkerkd69cfe82012-06-18 17:47:52 +0100463 if self._listener is not None:
464 self._listener.close()
465 self._listener = None
Benjamin Petersone711caf2008-06-11 16:44:04 +0000466
467 address = property(lambda self: self._listener._address)
468 last_accepted = property(lambda self: self._listener._last_accepted)
469
Richard Oudkerkd69cfe82012-06-18 17:47:52 +0100470 def __enter__(self):
471 return self
472
473 def __exit__(self, exc_type, exc_value, exc_tb):
474 self.close()
475
Benjamin Petersone711caf2008-06-11 16:44:04 +0000476
477def Client(address, family=None, authkey=None):
478 '''
479 Returns a connection to the address of a `Listener`
480 '''
481 family = family or address_type(address)
Antoine Pitrou709176f2012-04-01 17:19:09 +0200482 _validate_family(family)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000483 if family == 'AF_PIPE':
484 c = PipeClient(address)
485 else:
486 c = SocketClient(address)
487
488 if authkey is not None and not isinstance(authkey, bytes):
489 raise TypeError('authkey should be a byte string')
490
491 if authkey is not None:
492 answer_challenge(c, authkey)
493 deliver_challenge(c, authkey)
494
495 return c
496
497
498if sys.platform != 'win32':
499
500 def Pipe(duplex=True):
501 '''
502 Returns pair of connection objects at either end of a pipe
503 '''
504 if duplex:
505 s1, s2 = socket.socketpair()
Richard Oudkerkb15e6222012-07-27 14:19:00 +0100506 s1.setblocking(True)
507 s2.setblocking(True)
Antoine Pitrou5aa878c2011-05-09 21:00:28 +0200508 c1 = Connection(s1.detach())
509 c2 = Connection(s2.detach())
Benjamin Petersone711caf2008-06-11 16:44:04 +0000510 else:
Victor Stinnerdaf45552013-08-28 00:53:59 +0200511 fd1, fd2 = os.pipe()
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200512 c1 = Connection(fd1, writable=False)
513 c2 = Connection(fd2, readable=False)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000514
515 return c1, c2
516
517else:
518
Benjamin Petersone711caf2008-06-11 16:44:04 +0000519 def Pipe(duplex=True):
520 '''
521 Returns pair of connection objects at either end of a pipe
522 '''
523 address = arbitrary_address('AF_PIPE')
524 if duplex:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200525 openmode = _winapi.PIPE_ACCESS_DUPLEX
526 access = _winapi.GENERIC_READ | _winapi.GENERIC_WRITE
Benjamin Petersone711caf2008-06-11 16:44:04 +0000527 obsize, ibsize = BUFSIZE, BUFSIZE
528 else:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200529 openmode = _winapi.PIPE_ACCESS_INBOUND
530 access = _winapi.GENERIC_WRITE
Benjamin Petersone711caf2008-06-11 16:44:04 +0000531 obsize, ibsize = 0, BUFSIZE
532
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200533 h1 = _winapi.CreateNamedPipe(
534 address, openmode | _winapi.FILE_FLAG_OVERLAPPED |
535 _winapi.FILE_FLAG_FIRST_PIPE_INSTANCE,
536 _winapi.PIPE_TYPE_MESSAGE | _winapi.PIPE_READMODE_MESSAGE |
537 _winapi.PIPE_WAIT,
Victor Stinnerdaf45552013-08-28 00:53:59 +0200538 1, obsize, ibsize, _winapi.NMPWAIT_WAIT_FOREVER,
539 # default security descriptor: the handle cannot be inherited
540 _winapi.NULL
Benjamin Petersone711caf2008-06-11 16:44:04 +0000541 )
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200542 h2 = _winapi.CreateFile(
543 address, access, 0, _winapi.NULL, _winapi.OPEN_EXISTING,
544 _winapi.FILE_FLAG_OVERLAPPED, _winapi.NULL
Benjamin Petersone711caf2008-06-11 16:44:04 +0000545 )
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200546 _winapi.SetNamedPipeHandleState(
547 h2, _winapi.PIPE_READMODE_MESSAGE, None, None
Benjamin Petersone711caf2008-06-11 16:44:04 +0000548 )
549
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200550 overlapped = _winapi.ConnectNamedPipe(h1, overlapped=True)
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100551 _, err = overlapped.GetOverlappedResult(True)
552 assert err == 0
Benjamin Petersone711caf2008-06-11 16:44:04 +0000553
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200554 c1 = PipeConnection(h1, writable=duplex)
555 c2 = PipeConnection(h2, readable=duplex)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000556
557 return c1, c2
558
559#
560# Definitions for connections based on sockets
561#
562
563class SocketListener(object):
564 '''
Georg Brandl734e2682008-08-12 08:18:18 +0000565 Representation of a socket which is bound to an address and listening
Benjamin Petersone711caf2008-06-11 16:44:04 +0000566 '''
567 def __init__(self, address, family, backlog=1):
568 self._socket = socket.socket(getattr(socket, family))
Charles-François Nataliedc67fe2012-02-04 15:12:08 +0100569 try:
Charles-François Natalied4a8fc2012-02-08 21:15:58 +0100570 # SO_REUSEADDR has different semantics on Windows (issue #2550).
571 if os.name == 'posix':
572 self._socket.setsockopt(socket.SOL_SOCKET,
573 socket.SO_REUSEADDR, 1)
Richard Oudkerkb15e6222012-07-27 14:19:00 +0100574 self._socket.setblocking(True)
Charles-François Nataliedc67fe2012-02-04 15:12:08 +0100575 self._socket.bind(address)
576 self._socket.listen(backlog)
577 self._address = self._socket.getsockname()
578 except OSError:
579 self._socket.close()
580 raise
Benjamin Petersone711caf2008-06-11 16:44:04 +0000581 self._family = family
582 self._last_accepted = None
583
Benjamin Petersone711caf2008-06-11 16:44:04 +0000584 if family == 'AF_UNIX':
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100585 self._unlink = util.Finalize(
Georg Brandl2ee470f2008-07-16 12:55:28 +0000586 self, os.unlink, args=(address,), exitpriority=0
Benjamin Petersone711caf2008-06-11 16:44:04 +0000587 )
588 else:
589 self._unlink = None
590
591 def accept(self):
Richard Oudkerkcca8c532013-07-01 18:59:26 +0100592 while True:
593 try:
594 s, self._last_accepted = self._socket.accept()
595 except InterruptedError:
596 pass
597 else:
598 break
Richard Oudkerkb15e6222012-07-27 14:19:00 +0100599 s.setblocking(True)
Antoine Pitroudf97cbe2012-04-07 22:38:52 +0200600 return Connection(s.detach())
Benjamin Petersone711caf2008-06-11 16:44:04 +0000601
602 def close(self):
603 self._socket.close()
604 if self._unlink is not None:
605 self._unlink()
606
607
608def SocketClient(address):
609 '''
610 Return a connection object connected to the socket given by `address`
611 '''
612 family = address_type(address)
Victor Stinner2b695062011-01-03 15:47:59 +0000613 with socket.socket( getattr(socket, family) ) as s:
Richard Oudkerkb15e6222012-07-27 14:19:00 +0100614 s.setblocking(True)
Charles-François Natalie6eabd42011-11-19 09:59:43 +0100615 s.connect(address)
Antoine Pitroudf97cbe2012-04-07 22:38:52 +0200616 return Connection(s.detach())
Benjamin Petersone711caf2008-06-11 16:44:04 +0000617
618#
619# Definitions for connections based on named pipes
620#
621
622if sys.platform == 'win32':
623
624 class PipeListener(object):
625 '''
626 Representation of a named pipe
627 '''
628 def __init__(self, address, backlog=None):
629 self._address = address
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100630 self._handle_queue = [self._new_handle(first=True)]
631
Benjamin Petersone711caf2008-06-11 16:44:04 +0000632 self._last_accepted = None
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100633 util.sub_debug('listener created with address=%r', self._address)
634 self.close = util.Finalize(
Benjamin Petersone711caf2008-06-11 16:44:04 +0000635 self, PipeListener._finalize_pipe_listener,
636 args=(self._handle_queue, self._address), exitpriority=0
637 )
638
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100639 def _new_handle(self, first=False):
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200640 flags = _winapi.PIPE_ACCESS_DUPLEX | _winapi.FILE_FLAG_OVERLAPPED
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100641 if first:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200642 flags |= _winapi.FILE_FLAG_FIRST_PIPE_INSTANCE
643 return _winapi.CreateNamedPipe(
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100644 self._address, flags,
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200645 _winapi.PIPE_TYPE_MESSAGE | _winapi.PIPE_READMODE_MESSAGE |
646 _winapi.PIPE_WAIT,
647 _winapi.PIPE_UNLIMITED_INSTANCES, BUFSIZE, BUFSIZE,
648 _winapi.NMPWAIT_WAIT_FOREVER, _winapi.NULL
Benjamin Petersone711caf2008-06-11 16:44:04 +0000649 )
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100650
651 def accept(self):
652 self._handle_queue.append(self._new_handle())
Benjamin Petersone711caf2008-06-11 16:44:04 +0000653 handle = self._handle_queue.pop(0)
654 try:
Richard Oudkerkfdb8dcf2012-05-05 19:45:37 +0100655 ov = _winapi.ConnectNamedPipe(handle, overlapped=True)
656 except OSError as e:
657 if e.winerror != _winapi.ERROR_NO_DATA:
658 raise
659 # ERROR_NO_DATA can occur if a client has already connected,
660 # written data and then disconnected -- see Issue 14725.
661 else:
662 try:
663 res = _winapi.WaitForMultipleObjects(
664 [ov.event], False, INFINITE)
665 except:
666 ov.cancel()
667 _winapi.CloseHandle(handle)
668 raise
669 finally:
670 _, err = ov.GetOverlappedResult(True)
671 assert err == 0
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200672 return PipeConnection(handle)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000673
674 @staticmethod
675 def _finalize_pipe_listener(queue, address):
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100676 util.sub_debug('closing listener with address=%r', address)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000677 for handle in queue:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200678 _winapi.CloseHandle(handle)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000679
680 def PipeClient(address):
681 '''
682 Return a connection object connected to the pipe given by `address`
683 '''
Antoine Pitrou45d61a32009-11-13 22:35:18 +0000684 t = _init_timeout()
Benjamin Petersone711caf2008-06-11 16:44:04 +0000685 while 1:
686 try:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200687 _winapi.WaitNamedPipe(address, 1000)
688 h = _winapi.CreateFile(
689 address, _winapi.GENERIC_READ | _winapi.GENERIC_WRITE,
690 0, _winapi.NULL, _winapi.OPEN_EXISTING,
691 _winapi.FILE_FLAG_OVERLAPPED, _winapi.NULL
Benjamin Petersone711caf2008-06-11 16:44:04 +0000692 )
Andrew Svetlov2606a6f2012-12-19 14:33:35 +0200693 except OSError as e:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200694 if e.winerror not in (_winapi.ERROR_SEM_TIMEOUT,
695 _winapi.ERROR_PIPE_BUSY) or _check_timeout(t):
Benjamin Petersone711caf2008-06-11 16:44:04 +0000696 raise
697 else:
698 break
699 else:
700 raise
701
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200702 _winapi.SetNamedPipeHandleState(
703 h, _winapi.PIPE_READMODE_MESSAGE, None, None
Benjamin Petersone711caf2008-06-11 16:44:04 +0000704 )
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200705 return PipeConnection(h)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000706
707#
708# Authentication stuff
709#
710
711MESSAGE_LENGTH = 20
712
Benjamin Peterson1fcfe212008-06-25 12:54:22 +0000713CHALLENGE = b'#CHALLENGE#'
714WELCOME = b'#WELCOME#'
715FAILURE = b'#FAILURE#'
Benjamin Petersone711caf2008-06-11 16:44:04 +0000716
717def deliver_challenge(connection, authkey):
718 import hmac
719 assert isinstance(authkey, bytes)
720 message = os.urandom(MESSAGE_LENGTH)
721 connection.send_bytes(CHALLENGE + message)
722 digest = hmac.new(authkey, message).digest()
723 response = connection.recv_bytes(256) # reject large message
724 if response == digest:
725 connection.send_bytes(WELCOME)
726 else:
727 connection.send_bytes(FAILURE)
728 raise AuthenticationError('digest received was wrong')
729
730def answer_challenge(connection, authkey):
731 import hmac
732 assert isinstance(authkey, bytes)
733 message = connection.recv_bytes(256) # reject large message
734 assert message[:len(CHALLENGE)] == CHALLENGE, 'message = %r' % message
735 message = message[len(CHALLENGE):]
736 digest = hmac.new(authkey, message).digest()
737 connection.send_bytes(digest)
738 response = connection.recv_bytes(256) # reject large message
739 if response != WELCOME:
740 raise AuthenticationError('digest sent was rejected')
741
742#
743# Support for using xmlrpclib for serialization
744#
745
746class ConnectionWrapper(object):
747 def __init__(self, conn, dumps, loads):
748 self._conn = conn
749 self._dumps = dumps
750 self._loads = loads
751 for attr in ('fileno', 'close', 'poll', 'recv_bytes', 'send_bytes'):
752 obj = getattr(conn, attr)
753 setattr(self, attr, obj)
754 def send(self, obj):
755 s = self._dumps(obj)
756 self._conn.send_bytes(s)
757 def recv(self):
758 s = self._conn.recv_bytes()
759 return self._loads(s)
760
761def _xml_dumps(obj):
Marc-André Lemburg8f36af72011-02-25 15:42:01 +0000762 return xmlrpclib.dumps((obj,), None, None, None, 1).encode('utf-8')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000763
764def _xml_loads(s):
Marc-André Lemburg8f36af72011-02-25 15:42:01 +0000765 (obj,), method = xmlrpclib.loads(s.decode('utf-8'))
Benjamin Petersone711caf2008-06-11 16:44:04 +0000766 return obj
767
768class XmlListener(Listener):
769 def accept(self):
770 global xmlrpclib
771 import xmlrpc.client as xmlrpclib
772 obj = Listener.accept(self)
773 return ConnectionWrapper(obj, _xml_dumps, _xml_loads)
774
775def XmlClient(*args, **kwds):
776 global xmlrpclib
777 import xmlrpc.client as xmlrpclib
778 return ConnectionWrapper(Client(*args, **kwds), _xml_dumps, _xml_loads)
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200779
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100780#
781# Wait
782#
783
784if sys.platform == 'win32':
785
786 def _exhaustive_wait(handles, timeout):
787 # Return ALL handles which are currently signalled. (Only
788 # returning the first signalled might create starvation issues.)
789 L = list(handles)
790 ready = []
791 while L:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200792 res = _winapi.WaitForMultipleObjects(L, False, timeout)
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100793 if res == WAIT_TIMEOUT:
794 break
795 elif WAIT_OBJECT_0 <= res < WAIT_OBJECT_0 + len(L):
796 res -= WAIT_OBJECT_0
797 elif WAIT_ABANDONED_0 <= res < WAIT_ABANDONED_0 + len(L):
798 res -= WAIT_ABANDONED_0
799 else:
800 raise RuntimeError('Should not get here')
801 ready.append(L[res])
802 L = L[res+1:]
803 timeout = 0
804 return ready
805
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200806 _ready_errors = {_winapi.ERROR_BROKEN_PIPE, _winapi.ERROR_NETNAME_DELETED}
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100807
808 def wait(object_list, timeout=None):
809 '''
810 Wait till an object in object_list is ready/readable.
811
812 Returns list of those objects in object_list which are ready/readable.
813 '''
814 if timeout is None:
815 timeout = INFINITE
816 elif timeout < 0:
817 timeout = 0
818 else:
819 timeout = int(timeout * 1000 + 0.5)
820
821 object_list = list(object_list)
822 waithandle_to_obj = {}
823 ov_list = []
824 ready_objects = set()
825 ready_handles = set()
826
827 try:
828 for o in object_list:
829 try:
830 fileno = getattr(o, 'fileno')
831 except AttributeError:
832 waithandle_to_obj[o.__index__()] = o
833 else:
834 # start an overlapped read of length zero
835 try:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200836 ov, err = _winapi.ReadFile(fileno(), 0, True)
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100837 except OSError as e:
838 err = e.winerror
839 if err not in _ready_errors:
840 raise
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200841 if err == _winapi.ERROR_IO_PENDING:
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100842 ov_list.append(ov)
843 waithandle_to_obj[ov.event] = o
844 else:
845 # If o.fileno() is an overlapped pipe handle and
846 # err == 0 then there is a zero length message
847 # in the pipe, but it HAS NOT been consumed.
848 ready_objects.add(o)
849 timeout = 0
850
851 ready_handles = _exhaustive_wait(waithandle_to_obj.keys(), timeout)
852 finally:
853 # request that overlapped reads stop
854 for ov in ov_list:
855 ov.cancel()
856
857 # wait for all overlapped reads to stop
858 for ov in ov_list:
859 try:
860 _, err = ov.GetOverlappedResult(True)
861 except OSError as e:
862 err = e.winerror
863 if err not in _ready_errors:
864 raise
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200865 if err != _winapi.ERROR_OPERATION_ABORTED:
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100866 o = waithandle_to_obj[ov.event]
867 ready_objects.add(o)
868 if err == 0:
869 # If o.fileno() is an overlapped pipe handle then
870 # a zero length message HAS been consumed.
871 if hasattr(o, '_got_empty_message'):
872 o._got_empty_message = True
873
874 ready_objects.update(waithandle_to_obj[h] for h in ready_handles)
875 return [o for o in object_list if o in ready_objects]
876
877else:
878
Charles-François Natalie241ac92013-09-05 20:46:49 +0200879 import selectors
Giampaolo Rodola'0c8ad612013-01-14 02:24:05 +0100880
Charles-François Natali45e25512013-09-08 11:30:53 +0200881 # poll/select have the advantage of not requiring any extra file
882 # descriptor, contrarily to epoll/kqueue (also, they require a single
883 # syscall).
884 if hasattr(selectors, 'PollSelector'):
885 _WaitSelector = selectors.PollSelector
886 else:
887 _WaitSelector = selectors.SelectSelector
888
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100889 def wait(object_list, timeout=None):
890 '''
891 Wait till an object in object_list is ready/readable.
892
893 Returns list of those objects in object_list which are ready/readable.
894 '''
Charles-François Natali45e25512013-09-08 11:30:53 +0200895 with _WaitSelector() as selector:
Charles-François Natalie241ac92013-09-05 20:46:49 +0200896 for obj in object_list:
897 selector.register(obj, selectors.EVENT_READ)
898
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100899 if timeout is not None:
Charles-François Natalie241ac92013-09-05 20:46:49 +0200900 deadline = time.time() + timeout
901
902 while True:
903 ready = selector.select(timeout)
904 if ready:
905 return [key.fileobj for (key, events) in ready]
906 else:
907 if timeout is not None:
908 timeout = deadline - time.time()
909 if timeout < 0:
910 return ready
Antoine Pitrou5438ed12012-04-24 22:56:57 +0200911
912#
913# Make connection and socket objects sharable if possible
914#
915
916if sys.platform == 'win32':
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100917 def reduce_connection(conn):
918 handle = conn.fileno()
919 with socket.fromfd(handle, socket.AF_INET, socket.SOCK_STREAM) as s:
920 from . import resource_sharer
921 ds = resource_sharer.DupSocket(s)
922 return rebuild_connection, (ds, conn.readable, conn.writable)
923 def rebuild_connection(ds, readable, writable):
924 sock = ds.detach()
925 return Connection(sock.detach(), readable, writable)
926 reduction.register(Connection, reduce_connection)
927
928 def reduce_pipe_connection(conn):
929 access = ((_winapi.FILE_GENERIC_READ if conn.readable else 0) |
930 (_winapi.FILE_GENERIC_WRITE if conn.writable else 0))
931 dh = reduction.DupHandle(conn.fileno(), access)
932 return rebuild_pipe_connection, (dh, conn.readable, conn.writable)
933 def rebuild_pipe_connection(dh, readable, writable):
934 handle = dh.detach()
935 return PipeConnection(handle, readable, writable)
936 reduction.register(PipeConnection, reduce_pipe_connection)
937
Antoine Pitrou5438ed12012-04-24 22:56:57 +0200938else:
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100939 def reduce_connection(conn):
940 df = reduction.DupFd(conn.fileno())
941 return rebuild_connection, (df, conn.readable, conn.writable)
942 def rebuild_connection(df, readable, writable):
943 fd = df.detach()
944 return Connection(fd, readable, writable)
945 reduction.register(Connection, reduce_connection)