init

.venv/lib/python3.8/site-packages/mypy/ipc.py

@@ -0,0 +1,268 @@
+"""Cross platform abstractions for inter-process communication
+
+On Unix, this uses AF_UNIX sockets.
+On Windows, this uses NamedPipes.
+"""
+
+import base64
+import os
+import shutil
+import sys
+import tempfile
+
+from typing import Optional, Callable
+from typing_extensions import Final, Type
+
+from types import TracebackType
+
+if sys.platform == 'win32':
+    # This may be private, but it is needed for IPC on Windows, and is basically stable
+    import _winapi
+    import ctypes
+
+    _IPCHandle = int
+
+    kernel32 = ctypes.windll.kernel32
+    DisconnectNamedPipe: Callable[[_IPCHandle], int] = kernel32.DisconnectNamedPipe
+    FlushFileBuffers: Callable[[_IPCHandle], int] = kernel32.FlushFileBuffers
+else:
+    import socket
+    _IPCHandle = socket.socket
+
+
+class IPCException(Exception):
+    """Exception for IPC issues."""
+    pass
+
+
+class IPCBase:
+    """Base class for communication between the dmypy client and server.
+
+    This contains logic shared between the client and server, such as reading
+    and writing.
+    """
+
+    connection: _IPCHandle
+
+    def __init__(self, name: str, timeout: Optional[float]) -> None:
+        self.name = name
+        self.timeout = timeout
+
+    def read(self, size: int = 100000) -> bytes:
+        """Read bytes from an IPC connection until its empty."""
+        bdata = bytearray()
+        if sys.platform == 'win32':
+            while True:
+                ov, err = _winapi.ReadFile(self.connection, size, overlapped=True)
+                try:
+                    if err == _winapi.ERROR_IO_PENDING:
+                        timeout = int(self.timeout * 1000) if self.timeout else _winapi.INFINITE
+                        res = _winapi.WaitForSingleObject(ov.event, timeout)
+                        if res != _winapi.WAIT_OBJECT_0:
+                            raise IPCException("Bad result from I/O wait: {}".format(res))
+                except BaseException:
+                    ov.cancel()
+                    raise
+                _, err = ov.GetOverlappedResult(True)
+                more = ov.getbuffer()
+                if more:
+                    bdata.extend(more)
+                if err == 0:
+                    # we are done!
+                    break
+                elif err == _winapi.ERROR_MORE_DATA:
+                    # read again
+                    continue
+                elif err == _winapi.ERROR_OPERATION_ABORTED:
+                    raise IPCException("ReadFile operation aborted.")
+        else:
+            while True:
+                more = self.connection.recv(size)
+                if not more:
+                    break
+                bdata.extend(more)
+        return bytes(bdata)
+
+    def write(self, data: bytes) -> None:
+        """Write bytes to an IPC connection."""
+        if sys.platform == 'win32':
+            try:
+                ov, err = _winapi.WriteFile(self.connection, data, overlapped=True)
+                # TODO: remove once typeshed supports Literal types
+                assert isinstance(ov, _winapi.Overlapped)
+                assert isinstance(err, int)
+                try:
+                    if err == _winapi.ERROR_IO_PENDING:
+                        timeout = int(self.timeout * 1000) if self.timeout else _winapi.INFINITE
+                        res = _winapi.WaitForSingleObject(ov.event, timeout)
+                        if res != _winapi.WAIT_OBJECT_0:
+                            raise IPCException("Bad result from I/O wait: {}".format(res))
+                    elif err != 0:
+                        raise IPCException("Failed writing to pipe with error: {}".format(err))
+                except BaseException:
+                    ov.cancel()
+                    raise
+                bytes_written, err = ov.GetOverlappedResult(True)
+                assert err == 0, err
+                assert bytes_written == len(data)
+            except WindowsError as e:
+                raise IPCException("Failed to write with error: {}".format(e.winerror)) from e
+        else:
+            self.connection.sendall(data)
+            self.connection.shutdown(socket.SHUT_WR)
+
+    def close(self) -> None:
+        if sys.platform == 'win32':
+            if self.connection != _winapi.NULL:
+                _winapi.CloseHandle(self.connection)
+        else:
+            self.connection.close()
+
+
+class IPCClient(IPCBase):
+    """The client side of an IPC connection."""
+
+    def __init__(self, name: str, timeout: Optional[float]) -> None:
+        super().__init__(name, timeout)
+        if sys.platform == 'win32':
+            timeout = int(self.timeout * 1000) if self.timeout else _winapi.NMPWAIT_WAIT_FOREVER
+            try:
+                _winapi.WaitNamedPipe(self.name, timeout)
+            except FileNotFoundError as e:
+                raise IPCException("The NamedPipe at {} was not found.".format(self.name)) from e
+            except WindowsError as e:
+                if e.winerror == _winapi.ERROR_SEM_TIMEOUT:
+                    raise IPCException("Timed out waiting for connection.") from e
+                else:
+                    raise
+            try:
+                self.connection = _winapi.CreateFile(
+                    self.name,
+                    _winapi.GENERIC_READ | _winapi.GENERIC_WRITE,
+                    0,
+                    _winapi.NULL,
+                    _winapi.OPEN_EXISTING,
+                    _winapi.FILE_FLAG_OVERLAPPED,
+                    _winapi.NULL,
+                )
+            except WindowsError as e:
+                if e.winerror == _winapi.ERROR_PIPE_BUSY:
+                    raise IPCException("The connection is busy.") from e
+                else:
+                    raise
+            _winapi.SetNamedPipeHandleState(self.connection,
+                                            _winapi.PIPE_READMODE_MESSAGE,
+                                            None,
+                                            None)
+        else:
+            self.connection = socket.socket(socket.AF_UNIX)
+            self.connection.settimeout(timeout)
+            self.connection.connect(name)
+
+    def __enter__(self) -> 'IPCClient':
+        return self
+
+    def __exit__(self,
+                 exc_ty: 'Optional[Type[BaseException]]' = None,
+                 exc_val: Optional[BaseException] = None,
+                 exc_tb: Optional[TracebackType] = None,
+                 ) -> None:
+        self.close()
+
+
+class IPCServer(IPCBase):
+
+    BUFFER_SIZE: Final = 2 ** 16
+
+    def __init__(self, name: str, timeout: Optional[float] = None) -> None:
+        if sys.platform == 'win32':
+            name = r'\\.\pipe\{}-{}.pipe'.format(
+                name, base64.urlsafe_b64encode(os.urandom(6)).decode())
+        else:
+            name = '{}.sock'.format(name)
+        super().__init__(name, timeout)
+        if sys.platform == 'win32':
+            self.connection = _winapi.CreateNamedPipe(self.name,
+                _winapi.PIPE_ACCESS_DUPLEX
+                | _winapi.FILE_FLAG_FIRST_PIPE_INSTANCE
+                | _winapi.FILE_FLAG_OVERLAPPED,
+                _winapi.PIPE_READMODE_MESSAGE
+                | _winapi.PIPE_TYPE_MESSAGE
+                | _winapi.PIPE_WAIT
+                | 0x8,  # PIPE_REJECT_REMOTE_CLIENTS
+                1,  # one instance
+                self.BUFFER_SIZE,
+                self.BUFFER_SIZE,
+                _winapi.NMPWAIT_WAIT_FOREVER,
+                0,  # Use default security descriptor
+                                                      )
+            if self.connection == -1:  # INVALID_HANDLE_VALUE
+                err = _winapi.GetLastError()
+                raise IPCException('Invalid handle to pipe: {}'.format(err))
+        else:
+            self.sock_directory = tempfile.mkdtemp()
+            sockfile = os.path.join(self.sock_directory, self.name)
+            self.sock = socket.socket(socket.AF_UNIX)
+            self.sock.bind(sockfile)
+            self.sock.listen(1)
+            if timeout is not None:
+                self.sock.settimeout(timeout)
+
+    def __enter__(self) -> 'IPCServer':
+        if sys.platform == 'win32':
+            # NOTE: It is theoretically possible that this will hang forever if the
+            # client never connects, though this can be "solved" by killing the server
+            try:
+                ov = _winapi.ConnectNamedPipe(self.connection, overlapped=True)
+                # TODO: remove once typeshed supports Literal types
+                assert isinstance(ov, _winapi.Overlapped)
+            except WindowsError as e:
+                # Don't raise if the client already exists, or the client already connected
+                if e.winerror not in (_winapi.ERROR_PIPE_CONNECTED, _winapi.ERROR_NO_DATA):
+                    raise
+            else:
+                try:
+                    timeout = int(self.timeout * 1000) if self.timeout else _winapi.INFINITE
+                    res = _winapi.WaitForSingleObject(ov.event, timeout)
+                    assert res == _winapi.WAIT_OBJECT_0
+                except BaseException:
+                    ov.cancel()
+                    _winapi.CloseHandle(self.connection)
+                    raise
+                _, err = ov.GetOverlappedResult(True)
+                assert err == 0
+        else:
+            try:
+                self.connection, _ = self.sock.accept()
+            except socket.timeout as e:
+                raise IPCException('The socket timed out') from e
+        return self
+
+    def __exit__(self,
+                 exc_ty: 'Optional[Type[BaseException]]' = None,
+                 exc_val: Optional[BaseException] = None,
+                 exc_tb: Optional[TracebackType] = None,
+                 ) -> None:
+        if sys.platform == 'win32':
+            try:
+                # Wait for the client to finish reading the last write before disconnecting
+                if not FlushFileBuffers(self.connection):
+                    raise IPCException("Failed to flush NamedPipe buffer,"
+                                       "maybe the client hung up?")
+            finally:
+                DisconnectNamedPipe(self.connection)
+        else:
+            self.close()
+
+    def cleanup(self) -> None:
+        if sys.platform == 'win32':
+            self.close()
+        else:
+            shutil.rmtree(self.sock_directory)
+
+    @property
+    def connection_name(self) -> str:
+        if sys.platform == 'win32':
+            return self.name
+        else:
+            return self.sock.getsockname()