blob: acf43b1feffe326a77287eeb2ca548b42bc0558f [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
Antoine Pitroudd696492011-06-08 17:21:55 +020026from multiprocessing.util import (
27 get_temp_dir, Finalize, sub_debug, debug, _eintr_retry)
Antoine Pitrou5438ed12012-04-24 22:56:57 +020028from multiprocessing.forking import ForkingPickler
Antoine Pitrou87cf2202011-05-09 17:04:27 +020029try:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +020030 import _winapi
31 from _winapi import WAIT_OBJECT_0, WAIT_TIMEOUT, INFINITE
Antoine Pitrou87cf2202011-05-09 17:04:27 +020032except ImportError:
33 if sys.platform == 'win32':
34 raise
Antoine Pitrou23bba4c2012-04-18 20:51:15 +020035 _winapi = None
Benjamin Petersone711caf2008-06-11 16:44:04 +000036
37#
38#
39#
40
41BUFSIZE = 8192
Antoine Pitrou45d61a32009-11-13 22:35:18 +000042# A very generous timeout when it comes to local connections...
43CONNECTION_TIMEOUT = 20.
Benjamin Petersone711caf2008-06-11 16:44:04 +000044
45_mmap_counter = itertools.count()
46
47default_family = 'AF_INET'
48families = ['AF_INET']
49
50if hasattr(socket, 'AF_UNIX'):
51 default_family = 'AF_UNIX'
52 families += ['AF_UNIX']
53
54if sys.platform == 'win32':
55 default_family = 'AF_PIPE'
56 families += ['AF_PIPE']
57
Antoine Pitrou45d61a32009-11-13 22:35:18 +000058
59def _init_timeout(timeout=CONNECTION_TIMEOUT):
60 return time.time() + timeout
61
62def _check_timeout(t):
63 return time.time() > t
64
Benjamin Petersone711caf2008-06-11 16:44:04 +000065#
66#
67#
68
69def arbitrary_address(family):
70 '''
71 Return an arbitrary free address for the given family
72 '''
73 if family == 'AF_INET':
74 return ('localhost', 0)
75 elif family == 'AF_UNIX':
76 return tempfile.mktemp(prefix='listener-', dir=get_temp_dir())
77 elif family == 'AF_PIPE':
78 return tempfile.mktemp(prefix=r'\\.\pipe\pyc-%d-%d-' %
79 (os.getpid(), next(_mmap_counter)))
80 else:
81 raise ValueError('unrecognized family')
82
Antoine Pitrou709176f2012-04-01 17:19:09 +020083def _validate_family(family):
84 '''
85 Checks if the family is valid for the current environment.
86 '''
87 if sys.platform != 'win32' and family == 'AF_PIPE':
88 raise ValueError('Family %s is not recognized.' % family)
89
Antoine Pitrou6d20cba2012-04-03 20:12:23 +020090 if sys.platform == 'win32' and family == 'AF_UNIX':
91 # double check
92 if not hasattr(socket, family):
93 raise ValueError('Family %s is not recognized.' % family)
Benjamin Petersone711caf2008-06-11 16:44:04 +000094
95def address_type(address):
96 '''
97 Return the types of the address
98
99 This can be 'AF_INET', 'AF_UNIX', or 'AF_PIPE'
100 '''
101 if type(address) == tuple:
102 return 'AF_INET'
103 elif type(address) is str and address.startswith('\\\\'):
104 return 'AF_PIPE'
105 elif type(address) is str:
106 return 'AF_UNIX'
107 else:
108 raise ValueError('address type of %r unrecognized' % address)
109
110#
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200111# Connection classes
112#
113
114class _ConnectionBase:
115 _handle = None
116
117 def __init__(self, handle, readable=True, writable=True):
118 handle = handle.__index__()
119 if handle < 0:
120 raise ValueError("invalid handle")
121 if not readable and not writable:
122 raise ValueError(
123 "at least one of `readable` and `writable` must be True")
124 self._handle = handle
125 self._readable = readable
126 self._writable = writable
127
Antoine Pitrou60001202011-07-09 01:03:46 +0200128 # XXX should we use util.Finalize instead of a __del__?
129
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200130 def __del__(self):
131 if self._handle is not None:
132 self._close()
133
134 def _check_closed(self):
135 if self._handle is None:
136 raise IOError("handle is closed")
137
138 def _check_readable(self):
139 if not self._readable:
140 raise IOError("connection is write-only")
141
142 def _check_writable(self):
143 if not self._writable:
144 raise IOError("connection is read-only")
145
146 def _bad_message_length(self):
147 if self._writable:
148 self._readable = False
149 else:
150 self.close()
151 raise IOError("bad message length")
152
153 @property
154 def closed(self):
155 """True if the connection is closed"""
156 return self._handle is None
157
158 @property
159 def readable(self):
160 """True if the connection is readable"""
161 return self._readable
162
163 @property
164 def writable(self):
165 """True if the connection is writable"""
166 return self._writable
167
168 def fileno(self):
169 """File descriptor or handle of the connection"""
170 self._check_closed()
171 return self._handle
172
173 def close(self):
174 """Close the connection"""
175 if self._handle is not None:
176 try:
177 self._close()
178 finally:
179 self._handle = None
180
181 def send_bytes(self, buf, offset=0, size=None):
182 """Send the bytes data from a bytes-like object"""
183 self._check_closed()
184 self._check_writable()
185 m = memoryview(buf)
186 # HACK for byte-indexing of non-bytewise buffers (e.g. array.array)
187 if m.itemsize > 1:
188 m = memoryview(bytes(m))
189 n = len(m)
190 if offset < 0:
191 raise ValueError("offset is negative")
192 if n < offset:
193 raise ValueError("buffer length < offset")
194 if size is None:
195 size = n - offset
196 elif size < 0:
197 raise ValueError("size is negative")
198 elif offset + size > n:
199 raise ValueError("buffer length < offset + size")
200 self._send_bytes(m[offset:offset + size])
201
202 def send(self, obj):
203 """Send a (picklable) object"""
204 self._check_closed()
205 self._check_writable()
Antoine Pitrou5438ed12012-04-24 22:56:57 +0200206 buf = io.BytesIO()
207 ForkingPickler(buf, pickle.HIGHEST_PROTOCOL).dump(obj)
208 self._send_bytes(buf.getbuffer())
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200209
210 def recv_bytes(self, maxlength=None):
211 """
212 Receive bytes data as a bytes object.
213 """
214 self._check_closed()
215 self._check_readable()
216 if maxlength is not None and maxlength < 0:
217 raise ValueError("negative maxlength")
218 buf = self._recv_bytes(maxlength)
219 if buf is None:
220 self._bad_message_length()
221 return buf.getvalue()
222
223 def recv_bytes_into(self, buf, offset=0):
224 """
225 Receive bytes data into a writeable buffer-like object.
226 Return the number of bytes read.
227 """
228 self._check_closed()
229 self._check_readable()
230 with memoryview(buf) as m:
231 # Get bytesize of arbitrary buffer
232 itemsize = m.itemsize
233 bytesize = itemsize * len(m)
234 if offset < 0:
235 raise ValueError("negative offset")
236 elif offset > bytesize:
237 raise ValueError("offset too large")
238 result = self._recv_bytes()
239 size = result.tell()
240 if bytesize < offset + size:
241 raise BufferTooShort(result.getvalue())
242 # Message can fit in dest
243 result.seek(0)
244 result.readinto(m[offset // itemsize :
245 (offset + size) // itemsize])
246 return size
247
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100248 def recv(self):
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200249 """Receive a (picklable) object"""
250 self._check_closed()
251 self._check_readable()
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100252 buf = self._recv_bytes()
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200253 return pickle.loads(buf.getbuffer())
254
255 def poll(self, timeout=0.0):
256 """Whether there is any input available to be read"""
257 self._check_closed()
258 self._check_readable()
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200259 return self._poll(timeout)
260
261
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200262if _winapi:
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200263
264 class PipeConnection(_ConnectionBase):
265 """
266 Connection class based on a Windows named pipe.
Antoine Pitroudd696492011-06-08 17:21:55 +0200267 Overlapped I/O is used, so the handles must have been created
268 with FILE_FLAG_OVERLAPPED.
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200269 """
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100270 _got_empty_message = False
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200271
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200272 def _close(self, _CloseHandle=_winapi.CloseHandle):
Antoine Pitrou71a28a92011-07-09 01:03:00 +0200273 _CloseHandle(self._handle)
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200274
275 def _send_bytes(self, buf):
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200276 ov, err = _winapi.WriteFile(self._handle, buf, overlapped=True)
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100277 try:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200278 if err == _winapi.ERROR_IO_PENDING:
279 waitres = _winapi.WaitForMultipleObjects(
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100280 [ov.event], False, INFINITE)
281 assert waitres == WAIT_OBJECT_0
282 except:
283 ov.cancel()
284 raise
285 finally:
286 nwritten, err = ov.GetOverlappedResult(True)
287 assert err == 0
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200288 assert nwritten == len(buf)
289
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100290 def _recv_bytes(self, maxsize=None):
291 if self._got_empty_message:
292 self._got_empty_message = False
293 return io.BytesIO()
Antoine Pitroudd696492011-06-08 17:21:55 +0200294 else:
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100295 bsize = 128 if maxsize is None else min(maxsize, 128)
Antoine Pitroudd696492011-06-08 17:21:55 +0200296 try:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200297 ov, err = _winapi.ReadFile(self._handle, bsize,
298 overlapped=True)
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100299 try:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200300 if err == _winapi.ERROR_IO_PENDING:
301 waitres = _winapi.WaitForMultipleObjects(
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100302 [ov.event], False, INFINITE)
303 assert waitres == WAIT_OBJECT_0
304 except:
305 ov.cancel()
306 raise
307 finally:
308 nread, err = ov.GetOverlappedResult(True)
309 if err == 0:
310 f = io.BytesIO()
311 f.write(ov.getbuffer())
312 return f
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200313 elif err == _winapi.ERROR_MORE_DATA:
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100314 return self._get_more_data(ov, maxsize)
Antoine Pitroudd696492011-06-08 17:21:55 +0200315 except IOError as e:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200316 if e.winerror == _winapi.ERROR_BROKEN_PIPE:
Antoine Pitroudd696492011-06-08 17:21:55 +0200317 raise EOFError
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100318 else:
319 raise
320 raise RuntimeError("shouldn't get here; expected KeyboardInterrupt")
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200321
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100322 def _poll(self, timeout):
323 if (self._got_empty_message or
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200324 _winapi.PeekNamedPipe(self._handle)[0] != 0):
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200325 return True
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100326 if timeout < 0:
327 timeout = None
328 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:
366 n = write(self._handle, buf)
367 remaining -= n
368 if remaining == 0:
369 break
370 buf = buf[n:]
371
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100372 def _recv(self, size, read=_read):
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200373 buf = io.BytesIO()
Antoine Pitroudd696492011-06-08 17:21:55 +0200374 handle = self._handle
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200375 remaining = size
376 while remaining > 0:
Antoine Pitroudd696492011-06-08 17:21:55 +0200377 chunk = read(handle, remaining)
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200378 n = len(chunk)
379 if n == 0:
380 if remaining == size:
381 raise EOFError
382 else:
383 raise IOError("got end of file during message")
384 buf.write(chunk)
385 remaining -= n
386 return buf
387
388 def _send_bytes(self, buf):
389 # For wire compatibility with 3.2 and lower
390 n = len(buf)
Charles-François Natali225aa4f2011-09-20 19:27:39 +0200391 self._send(struct.pack("!i", n))
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200392 # The condition is necessary to avoid "broken pipe" errors
393 # when sending a 0-length buffer if the other end closed the pipe.
394 if n > 0:
395 self._send(buf)
396
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100397 def _recv_bytes(self, maxsize=None):
398 buf = self._recv(4)
Charles-François Natali225aa4f2011-09-20 19:27:39 +0200399 size, = struct.unpack("!i", buf.getvalue())
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200400 if maxsize is not None and size > maxsize:
401 return None
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100402 return self._recv(size)
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200403
404 def _poll(self, timeout):
Antoine Pitroudd696492011-06-08 17:21:55 +0200405 if timeout < 0.0:
406 timeout = None
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100407 r = wait([self._handle], timeout)
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200408 return bool(r)
409
410
411#
Benjamin Petersone711caf2008-06-11 16:44:04 +0000412# Public functions
413#
414
415class Listener(object):
416 '''
417 Returns a listener object.
418
419 This is a wrapper for a bound socket which is 'listening' for
420 connections, or for a Windows named pipe.
421 '''
422 def __init__(self, address=None, family=None, backlog=1, authkey=None):
423 family = family or (address and address_type(address)) \
424 or default_family
425 address = address or arbitrary_address(family)
426
Antoine Pitrou709176f2012-04-01 17:19:09 +0200427 _validate_family(family)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000428 if family == 'AF_PIPE':
429 self._listener = PipeListener(address, backlog)
430 else:
431 self._listener = SocketListener(address, family, backlog)
432
433 if authkey is not None and not isinstance(authkey, bytes):
434 raise TypeError('authkey should be a byte string')
435
436 self._authkey = authkey
437
438 def accept(self):
439 '''
440 Accept a connection on the bound socket or named pipe of `self`.
441
442 Returns a `Connection` object.
443 '''
444 c = self._listener.accept()
445 if self._authkey:
446 deliver_challenge(c, self._authkey)
447 answer_challenge(c, self._authkey)
448 return c
449
450 def close(self):
451 '''
452 Close the bound socket or named pipe of `self`.
453 '''
454 return self._listener.close()
455
456 address = property(lambda self: self._listener._address)
457 last_accepted = property(lambda self: self._listener._last_accepted)
458
459
460def Client(address, family=None, authkey=None):
461 '''
462 Returns a connection to the address of a `Listener`
463 '''
464 family = family or address_type(address)
Antoine Pitrou709176f2012-04-01 17:19:09 +0200465 _validate_family(family)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000466 if family == 'AF_PIPE':
467 c = PipeClient(address)
468 else:
469 c = SocketClient(address)
470
471 if authkey is not None and not isinstance(authkey, bytes):
472 raise TypeError('authkey should be a byte string')
473
474 if authkey is not None:
475 answer_challenge(c, authkey)
476 deliver_challenge(c, authkey)
477
478 return c
479
480
481if sys.platform != 'win32':
482
483 def Pipe(duplex=True):
484 '''
485 Returns pair of connection objects at either end of a pipe
486 '''
487 if duplex:
488 s1, s2 = socket.socketpair()
Antoine Pitrou5aa878c2011-05-09 21:00:28 +0200489 c1 = Connection(s1.detach())
490 c2 = Connection(s2.detach())
Benjamin Petersone711caf2008-06-11 16:44:04 +0000491 else:
492 fd1, fd2 = os.pipe()
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200493 c1 = Connection(fd1, writable=False)
494 c2 = Connection(fd2, readable=False)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000495
496 return c1, c2
497
498else:
499
Benjamin Petersone711caf2008-06-11 16:44:04 +0000500 def Pipe(duplex=True):
501 '''
502 Returns pair of connection objects at either end of a pipe
503 '''
504 address = arbitrary_address('AF_PIPE')
505 if duplex:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200506 openmode = _winapi.PIPE_ACCESS_DUPLEX
507 access = _winapi.GENERIC_READ | _winapi.GENERIC_WRITE
Benjamin Petersone711caf2008-06-11 16:44:04 +0000508 obsize, ibsize = BUFSIZE, BUFSIZE
509 else:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200510 openmode = _winapi.PIPE_ACCESS_INBOUND
511 access = _winapi.GENERIC_WRITE
Benjamin Petersone711caf2008-06-11 16:44:04 +0000512 obsize, ibsize = 0, BUFSIZE
513
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200514 h1 = _winapi.CreateNamedPipe(
515 address, openmode | _winapi.FILE_FLAG_OVERLAPPED |
516 _winapi.FILE_FLAG_FIRST_PIPE_INSTANCE,
517 _winapi.PIPE_TYPE_MESSAGE | _winapi.PIPE_READMODE_MESSAGE |
518 _winapi.PIPE_WAIT,
519 1, obsize, ibsize, _winapi.NMPWAIT_WAIT_FOREVER, _winapi.NULL
Benjamin Petersone711caf2008-06-11 16:44:04 +0000520 )
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200521 h2 = _winapi.CreateFile(
522 address, access, 0, _winapi.NULL, _winapi.OPEN_EXISTING,
523 _winapi.FILE_FLAG_OVERLAPPED, _winapi.NULL
Benjamin Petersone711caf2008-06-11 16:44:04 +0000524 )
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200525 _winapi.SetNamedPipeHandleState(
526 h2, _winapi.PIPE_READMODE_MESSAGE, None, None
Benjamin Petersone711caf2008-06-11 16:44:04 +0000527 )
528
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200529 overlapped = _winapi.ConnectNamedPipe(h1, overlapped=True)
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100530 _, err = overlapped.GetOverlappedResult(True)
531 assert err == 0
Benjamin Petersone711caf2008-06-11 16:44:04 +0000532
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200533 c1 = PipeConnection(h1, writable=duplex)
534 c2 = PipeConnection(h2, readable=duplex)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000535
536 return c1, c2
537
538#
539# Definitions for connections based on sockets
540#
541
542class SocketListener(object):
543 '''
Georg Brandl734e2682008-08-12 08:18:18 +0000544 Representation of a socket which is bound to an address and listening
Benjamin Petersone711caf2008-06-11 16:44:04 +0000545 '''
546 def __init__(self, address, family, backlog=1):
547 self._socket = socket.socket(getattr(socket, family))
Charles-François Nataliedc67fe2012-02-04 15:12:08 +0100548 try:
Charles-François Natalied4a8fc2012-02-08 21:15:58 +0100549 # SO_REUSEADDR has different semantics on Windows (issue #2550).
550 if os.name == 'posix':
551 self._socket.setsockopt(socket.SOL_SOCKET,
552 socket.SO_REUSEADDR, 1)
Charles-François Nataliedc67fe2012-02-04 15:12:08 +0100553 self._socket.bind(address)
554 self._socket.listen(backlog)
555 self._address = self._socket.getsockname()
556 except OSError:
557 self._socket.close()
558 raise
Benjamin Petersone711caf2008-06-11 16:44:04 +0000559 self._family = family
560 self._last_accepted = None
561
Benjamin Petersone711caf2008-06-11 16:44:04 +0000562 if family == 'AF_UNIX':
563 self._unlink = Finalize(
Georg Brandl2ee470f2008-07-16 12:55:28 +0000564 self, os.unlink, args=(address,), exitpriority=0
Benjamin Petersone711caf2008-06-11 16:44:04 +0000565 )
566 else:
567 self._unlink = None
568
569 def accept(self):
570 s, self._last_accepted = self._socket.accept()
Antoine Pitroudf97cbe2012-04-07 22:38:52 +0200571 return Connection(s.detach())
Benjamin Petersone711caf2008-06-11 16:44:04 +0000572
573 def close(self):
574 self._socket.close()
575 if self._unlink is not None:
576 self._unlink()
577
578
579def SocketClient(address):
580 '''
581 Return a connection object connected to the socket given by `address`
582 '''
583 family = address_type(address)
Victor Stinner2b695062011-01-03 15:47:59 +0000584 with socket.socket( getattr(socket, family) ) as s:
Charles-François Natalie6eabd42011-11-19 09:59:43 +0100585 s.connect(address)
Antoine Pitroudf97cbe2012-04-07 22:38:52 +0200586 return Connection(s.detach())
Benjamin Petersone711caf2008-06-11 16:44:04 +0000587
588#
589# Definitions for connections based on named pipes
590#
591
592if sys.platform == 'win32':
593
594 class PipeListener(object):
595 '''
596 Representation of a named pipe
597 '''
598 def __init__(self, address, backlog=None):
599 self._address = address
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100600 self._handle_queue = [self._new_handle(first=True)]
601
Benjamin Petersone711caf2008-06-11 16:44:04 +0000602 self._last_accepted = None
Benjamin Petersone711caf2008-06-11 16:44:04 +0000603 sub_debug('listener created with address=%r', self._address)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000604 self.close = Finalize(
605 self, PipeListener._finalize_pipe_listener,
606 args=(self._handle_queue, self._address), exitpriority=0
607 )
608
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100609 def _new_handle(self, first=False):
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200610 flags = _winapi.PIPE_ACCESS_DUPLEX | _winapi.FILE_FLAG_OVERLAPPED
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100611 if first:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200612 flags |= _winapi.FILE_FLAG_FIRST_PIPE_INSTANCE
613 return _winapi.CreateNamedPipe(
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100614 self._address, flags,
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200615 _winapi.PIPE_TYPE_MESSAGE | _winapi.PIPE_READMODE_MESSAGE |
616 _winapi.PIPE_WAIT,
617 _winapi.PIPE_UNLIMITED_INSTANCES, BUFSIZE, BUFSIZE,
618 _winapi.NMPWAIT_WAIT_FOREVER, _winapi.NULL
Benjamin Petersone711caf2008-06-11 16:44:04 +0000619 )
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100620
621 def accept(self):
622 self._handle_queue.append(self._new_handle())
Benjamin Petersone711caf2008-06-11 16:44:04 +0000623 handle = self._handle_queue.pop(0)
624 try:
Richard Oudkerkfdb8dcf2012-05-05 19:45:37 +0100625 ov = _winapi.ConnectNamedPipe(handle, overlapped=True)
626 except OSError as e:
627 if e.winerror != _winapi.ERROR_NO_DATA:
628 raise
629 # ERROR_NO_DATA can occur if a client has already connected,
630 # written data and then disconnected -- see Issue 14725.
631 else:
632 try:
633 res = _winapi.WaitForMultipleObjects(
634 [ov.event], False, INFINITE)
635 except:
636 ov.cancel()
637 _winapi.CloseHandle(handle)
638 raise
639 finally:
640 _, err = ov.GetOverlappedResult(True)
641 assert err == 0
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200642 return PipeConnection(handle)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000643
644 @staticmethod
645 def _finalize_pipe_listener(queue, address):
646 sub_debug('closing listener with address=%r', address)
647 for handle in queue:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200648 _winapi.CloseHandle(handle)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000649
650 def PipeClient(address):
651 '''
652 Return a connection object connected to the pipe given by `address`
653 '''
Antoine Pitrou45d61a32009-11-13 22:35:18 +0000654 t = _init_timeout()
Benjamin Petersone711caf2008-06-11 16:44:04 +0000655 while 1:
656 try:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200657 _winapi.WaitNamedPipe(address, 1000)
658 h = _winapi.CreateFile(
659 address, _winapi.GENERIC_READ | _winapi.GENERIC_WRITE,
660 0, _winapi.NULL, _winapi.OPEN_EXISTING,
661 _winapi.FILE_FLAG_OVERLAPPED, _winapi.NULL
Benjamin Petersone711caf2008-06-11 16:44:04 +0000662 )
663 except WindowsError as e:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200664 if e.winerror not in (_winapi.ERROR_SEM_TIMEOUT,
665 _winapi.ERROR_PIPE_BUSY) or _check_timeout(t):
Benjamin Petersone711caf2008-06-11 16:44:04 +0000666 raise
667 else:
668 break
669 else:
670 raise
671
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200672 _winapi.SetNamedPipeHandleState(
673 h, _winapi.PIPE_READMODE_MESSAGE, None, None
Benjamin Petersone711caf2008-06-11 16:44:04 +0000674 )
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200675 return PipeConnection(h)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000676
677#
678# Authentication stuff
679#
680
681MESSAGE_LENGTH = 20
682
Benjamin Peterson1fcfe212008-06-25 12:54:22 +0000683CHALLENGE = b'#CHALLENGE#'
684WELCOME = b'#WELCOME#'
685FAILURE = b'#FAILURE#'
Benjamin Petersone711caf2008-06-11 16:44:04 +0000686
687def deliver_challenge(connection, authkey):
688 import hmac
689 assert isinstance(authkey, bytes)
690 message = os.urandom(MESSAGE_LENGTH)
691 connection.send_bytes(CHALLENGE + message)
692 digest = hmac.new(authkey, message).digest()
693 response = connection.recv_bytes(256) # reject large message
694 if response == digest:
695 connection.send_bytes(WELCOME)
696 else:
697 connection.send_bytes(FAILURE)
698 raise AuthenticationError('digest received was wrong')
699
700def answer_challenge(connection, authkey):
701 import hmac
702 assert isinstance(authkey, bytes)
703 message = connection.recv_bytes(256) # reject large message
704 assert message[:len(CHALLENGE)] == CHALLENGE, 'message = %r' % message
705 message = message[len(CHALLENGE):]
706 digest = hmac.new(authkey, message).digest()
707 connection.send_bytes(digest)
708 response = connection.recv_bytes(256) # reject large message
709 if response != WELCOME:
710 raise AuthenticationError('digest sent was rejected')
711
712#
713# Support for using xmlrpclib for serialization
714#
715
716class ConnectionWrapper(object):
717 def __init__(self, conn, dumps, loads):
718 self._conn = conn
719 self._dumps = dumps
720 self._loads = loads
721 for attr in ('fileno', 'close', 'poll', 'recv_bytes', 'send_bytes'):
722 obj = getattr(conn, attr)
723 setattr(self, attr, obj)
724 def send(self, obj):
725 s = self._dumps(obj)
726 self._conn.send_bytes(s)
727 def recv(self):
728 s = self._conn.recv_bytes()
729 return self._loads(s)
730
731def _xml_dumps(obj):
Marc-André Lemburg8f36af72011-02-25 15:42:01 +0000732 return xmlrpclib.dumps((obj,), None, None, None, 1).encode('utf-8')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000733
734def _xml_loads(s):
Marc-André Lemburg8f36af72011-02-25 15:42:01 +0000735 (obj,), method = xmlrpclib.loads(s.decode('utf-8'))
Benjamin Petersone711caf2008-06-11 16:44:04 +0000736 return obj
737
738class XmlListener(Listener):
739 def accept(self):
740 global xmlrpclib
741 import xmlrpc.client as xmlrpclib
742 obj = Listener.accept(self)
743 return ConnectionWrapper(obj, _xml_dumps, _xml_loads)
744
745def XmlClient(*args, **kwds):
746 global xmlrpclib
747 import xmlrpc.client as xmlrpclib
748 return ConnectionWrapper(Client(*args, **kwds), _xml_dumps, _xml_loads)
Antoine Pitrou87cf2202011-05-09 17:04:27 +0200749
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100750#
751# Wait
752#
753
754if sys.platform == 'win32':
755
756 def _exhaustive_wait(handles, timeout):
757 # Return ALL handles which are currently signalled. (Only
758 # returning the first signalled might create starvation issues.)
759 L = list(handles)
760 ready = []
761 while L:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200762 res = _winapi.WaitForMultipleObjects(L, False, timeout)
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100763 if res == WAIT_TIMEOUT:
764 break
765 elif WAIT_OBJECT_0 <= res < WAIT_OBJECT_0 + len(L):
766 res -= WAIT_OBJECT_0
767 elif WAIT_ABANDONED_0 <= res < WAIT_ABANDONED_0 + len(L):
768 res -= WAIT_ABANDONED_0
769 else:
770 raise RuntimeError('Should not get here')
771 ready.append(L[res])
772 L = L[res+1:]
773 timeout = 0
774 return ready
775
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200776 _ready_errors = {_winapi.ERROR_BROKEN_PIPE, _winapi.ERROR_NETNAME_DELETED}
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100777
778 def wait(object_list, timeout=None):
779 '''
780 Wait till an object in object_list is ready/readable.
781
782 Returns list of those objects in object_list which are ready/readable.
783 '''
784 if timeout is None:
785 timeout = INFINITE
786 elif timeout < 0:
787 timeout = 0
788 else:
789 timeout = int(timeout * 1000 + 0.5)
790
791 object_list = list(object_list)
792 waithandle_to_obj = {}
793 ov_list = []
794 ready_objects = set()
795 ready_handles = set()
796
797 try:
798 for o in object_list:
799 try:
800 fileno = getattr(o, 'fileno')
801 except AttributeError:
802 waithandle_to_obj[o.__index__()] = o
803 else:
804 # start an overlapped read of length zero
805 try:
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200806 ov, err = _winapi.ReadFile(fileno(), 0, True)
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100807 except OSError as e:
808 err = e.winerror
809 if err not in _ready_errors:
810 raise
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200811 if err == _winapi.ERROR_IO_PENDING:
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100812 ov_list.append(ov)
813 waithandle_to_obj[ov.event] = o
814 else:
815 # If o.fileno() is an overlapped pipe handle and
816 # err == 0 then there is a zero length message
817 # in the pipe, but it HAS NOT been consumed.
818 ready_objects.add(o)
819 timeout = 0
820
821 ready_handles = _exhaustive_wait(waithandle_to_obj.keys(), timeout)
822 finally:
823 # request that overlapped reads stop
824 for ov in ov_list:
825 ov.cancel()
826
827 # wait for all overlapped reads to stop
828 for ov in ov_list:
829 try:
830 _, err = ov.GetOverlappedResult(True)
831 except OSError as e:
832 err = e.winerror
833 if err not in _ready_errors:
834 raise
Antoine Pitrou23bba4c2012-04-18 20:51:15 +0200835 if err != _winapi.ERROR_OPERATION_ABORTED:
Antoine Pitroubdb1cf12012-03-05 19:28:37 +0100836 o = waithandle_to_obj[ov.event]
837 ready_objects.add(o)
838 if err == 0:
839 # If o.fileno() is an overlapped pipe handle then
840 # a zero length message HAS been consumed.
841 if hasattr(o, '_got_empty_message'):
842 o._got_empty_message = True
843
844 ready_objects.update(waithandle_to_obj[h] for h in ready_handles)
845 return [o for o in object_list if o in ready_objects]
846
847else:
848
849 def wait(object_list, timeout=None):
850 '''
851 Wait till an object in object_list is ready/readable.
852
853 Returns list of those objects in object_list which are ready/readable.
854 '''
855 if timeout is not None:
856 if timeout <= 0:
857 return select.select(object_list, [], [], 0)[0]
858 else:
859 deadline = time.time() + timeout
860 while True:
861 try:
862 return select.select(object_list, [], [], timeout)[0]
863 except OSError as e:
864 if e.errno != errno.EINTR:
865 raise
866 if timeout is not None:
867 timeout = deadline - time.time()
Antoine Pitrou5438ed12012-04-24 22:56:57 +0200868
869#
870# Make connection and socket objects sharable if possible
871#
872
873if sys.platform == 'win32':
874 from . import reduction
875 ForkingPickler.register(socket.socket, reduction.reduce_socket)
876 ForkingPickler.register(Connection, reduction.reduce_connection)
877 ForkingPickler.register(PipeConnection, reduction.reduce_pipe_connection)
878else:
879 try:
880 from . import reduction
881 except ImportError:
882 pass
883 else:
884 ForkingPickler.register(socket.socket, reduction.reduce_socket)
885 ForkingPickler.register(Connection, reduction.reduce_connection)