jpayne@68: #
jpayne@68: # A higher level module for using sockets (or Windows named pipes)
jpayne@68: #
jpayne@68: # multiprocessing/connection.py
jpayne@68: #
jpayne@68: # Copyright (c) 2006-2008, R Oudkerk
jpayne@68: # Licensed to PSF under a Contributor Agreement.
jpayne@68: #
jpayne@68: 
jpayne@68: __all__ = [ 'Client', 'Listener', 'Pipe', 'wait' ]
jpayne@68: 
jpayne@68: import io
jpayne@68: import os
jpayne@68: import sys
jpayne@68: import socket
jpayne@68: import struct
jpayne@68: import time
jpayne@68: import tempfile
jpayne@68: import itertools
jpayne@68: 
jpayne@68: import _multiprocessing
jpayne@68: 
jpayne@68: from . import util
jpayne@68: 
jpayne@68: from . import AuthenticationError, BufferTooShort
jpayne@68: from .context import reduction
jpayne@68: _ForkingPickler = reduction.ForkingPickler
jpayne@68: 
jpayne@68: try:
jpayne@68:     import _winapi
jpayne@68:     from _winapi import WAIT_OBJECT_0, WAIT_ABANDONED_0, WAIT_TIMEOUT, INFINITE
jpayne@68: except ImportError:
jpayne@68:     if sys.platform == 'win32':
jpayne@68:         raise
jpayne@68:     _winapi = None
jpayne@68: 
jpayne@68: #
jpayne@68: #
jpayne@68: #
jpayne@68: 
jpayne@68: BUFSIZE = 8192
jpayne@68: # A very generous timeout when it comes to local connections...
jpayne@68: CONNECTION_TIMEOUT = 20.
jpayne@68: 
jpayne@68: _mmap_counter = itertools.count()
jpayne@68: 
jpayne@68: default_family = 'AF_INET'
jpayne@68: families = ['AF_INET']
jpayne@68: 
jpayne@68: if hasattr(socket, 'AF_UNIX'):
jpayne@68:     default_family = 'AF_UNIX'
jpayne@68:     families += ['AF_UNIX']
jpayne@68: 
jpayne@68: if sys.platform == 'win32':
jpayne@68:     default_family = 'AF_PIPE'
jpayne@68:     families += ['AF_PIPE']
jpayne@68: 
jpayne@68: 
jpayne@68: def _init_timeout(timeout=CONNECTION_TIMEOUT):
jpayne@68:     return time.monotonic() + timeout
jpayne@68: 
jpayne@68: def _check_timeout(t):
jpayne@68:     return time.monotonic() > t
jpayne@68: 
jpayne@68: #
jpayne@68: #
jpayne@68: #
jpayne@68: 
jpayne@68: def arbitrary_address(family):
jpayne@68:     '''
jpayne@68:     Return an arbitrary free address for the given family
jpayne@68:     '''
jpayne@68:     if family == 'AF_INET':
jpayne@68:         return ('localhost', 0)
jpayne@68:     elif family == 'AF_UNIX':
jpayne@68:         return tempfile.mktemp(prefix='listener-', dir=util.get_temp_dir())
jpayne@68:     elif family == 'AF_PIPE':
jpayne@68:         return tempfile.mktemp(prefix=r'\\.\pipe\pyc-%d-%d-' %
jpayne@68:                                (os.getpid(), next(_mmap_counter)), dir="")
jpayne@68:     else:
jpayne@68:         raise ValueError('unrecognized family')
jpayne@68: 
jpayne@68: def _validate_family(family):
jpayne@68:     '''
jpayne@68:     Checks if the family is valid for the current environment.
jpayne@68:     '''
jpayne@68:     if sys.platform != 'win32' and family == 'AF_PIPE':
jpayne@68:         raise ValueError('Family %s is not recognized.' % family)
jpayne@68: 
jpayne@68:     if sys.platform == 'win32' and family == 'AF_UNIX':
jpayne@68:         # double check
jpayne@68:         if not hasattr(socket, family):
jpayne@68:             raise ValueError('Family %s is not recognized.' % family)
jpayne@68: 
jpayne@68: def address_type(address):
jpayne@68:     '''
jpayne@68:     Return the types of the address
jpayne@68: 
jpayne@68:     This can be 'AF_INET', 'AF_UNIX', or 'AF_PIPE'
jpayne@68:     '''
jpayne@68:     if type(address) == tuple:
jpayne@68:         return 'AF_INET'
jpayne@68:     elif type(address) is str and address.startswith('\\\\'):
jpayne@68:         return 'AF_PIPE'
jpayne@68:     elif type(address) is str:
jpayne@68:         return 'AF_UNIX'
jpayne@68:     else:
jpayne@68:         raise ValueError('address type of %r unrecognized' % address)
jpayne@68: 
jpayne@68: #
jpayne@68: # Connection classes
jpayne@68: #
jpayne@68: 
jpayne@68: class _ConnectionBase:
jpayne@68:     _handle = None
jpayne@68: 
jpayne@68:     def __init__(self, handle, readable=True, writable=True):
jpayne@68:         handle = handle.__index__()
jpayne@68:         if handle < 0:
jpayne@68:             raise ValueError("invalid handle")
jpayne@68:         if not readable and not writable:
jpayne@68:             raise ValueError(
jpayne@68:                 "at least one of `readable` and `writable` must be True")
jpayne@68:         self._handle = handle
jpayne@68:         self._readable = readable
jpayne@68:         self._writable = writable
jpayne@68: 
jpayne@68:     # XXX should we use util.Finalize instead of a __del__?
jpayne@68: 
jpayne@68:     def __del__(self):
jpayne@68:         if self._handle is not None:
jpayne@68:             self._close()
jpayne@68: 
jpayne@68:     def _check_closed(self):
jpayne@68:         if self._handle is None:
jpayne@68:             raise OSError("handle is closed")
jpayne@68: 
jpayne@68:     def _check_readable(self):
jpayne@68:         if not self._readable:
jpayne@68:             raise OSError("connection is write-only")
jpayne@68: 
jpayne@68:     def _check_writable(self):
jpayne@68:         if not self._writable:
jpayne@68:             raise OSError("connection is read-only")
jpayne@68: 
jpayne@68:     def _bad_message_length(self):
jpayne@68:         if self._writable:
jpayne@68:             self._readable = False
jpayne@68:         else:
jpayne@68:             self.close()
jpayne@68:         raise OSError("bad message length")
jpayne@68: 
jpayne@68:     @property
jpayne@68:     def closed(self):
jpayne@68:         """True if the connection is closed"""
jpayne@68:         return self._handle is None
jpayne@68: 
jpayne@68:     @property
jpayne@68:     def readable(self):
jpayne@68:         """True if the connection is readable"""
jpayne@68:         return self._readable
jpayne@68: 
jpayne@68:     @property
jpayne@68:     def writable(self):
jpayne@68:         """True if the connection is writable"""
jpayne@68:         return self._writable
jpayne@68: 
jpayne@68:     def fileno(self):
jpayne@68:         """File descriptor or handle of the connection"""
jpayne@68:         self._check_closed()
jpayne@68:         return self._handle
jpayne@68: 
jpayne@68:     def close(self):
jpayne@68:         """Close the connection"""
jpayne@68:         if self._handle is not None:
jpayne@68:             try:
jpayne@68:                 self._close()
jpayne@68:             finally:
jpayne@68:                 self._handle = None
jpayne@68: 
jpayne@68:     def send_bytes(self, buf, offset=0, size=None):
jpayne@68:         """Send the bytes data from a bytes-like object"""
jpayne@68:         self._check_closed()
jpayne@68:         self._check_writable()
jpayne@68:         m = memoryview(buf)
jpayne@68:         # HACK for byte-indexing of non-bytewise buffers (e.g. array.array)
jpayne@68:         if m.itemsize > 1:
jpayne@68:             m = memoryview(bytes(m))
jpayne@68:         n = len(m)
jpayne@68:         if offset < 0:
jpayne@68:             raise ValueError("offset is negative")
jpayne@68:         if n < offset:
jpayne@68:             raise ValueError("buffer length < offset")
jpayne@68:         if size is None:
jpayne@68:             size = n - offset
jpayne@68:         elif size < 0:
jpayne@68:             raise ValueError("size is negative")
jpayne@68:         elif offset + size > n:
jpayne@68:             raise ValueError("buffer length < offset + size")
jpayne@68:         self._send_bytes(m[offset:offset + size])
jpayne@68: 
jpayne@68:     def send(self, obj):
jpayne@68:         """Send a (picklable) object"""
jpayne@68:         self._check_closed()
jpayne@68:         self._check_writable()
jpayne@68:         self._send_bytes(_ForkingPickler.dumps(obj))
jpayne@68: 
jpayne@68:     def recv_bytes(self, maxlength=None):
jpayne@68:         """
jpayne@68:         Receive bytes data as a bytes object.
jpayne@68:         """
jpayne@68:         self._check_closed()
jpayne@68:         self._check_readable()
jpayne@68:         if maxlength is not None and maxlength < 0:
jpayne@68:             raise ValueError("negative maxlength")
jpayne@68:         buf = self._recv_bytes(maxlength)
jpayne@68:         if buf is None:
jpayne@68:             self._bad_message_length()
jpayne@68:         return buf.getvalue()
jpayne@68: 
jpayne@68:     def recv_bytes_into(self, buf, offset=0):
jpayne@68:         """
jpayne@68:         Receive bytes data into a writeable bytes-like object.
jpayne@68:         Return the number of bytes read.
jpayne@68:         """
jpayne@68:         self._check_closed()
jpayne@68:         self._check_readable()
jpayne@68:         with memoryview(buf) as m:
jpayne@68:             # Get bytesize of arbitrary buffer
jpayne@68:             itemsize = m.itemsize
jpayne@68:             bytesize = itemsize * len(m)
jpayne@68:             if offset < 0:
jpayne@68:                 raise ValueError("negative offset")
jpayne@68:             elif offset > bytesize:
jpayne@68:                 raise ValueError("offset too large")
jpayne@68:             result = self._recv_bytes()
jpayne@68:             size = result.tell()
jpayne@68:             if bytesize < offset + size:
jpayne@68:                 raise BufferTooShort(result.getvalue())
jpayne@68:             # Message can fit in dest
jpayne@68:             result.seek(0)
jpayne@68:             result.readinto(m[offset // itemsize :
jpayne@68:                               (offset + size) // itemsize])
jpayne@68:             return size
jpayne@68: 
jpayne@68:     def recv(self):
jpayne@68:         """Receive a (picklable) object"""
jpayne@68:         self._check_closed()
jpayne@68:         self._check_readable()
jpayne@68:         buf = self._recv_bytes()
jpayne@68:         return _ForkingPickler.loads(buf.getbuffer())
jpayne@68: 
jpayne@68:     def poll(self, timeout=0.0):
jpayne@68:         """Whether there is any input available to be read"""
jpayne@68:         self._check_closed()
jpayne@68:         self._check_readable()
jpayne@68:         return self._poll(timeout)
jpayne@68: 
jpayne@68:     def __enter__(self):
jpayne@68:         return self
jpayne@68: 
jpayne@68:     def __exit__(self, exc_type, exc_value, exc_tb):
jpayne@68:         self.close()
jpayne@68: 
jpayne@68: 
jpayne@68: if _winapi:
jpayne@68: 
jpayne@68:     class PipeConnection(_ConnectionBase):
jpayne@68:         """
jpayne@68:         Connection class based on a Windows named pipe.
jpayne@68:         Overlapped I/O is used, so the handles must have been created
jpayne@68:         with FILE_FLAG_OVERLAPPED.
jpayne@68:         """
jpayne@68:         _got_empty_message = False
jpayne@68: 
jpayne@68:         def _close(self, _CloseHandle=_winapi.CloseHandle):
jpayne@68:             _CloseHandle(self._handle)
jpayne@68: 
jpayne@68:         def _send_bytes(self, buf):
jpayne@68:             ov, err = _winapi.WriteFile(self._handle, buf, overlapped=True)
jpayne@68:             try:
jpayne@68:                 if err == _winapi.ERROR_IO_PENDING:
jpayne@68:                     waitres = _winapi.WaitForMultipleObjects(
jpayne@68:                         [ov.event], False, INFINITE)
jpayne@68:                     assert waitres == WAIT_OBJECT_0
jpayne@68:             except:
jpayne@68:                 ov.cancel()
jpayne@68:                 raise
jpayne@68:             finally:
jpayne@68:                 nwritten, err = ov.GetOverlappedResult(True)
jpayne@68:             assert err == 0
jpayne@68:             assert nwritten == len(buf)
jpayne@68: 
jpayne@68:         def _recv_bytes(self, maxsize=None):
jpayne@68:             if self._got_empty_message:
jpayne@68:                 self._got_empty_message = False
jpayne@68:                 return io.BytesIO()
jpayne@68:             else:
jpayne@68:                 bsize = 128 if maxsize is None else min(maxsize, 128)
jpayne@68:                 try:
jpayne@68:                     ov, err = _winapi.ReadFile(self._handle, bsize,
jpayne@68:                                                 overlapped=True)
jpayne@68:                     try:
jpayne@68:                         if err == _winapi.ERROR_IO_PENDING:
jpayne@68:                             waitres = _winapi.WaitForMultipleObjects(
jpayne@68:                                 [ov.event], False, INFINITE)
jpayne@68:                             assert waitres == WAIT_OBJECT_0
jpayne@68:                     except:
jpayne@68:                         ov.cancel()
jpayne@68:                         raise
jpayne@68:                     finally:
jpayne@68:                         nread, err = ov.GetOverlappedResult(True)
jpayne@68:                         if err == 0:
jpayne@68:                             f = io.BytesIO()
jpayne@68:                             f.write(ov.getbuffer())
jpayne@68:                             return f
jpayne@68:                         elif err == _winapi.ERROR_MORE_DATA:
jpayne@68:                             return self._get_more_data(ov, maxsize)
jpayne@68:                 except OSError as e:
jpayne@68:                     if e.winerror == _winapi.ERROR_BROKEN_PIPE:
jpayne@68:                         raise EOFError
jpayne@68:                     else:
jpayne@68:                         raise
jpayne@68:             raise RuntimeError("shouldn't get here; expected KeyboardInterrupt")
jpayne@68: 
jpayne@68:         def _poll(self, timeout):
jpayne@68:             if (self._got_empty_message or
jpayne@68:                         _winapi.PeekNamedPipe(self._handle)[0] != 0):
jpayne@68:                 return True
jpayne@68:             return bool(wait([self], timeout))
jpayne@68: 
jpayne@68:         def _get_more_data(self, ov, maxsize):
jpayne@68:             buf = ov.getbuffer()
jpayne@68:             f = io.BytesIO()
jpayne@68:             f.write(buf)
jpayne@68:             left = _winapi.PeekNamedPipe(self._handle)[1]
jpayne@68:             assert left > 0
jpayne@68:             if maxsize is not None and len(buf) + left > maxsize:
jpayne@68:                 self._bad_message_length()
jpayne@68:             ov, err = _winapi.ReadFile(self._handle, left, overlapped=True)
jpayne@68:             rbytes, err = ov.GetOverlappedResult(True)
jpayne@68:             assert err == 0
jpayne@68:             assert rbytes == left
jpayne@68:             f.write(ov.getbuffer())
jpayne@68:             return f
jpayne@68: 
jpayne@68: 
jpayne@68: class Connection(_ConnectionBase):
jpayne@68:     """
jpayne@68:     Connection class based on an arbitrary file descriptor (Unix only), or
jpayne@68:     a socket handle (Windows).
jpayne@68:     """
jpayne@68: 
jpayne@68:     if _winapi:
jpayne@68:         def _close(self, _close=_multiprocessing.closesocket):
jpayne@68:             _close(self._handle)
jpayne@68:         _write = _multiprocessing.send
jpayne@68:         _read = _multiprocessing.recv
jpayne@68:     else:
jpayne@68:         def _close(self, _close=os.close):
jpayne@68:             _close(self._handle)
jpayne@68:         _write = os.write
jpayne@68:         _read = os.read
jpayne@68: 
jpayne@68:     def _send(self, buf, write=_write):
jpayne@68:         remaining = len(buf)
jpayne@68:         while True:
jpayne@68:             n = write(self._handle, buf)
jpayne@68:             remaining -= n
jpayne@68:             if remaining == 0:
jpayne@68:                 break
jpayne@68:             buf = buf[n:]
jpayne@68: 
jpayne@68:     def _recv(self, size, read=_read):
jpayne@68:         buf = io.BytesIO()
jpayne@68:         handle = self._handle
jpayne@68:         remaining = size
jpayne@68:         while remaining > 0:
jpayne@68:             chunk = read(handle, remaining)
jpayne@68:             n = len(chunk)
jpayne@68:             if n == 0:
jpayne@68:                 if remaining == size:
jpayne@68:                     raise EOFError
jpayne@68:                 else:
jpayne@68:                     raise OSError("got end of file during message")
jpayne@68:             buf.write(chunk)
jpayne@68:             remaining -= n
jpayne@68:         return buf
jpayne@68: 
jpayne@68:     def _send_bytes(self, buf):
jpayne@68:         n = len(buf)
jpayne@68:         if n > 0x7fffffff:
jpayne@68:             pre_header = struct.pack("!i", -1)
jpayne@68:             header = struct.pack("!Q", n)
jpayne@68:             self._send(pre_header)
jpayne@68:             self._send(header)
jpayne@68:             self._send(buf)
jpayne@68:         else:
jpayne@68:             # For wire compatibility with 3.7 and lower
jpayne@68:             header = struct.pack("!i", n)
jpayne@68:             if n > 16384:
jpayne@68:                 # The payload is large so Nagle's algorithm won't be triggered
jpayne@68:                 # and we'd better avoid the cost of concatenation.
jpayne@68:                 self._send(header)
jpayne@68:                 self._send(buf)
jpayne@68:             else:
jpayne@68:                 # Issue #20540: concatenate before sending, to avoid delays due
jpayne@68:                 # to Nagle's algorithm on a TCP socket.
jpayne@68:                 # Also note we want to avoid sending a 0-length buffer separately,
jpayne@68:                 # to avoid "broken pipe" errors if the other end closed the pipe.
jpayne@68:                 self._send(header + buf)
jpayne@68: 
jpayne@68:     def _recv_bytes(self, maxsize=None):
jpayne@68:         buf = self._recv(4)
jpayne@68:         size, = struct.unpack("!i", buf.getvalue())
jpayne@68:         if size == -1:
jpayne@68:             buf = self._recv(8)
jpayne@68:             size, = struct.unpack("!Q", buf.getvalue())
jpayne@68:         if maxsize is not None and size > maxsize:
jpayne@68:             return None
jpayne@68:         return self._recv(size)
jpayne@68: 
jpayne@68:     def _poll(self, timeout):
jpayne@68:         r = wait([self], timeout)
jpayne@68:         return bool(r)
jpayne@68: 
jpayne@68: 
jpayne@68: #
jpayne@68: # Public functions
jpayne@68: #
jpayne@68: 
jpayne@68: class Listener(object):
jpayne@68:     '''
jpayne@68:     Returns a listener object.
jpayne@68: 
jpayne@68:     This is a wrapper for a bound socket which is 'listening' for
jpayne@68:     connections, or for a Windows named pipe.
jpayne@68:     '''
jpayne@68:     def __init__(self, address=None, family=None, backlog=1, authkey=None):
jpayne@68:         family = family or (address and address_type(address)) \
jpayne@68:                  or default_family
jpayne@68:         address = address or arbitrary_address(family)
jpayne@68: 
jpayne@68:         _validate_family(family)
jpayne@68:         if family == 'AF_PIPE':
jpayne@68:             self._listener = PipeListener(address, backlog)
jpayne@68:         else:
jpayne@68:             self._listener = SocketListener(address, family, backlog)
jpayne@68: 
jpayne@68:         if authkey is not None and not isinstance(authkey, bytes):
jpayne@68:             raise TypeError('authkey should be a byte string')
jpayne@68: 
jpayne@68:         self._authkey = authkey
jpayne@68: 
jpayne@68:     def accept(self):
jpayne@68:         '''
jpayne@68:         Accept a connection on the bound socket or named pipe of `self`.
jpayne@68: 
jpayne@68:         Returns a `Connection` object.
jpayne@68:         '''
jpayne@68:         if self._listener is None:
jpayne@68:             raise OSError('listener is closed')
jpayne@68:         c = self._listener.accept()
jpayne@68:         if self._authkey:
jpayne@68:             deliver_challenge(c, self._authkey)
jpayne@68:             answer_challenge(c, self._authkey)
jpayne@68:         return c
jpayne@68: 
jpayne@68:     def close(self):
jpayne@68:         '''
jpayne@68:         Close the bound socket or named pipe of `self`.
jpayne@68:         '''
jpayne@68:         listener = self._listener
jpayne@68:         if listener is not None:
jpayne@68:             self._listener = None
jpayne@68:             listener.close()
jpayne@68: 
jpayne@68:     @property
jpayne@68:     def address(self):
jpayne@68:         return self._listener._address
jpayne@68: 
jpayne@68:     @property
jpayne@68:     def last_accepted(self):
jpayne@68:         return self._listener._last_accepted
jpayne@68: 
jpayne@68:     def __enter__(self):
jpayne@68:         return self
jpayne@68: 
jpayne@68:     def __exit__(self, exc_type, exc_value, exc_tb):
jpayne@68:         self.close()
jpayne@68: 
jpayne@68: 
jpayne@68: def Client(address, family=None, authkey=None):
jpayne@68:     '''
jpayne@68:     Returns a connection to the address of a `Listener`
jpayne@68:     '''
jpayne@68:     family = family or address_type(address)
jpayne@68:     _validate_family(family)
jpayne@68:     if family == 'AF_PIPE':
jpayne@68:         c = PipeClient(address)
jpayne@68:     else:
jpayne@68:         c = SocketClient(address)
jpayne@68: 
jpayne@68:     if authkey is not None and not isinstance(authkey, bytes):
jpayne@68:         raise TypeError('authkey should be a byte string')
jpayne@68: 
jpayne@68:     if authkey is not None:
jpayne@68:         answer_challenge(c, authkey)
jpayne@68:         deliver_challenge(c, authkey)
jpayne@68: 
jpayne@68:     return c
jpayne@68: 
jpayne@68: 
jpayne@68: if sys.platform != 'win32':
jpayne@68: 
jpayne@68:     def Pipe(duplex=True):
jpayne@68:         '''
jpayne@68:         Returns pair of connection objects at either end of a pipe
jpayne@68:         '''
jpayne@68:         if duplex:
jpayne@68:             s1, s2 = socket.socketpair()
jpayne@68:             s1.setblocking(True)
jpayne@68:             s2.setblocking(True)
jpayne@68:             c1 = Connection(s1.detach())
jpayne@68:             c2 = Connection(s2.detach())
jpayne@68:         else:
jpayne@68:             fd1, fd2 = os.pipe()
jpayne@68:             c1 = Connection(fd1, writable=False)
jpayne@68:             c2 = Connection(fd2, readable=False)
jpayne@68: 
jpayne@68:         return c1, c2
jpayne@68: 
jpayne@68: else:
jpayne@68: 
jpayne@68:     def Pipe(duplex=True):
jpayne@68:         '''
jpayne@68:         Returns pair of connection objects at either end of a pipe
jpayne@68:         '''
jpayne@68:         address = arbitrary_address('AF_PIPE')
jpayne@68:         if duplex:
jpayne@68:             openmode = _winapi.PIPE_ACCESS_DUPLEX
jpayne@68:             access = _winapi.GENERIC_READ | _winapi.GENERIC_WRITE
jpayne@68:             obsize, ibsize = BUFSIZE, BUFSIZE
jpayne@68:         else:
jpayne@68:             openmode = _winapi.PIPE_ACCESS_INBOUND
jpayne@68:             access = _winapi.GENERIC_WRITE
jpayne@68:             obsize, ibsize = 0, BUFSIZE
jpayne@68: 
jpayne@68:         h1 = _winapi.CreateNamedPipe(
jpayne@68:             address, openmode | _winapi.FILE_FLAG_OVERLAPPED |
jpayne@68:             _winapi.FILE_FLAG_FIRST_PIPE_INSTANCE,
jpayne@68:             _winapi.PIPE_TYPE_MESSAGE | _winapi.PIPE_READMODE_MESSAGE |
jpayne@68:             _winapi.PIPE_WAIT,
jpayne@68:             1, obsize, ibsize, _winapi.NMPWAIT_WAIT_FOREVER,
jpayne@68:             # default security descriptor: the handle cannot be inherited
jpayne@68:             _winapi.NULL
jpayne@68:             )
jpayne@68:         h2 = _winapi.CreateFile(
jpayne@68:             address, access, 0, _winapi.NULL, _winapi.OPEN_EXISTING,
jpayne@68:             _winapi.FILE_FLAG_OVERLAPPED, _winapi.NULL
jpayne@68:             )
jpayne@68:         _winapi.SetNamedPipeHandleState(
jpayne@68:             h2, _winapi.PIPE_READMODE_MESSAGE, None, None
jpayne@68:             )
jpayne@68: 
jpayne@68:         overlapped = _winapi.ConnectNamedPipe(h1, overlapped=True)
jpayne@68:         _, err = overlapped.GetOverlappedResult(True)
jpayne@68:         assert err == 0
jpayne@68: 
jpayne@68:         c1 = PipeConnection(h1, writable=duplex)
jpayne@68:         c2 = PipeConnection(h2, readable=duplex)
jpayne@68: 
jpayne@68:         return c1, c2
jpayne@68: 
jpayne@68: #
jpayne@68: # Definitions for connections based on sockets
jpayne@68: #
jpayne@68: 
jpayne@68: class SocketListener(object):
jpayne@68:     '''
jpayne@68:     Representation of a socket which is bound to an address and listening
jpayne@68:     '''
jpayne@68:     def __init__(self, address, family, backlog=1):
jpayne@68:         self._socket = socket.socket(getattr(socket, family))
jpayne@68:         try:
jpayne@68:             # SO_REUSEADDR has different semantics on Windows (issue #2550).
jpayne@68:             if os.name == 'posix':
jpayne@68:                 self._socket.setsockopt(socket.SOL_SOCKET,
jpayne@68:                                         socket.SO_REUSEADDR, 1)
jpayne@68:             self._socket.setblocking(True)
jpayne@68:             self._socket.bind(address)
jpayne@68:             self._socket.listen(backlog)
jpayne@68:             self._address = self._socket.getsockname()
jpayne@68:         except OSError:
jpayne@68:             self._socket.close()
jpayne@68:             raise
jpayne@68:         self._family = family
jpayne@68:         self._last_accepted = None
jpayne@68: 
jpayne@68:         if family == 'AF_UNIX':
jpayne@68:             self._unlink = util.Finalize(
jpayne@68:                 self, os.unlink, args=(address,), exitpriority=0
jpayne@68:                 )
jpayne@68:         else:
jpayne@68:             self._unlink = None
jpayne@68: 
jpayne@68:     def accept(self):
jpayne@68:         s, self._last_accepted = self._socket.accept()
jpayne@68:         s.setblocking(True)
jpayne@68:         return Connection(s.detach())
jpayne@68: 
jpayne@68:     def close(self):
jpayne@68:         try:
jpayne@68:             self._socket.close()
jpayne@68:         finally:
jpayne@68:             unlink = self._unlink
jpayne@68:             if unlink is not None:
jpayne@68:                 self._unlink = None
jpayne@68:                 unlink()
jpayne@68: 
jpayne@68: 
jpayne@68: def SocketClient(address):
jpayne@68:     '''
jpayne@68:     Return a connection object connected to the socket given by `address`
jpayne@68:     '''
jpayne@68:     family = address_type(address)
jpayne@68:     with socket.socket( getattr(socket, family) ) as s:
jpayne@68:         s.setblocking(True)
jpayne@68:         s.connect(address)
jpayne@68:         return Connection(s.detach())
jpayne@68: 
jpayne@68: #
jpayne@68: # Definitions for connections based on named pipes
jpayne@68: #
jpayne@68: 
jpayne@68: if sys.platform == 'win32':
jpayne@68: 
jpayne@68:     class PipeListener(object):
jpayne@68:         '''
jpayne@68:         Representation of a named pipe
jpayne@68:         '''
jpayne@68:         def __init__(self, address, backlog=None):
jpayne@68:             self._address = address
jpayne@68:             self._handle_queue = [self._new_handle(first=True)]
jpayne@68: 
jpayne@68:             self._last_accepted = None
jpayne@68:             util.sub_debug('listener created with address=%r', self._address)
jpayne@68:             self.close = util.Finalize(
jpayne@68:                 self, PipeListener._finalize_pipe_listener,
jpayne@68:                 args=(self._handle_queue, self._address), exitpriority=0
jpayne@68:                 )
jpayne@68: 
jpayne@68:         def _new_handle(self, first=False):
jpayne@68:             flags = _winapi.PIPE_ACCESS_DUPLEX | _winapi.FILE_FLAG_OVERLAPPED
jpayne@68:             if first:
jpayne@68:                 flags |= _winapi.FILE_FLAG_FIRST_PIPE_INSTANCE
jpayne@68:             return _winapi.CreateNamedPipe(
jpayne@68:                 self._address, flags,
jpayne@68:                 _winapi.PIPE_TYPE_MESSAGE | _winapi.PIPE_READMODE_MESSAGE |
jpayne@68:                 _winapi.PIPE_WAIT,
jpayne@68:                 _winapi.PIPE_UNLIMITED_INSTANCES, BUFSIZE, BUFSIZE,
jpayne@68:                 _winapi.NMPWAIT_WAIT_FOREVER, _winapi.NULL
jpayne@68:                 )
jpayne@68: 
jpayne@68:         def accept(self):
jpayne@68:             self._handle_queue.append(self._new_handle())
jpayne@68:             handle = self._handle_queue.pop(0)
jpayne@68:             try:
jpayne@68:                 ov = _winapi.ConnectNamedPipe(handle, overlapped=True)
jpayne@68:             except OSError as e:
jpayne@68:                 if e.winerror != _winapi.ERROR_NO_DATA:
jpayne@68:                     raise
jpayne@68:                 # ERROR_NO_DATA can occur if a client has already connected,
jpayne@68:                 # written data and then disconnected -- see Issue 14725.
jpayne@68:             else:
jpayne@68:                 try:
jpayne@68:                     res = _winapi.WaitForMultipleObjects(
jpayne@68:                         [ov.event], False, INFINITE)
jpayne@68:                 except:
jpayne@68:                     ov.cancel()
jpayne@68:                     _winapi.CloseHandle(handle)
jpayne@68:                     raise
jpayne@68:                 finally:
jpayne@68:                     _, err = ov.GetOverlappedResult(True)
jpayne@68:                     assert err == 0
jpayne@68:             return PipeConnection(handle)
jpayne@68: 
jpayne@68:         @staticmethod
jpayne@68:         def _finalize_pipe_listener(queue, address):
jpayne@68:             util.sub_debug('closing listener with address=%r', address)
jpayne@68:             for handle in queue:
jpayne@68:                 _winapi.CloseHandle(handle)
jpayne@68: 
jpayne@68:     def PipeClient(address):
jpayne@68:         '''
jpayne@68:         Return a connection object connected to the pipe given by `address`
jpayne@68:         '''
jpayne@68:         t = _init_timeout()
jpayne@68:         while 1:
jpayne@68:             try:
jpayne@68:                 _winapi.WaitNamedPipe(address, 1000)
jpayne@68:                 h = _winapi.CreateFile(
jpayne@68:                     address, _winapi.GENERIC_READ | _winapi.GENERIC_WRITE,
jpayne@68:                     0, _winapi.NULL, _winapi.OPEN_EXISTING,
jpayne@68:                     _winapi.FILE_FLAG_OVERLAPPED, _winapi.NULL
jpayne@68:                     )
jpayne@68:             except OSError as e:
jpayne@68:                 if e.winerror not in (_winapi.ERROR_SEM_TIMEOUT,
jpayne@68:                                       _winapi.ERROR_PIPE_BUSY) or _check_timeout(t):
jpayne@68:                     raise
jpayne@68:             else:
jpayne@68:                 break
jpayne@68:         else:
jpayne@68:             raise
jpayne@68: 
jpayne@68:         _winapi.SetNamedPipeHandleState(
jpayne@68:             h, _winapi.PIPE_READMODE_MESSAGE, None, None
jpayne@68:             )
jpayne@68:         return PipeConnection(h)
jpayne@68: 
jpayne@68: #
jpayne@68: # Authentication stuff
jpayne@68: #
jpayne@68: 
jpayne@68: MESSAGE_LENGTH = 20
jpayne@68: 
jpayne@68: CHALLENGE = b'#CHALLENGE#'
jpayne@68: WELCOME = b'#WELCOME#'
jpayne@68: FAILURE = b'#FAILURE#'
jpayne@68: 
jpayne@68: def deliver_challenge(connection, authkey):
jpayne@68:     import hmac
jpayne@68:     if not isinstance(authkey, bytes):
jpayne@68:         raise ValueError(
jpayne@68:             "Authkey must be bytes, not {0!s}".format(type(authkey)))
jpayne@68:     message = os.urandom(MESSAGE_LENGTH)
jpayne@68:     connection.send_bytes(CHALLENGE + message)
jpayne@68:     digest = hmac.new(authkey, message, 'md5').digest()
jpayne@68:     response = connection.recv_bytes(256)        # reject large message
jpayne@68:     if response == digest:
jpayne@68:         connection.send_bytes(WELCOME)
jpayne@68:     else:
jpayne@68:         connection.send_bytes(FAILURE)
jpayne@68:         raise AuthenticationError('digest received was wrong')
jpayne@68: 
jpayne@68: def answer_challenge(connection, authkey):
jpayne@68:     import hmac
jpayne@68:     if not isinstance(authkey, bytes):
jpayne@68:         raise ValueError(
jpayne@68:             "Authkey must be bytes, not {0!s}".format(type(authkey)))
jpayne@68:     message = connection.recv_bytes(256)         # reject large message
jpayne@68:     assert message[:len(CHALLENGE)] == CHALLENGE, 'message = %r' % message
jpayne@68:     message = message[len(CHALLENGE):]
jpayne@68:     digest = hmac.new(authkey, message, 'md5').digest()
jpayne@68:     connection.send_bytes(digest)
jpayne@68:     response = connection.recv_bytes(256)        # reject large message
jpayne@68:     if response != WELCOME:
jpayne@68:         raise AuthenticationError('digest sent was rejected')
jpayne@68: 
jpayne@68: #
jpayne@68: # Support for using xmlrpclib for serialization
jpayne@68: #
jpayne@68: 
jpayne@68: class ConnectionWrapper(object):
jpayne@68:     def __init__(self, conn, dumps, loads):
jpayne@68:         self._conn = conn
jpayne@68:         self._dumps = dumps
jpayne@68:         self._loads = loads
jpayne@68:         for attr in ('fileno', 'close', 'poll', 'recv_bytes', 'send_bytes'):
jpayne@68:             obj = getattr(conn, attr)
jpayne@68:             setattr(self, attr, obj)
jpayne@68:     def send(self, obj):
jpayne@68:         s = self._dumps(obj)
jpayne@68:         self._conn.send_bytes(s)
jpayne@68:     def recv(self):
jpayne@68:         s = self._conn.recv_bytes()
jpayne@68:         return self._loads(s)
jpayne@68: 
jpayne@68: def _xml_dumps(obj):
jpayne@68:     return xmlrpclib.dumps((obj,), None, None, None, 1).encode('utf-8')
jpayne@68: 
jpayne@68: def _xml_loads(s):
jpayne@68:     (obj,), method = xmlrpclib.loads(s.decode('utf-8'))
jpayne@68:     return obj
jpayne@68: 
jpayne@68: class XmlListener(Listener):
jpayne@68:     def accept(self):
jpayne@68:         global xmlrpclib
jpayne@68:         import xmlrpc.client as xmlrpclib
jpayne@68:         obj = Listener.accept(self)
jpayne@68:         return ConnectionWrapper(obj, _xml_dumps, _xml_loads)
jpayne@68: 
jpayne@68: def XmlClient(*args, **kwds):
jpayne@68:     global xmlrpclib
jpayne@68:     import xmlrpc.client as xmlrpclib
jpayne@68:     return ConnectionWrapper(Client(*args, **kwds), _xml_dumps, _xml_loads)
jpayne@68: 
jpayne@68: #
jpayne@68: # Wait
jpayne@68: #
jpayne@68: 
jpayne@68: if sys.platform == 'win32':
jpayne@68: 
jpayne@68:     def _exhaustive_wait(handles, timeout):
jpayne@68:         # Return ALL handles which are currently signalled.  (Only
jpayne@68:         # returning the first signalled might create starvation issues.)
jpayne@68:         L = list(handles)
jpayne@68:         ready = []
jpayne@68:         while L:
jpayne@68:             res = _winapi.WaitForMultipleObjects(L, False, timeout)
jpayne@68:             if res == WAIT_TIMEOUT:
jpayne@68:                 break
jpayne@68:             elif WAIT_OBJECT_0 <= res < WAIT_OBJECT_0 + len(L):
jpayne@68:                 res -= WAIT_OBJECT_0
jpayne@68:             elif WAIT_ABANDONED_0 <= res < WAIT_ABANDONED_0 + len(L):
jpayne@68:                 res -= WAIT_ABANDONED_0
jpayne@68:             else:
jpayne@68:                 raise RuntimeError('Should not get here')
jpayne@68:             ready.append(L[res])
jpayne@68:             L = L[res+1:]
jpayne@68:             timeout = 0
jpayne@68:         return ready
jpayne@68: 
jpayne@68:     _ready_errors = {_winapi.ERROR_BROKEN_PIPE, _winapi.ERROR_NETNAME_DELETED}
jpayne@68: 
jpayne@68:     def wait(object_list, timeout=None):
jpayne@68:         '''
jpayne@68:         Wait till an object in object_list is ready/readable.
jpayne@68: 
jpayne@68:         Returns list of those objects in object_list which are ready/readable.
jpayne@68:         '''
jpayne@68:         if timeout is None:
jpayne@68:             timeout = INFINITE
jpayne@68:         elif timeout < 0:
jpayne@68:             timeout = 0
jpayne@68:         else:
jpayne@68:             timeout = int(timeout * 1000 + 0.5)
jpayne@68: 
jpayne@68:         object_list = list(object_list)
jpayne@68:         waithandle_to_obj = {}
jpayne@68:         ov_list = []
jpayne@68:         ready_objects = set()
jpayne@68:         ready_handles = set()
jpayne@68: 
jpayne@68:         try:
jpayne@68:             for o in object_list:
jpayne@68:                 try:
jpayne@68:                     fileno = getattr(o, 'fileno')
jpayne@68:                 except AttributeError:
jpayne@68:                     waithandle_to_obj[o.__index__()] = o
jpayne@68:                 else:
jpayne@68:                     # start an overlapped read of length zero
jpayne@68:                     try:
jpayne@68:                         ov, err = _winapi.ReadFile(fileno(), 0, True)
jpayne@68:                     except OSError as e:
jpayne@68:                         ov, err = None, e.winerror
jpayne@68:                         if err not in _ready_errors:
jpayne@68:                             raise
jpayne@68:                     if err == _winapi.ERROR_IO_PENDING:
jpayne@68:                         ov_list.append(ov)
jpayne@68:                         waithandle_to_obj[ov.event] = o
jpayne@68:                     else:
jpayne@68:                         # If o.fileno() is an overlapped pipe handle and
jpayne@68:                         # err == 0 then there is a zero length message
jpayne@68:                         # in the pipe, but it HAS NOT been consumed...
jpayne@68:                         if ov and sys.getwindowsversion()[:2] >= (6, 2):
jpayne@68:                             # ... except on Windows 8 and later, where
jpayne@68:                             # the message HAS been consumed.
jpayne@68:                             try:
jpayne@68:                                 _, err = ov.GetOverlappedResult(False)
jpayne@68:                             except OSError as e:
jpayne@68:                                 err = e.winerror
jpayne@68:                             if not err and hasattr(o, '_got_empty_message'):
jpayne@68:                                 o._got_empty_message = True
jpayne@68:                         ready_objects.add(o)
jpayne@68:                         timeout = 0
jpayne@68: 
jpayne@68:             ready_handles = _exhaustive_wait(waithandle_to_obj.keys(), timeout)
jpayne@68:         finally:
jpayne@68:             # request that overlapped reads stop
jpayne@68:             for ov in ov_list:
jpayne@68:                 ov.cancel()
jpayne@68: 
jpayne@68:             # wait for all overlapped reads to stop
jpayne@68:             for ov in ov_list:
jpayne@68:                 try:
jpayne@68:                     _, err = ov.GetOverlappedResult(True)
jpayne@68:                 except OSError as e:
jpayne@68:                     err = e.winerror
jpayne@68:                     if err not in _ready_errors:
jpayne@68:                         raise
jpayne@68:                 if err != _winapi.ERROR_OPERATION_ABORTED:
jpayne@68:                     o = waithandle_to_obj[ov.event]
jpayne@68:                     ready_objects.add(o)
jpayne@68:                     if err == 0:
jpayne@68:                         # If o.fileno() is an overlapped pipe handle then
jpayne@68:                         # a zero length message HAS been consumed.
jpayne@68:                         if hasattr(o, '_got_empty_message'):
jpayne@68:                             o._got_empty_message = True
jpayne@68: 
jpayne@68:         ready_objects.update(waithandle_to_obj[h] for h in ready_handles)
jpayne@68:         return [o for o in object_list if o in ready_objects]
jpayne@68: 
jpayne@68: else:
jpayne@68: 
jpayne@68:     import selectors
jpayne@68: 
jpayne@68:     # poll/select have the advantage of not requiring any extra file
jpayne@68:     # descriptor, contrarily to epoll/kqueue (also, they require a single
jpayne@68:     # syscall).
jpayne@68:     if hasattr(selectors, 'PollSelector'):
jpayne@68:         _WaitSelector = selectors.PollSelector
jpayne@68:     else:
jpayne@68:         _WaitSelector = selectors.SelectSelector
jpayne@68: 
jpayne@68:     def wait(object_list, timeout=None):
jpayne@68:         '''
jpayne@68:         Wait till an object in object_list is ready/readable.
jpayne@68: 
jpayne@68:         Returns list of those objects in object_list which are ready/readable.
jpayne@68:         '''
jpayne@68:         with _WaitSelector() as selector:
jpayne@68:             for obj in object_list:
jpayne@68:                 selector.register(obj, selectors.EVENT_READ)
jpayne@68: 
jpayne@68:             if timeout is not None:
jpayne@68:                 deadline = time.monotonic() + timeout
jpayne@68: 
jpayne@68:             while True:
jpayne@68:                 ready = selector.select(timeout)
jpayne@68:                 if ready:
jpayne@68:                     return [key.fileobj for (key, events) in ready]
jpayne@68:                 else:
jpayne@68:                     if timeout is not None:
jpayne@68:                         timeout = deadline - time.monotonic()
jpayne@68:                         if timeout < 0:
jpayne@68:                             return ready
jpayne@68: 
jpayne@68: #
jpayne@68: # Make connection and socket objects sharable if possible
jpayne@68: #
jpayne@68: 
jpayne@68: if sys.platform == 'win32':
jpayne@68:     def reduce_connection(conn):
jpayne@68:         handle = conn.fileno()
jpayne@68:         with socket.fromfd(handle, socket.AF_INET, socket.SOCK_STREAM) as s:
jpayne@68:             from . import resource_sharer
jpayne@68:             ds = resource_sharer.DupSocket(s)
jpayne@68:             return rebuild_connection, (ds, conn.readable, conn.writable)
jpayne@68:     def rebuild_connection(ds, readable, writable):
jpayne@68:         sock = ds.detach()
jpayne@68:         return Connection(sock.detach(), readable, writable)
jpayne@68:     reduction.register(Connection, reduce_connection)
jpayne@68: 
jpayne@68:     def reduce_pipe_connection(conn):
jpayne@68:         access = ((_winapi.FILE_GENERIC_READ if conn.readable else 0) |
jpayne@68:                   (_winapi.FILE_GENERIC_WRITE if conn.writable else 0))
jpayne@68:         dh = reduction.DupHandle(conn.fileno(), access)
jpayne@68:         return rebuild_pipe_connection, (dh, conn.readable, conn.writable)
jpayne@68:     def rebuild_pipe_connection(dh, readable, writable):
jpayne@68:         handle = dh.detach()
jpayne@68:         return PipeConnection(handle, readable, writable)
jpayne@68:     reduction.register(PipeConnection, reduce_pipe_connection)
jpayne@68: 
jpayne@68: else:
jpayne@68:     def reduce_connection(conn):
jpayne@68:         df = reduction.DupFd(conn.fileno())
jpayne@68:         return rebuild_connection, (df, conn.readable, conn.writable)
jpayne@68:     def rebuild_connection(df, readable, writable):
jpayne@68:         fd = df.detach()
jpayne@68:         return Connection(fd, readable, writable)
jpayne@68:     reduction.register(Connection, reduce_connection)