kafka-python 3.0.0__py3-none-any.whl

kafka/net/socks5.py

@@ -0,0 +1,262 @@
+import errno
+import logging
+import random
+import socket
+import struct
+from urllib.parse import urlparse
+
+from kafka.errors import KafkaConnectionError
+from kafka.net.inet import KafkaNetSocket
+
+
+log = logging.getLogger(__name__)
+
+
+class ProxyConnectionStates:
+    DISCONNECTED = '<disconnected>'
+    CONNECTING = '<connecting>'
+    NEGOTIATE_PROPOSE = '<negotiate_propose>'
+    NEGOTIATING = '<negotiating>'
+    AUTHENTICATING = '<authenticating>'
+    REQUEST_SUBMIT = '<request_submit>'
+    REQUESTING = '<requesting>'
+    READ_ADDRESS = '<read_address>'
+    COMPLETE = '<complete>'
+
+
+class Socks5Proxy(KafkaNetSocket):
+    """Socks5 proxy
+
+    Manages connection through socks5 proxy with support for username/password
+    authentication.
+    """
+    # socks5h for remote dns
+    SCHEMES = ('socks5', 'socks5h')
+
+    def __init__(self, proxy_url):
+        self._buffer_in = b''
+        self._buffer_out = b''
+        self._proxy_url = urlparse(proxy_url)
+        if self._proxy_url.scheme not in self.SCHEMES:
+            raise ValueError('Unsupported proxy scheme: %s' % (self._proxy_url.scheme,))
+        self._sock = None
+        self._state = ProxyConnectionStates.DISCONNECTED
+        self._target_afi = socket.AF_UNSPEC
+        self._proxy_addr = self._get_proxy_addr()
+
+    def _get_proxy_addr(self):
+        proxy_addrs = self.dns_lookup(self._proxy_url.hostname, self._proxy_url.port, proxy=True)
+        if not proxy_addrs:
+            raise KafkaConnectionError('Unable to resolve proxy_url via dns')
+        return random.choice(proxy_addrs)
+
+    def _use_remote_lookup(self):
+        return self._proxy_url.scheme == 'socks5h'
+
+    def dns_lookup(self, host, port, proxy=False):
+        if proxy:
+            return super().dns_lookup(host, port, raise_error=True)
+        elif self._use_remote_lookup():
+            return [(socket.AF_UNSPEC, socket.SOCK_STREAM, socket.IPPROTO_TCP, '', (host, port))]
+        else:
+            return super().dns_lookup(host, port)
+
+    def socket(self, family=socket.AF_UNSPEC, sock_type=socket.SOCK_STREAM, proto=socket.IPPROTO_TCP):
+        """Open and record a socket.
+
+        Returns the actual underlying socket
+        object to ensure e.g. selects and ssl wrapping works as expected.
+        """
+        self._target_afi = family  # Store the address family of the target
+        proxy_family, _, _, _, _ = self._proxy_addr
+        self._sock = socket.socket(proxy_family, sock_type, proto)
+        return self._sock
+
+    def _flush_buf(self):
+        """Send out all data that is stored in the outgoing buffer.
+
+        It is expected that the caller handles error handling, including non-blocking
+        as well as connection failure exceptions.
+        """
+        while self._buffer_out:
+            sent_bytes = self._sock.send(self._buffer_out)
+            self._buffer_out = self._buffer_out[sent_bytes:]
+
+    def _peek_buf(self, datalen):
+        """Ensure local inbound buffer has enough data, and return that data without
+        consuming the local buffer
+
+        It's expected that the caller handles e.g. blocking exceptions"""
+        while True:
+            bytes_remaining = datalen - len(self._buffer_in)
+            if bytes_remaining <= 0:
+                break
+            data = self._sock.recv(bytes_remaining)
+            if not data:
+                break
+            self._buffer_in = self._buffer_in + data
+
+        return self._buffer_in[:datalen]
+
+    def _read_buf(self, datalen):
+        """Read and consume bytes from socket connection
+
+        It's expected that the caller handles e.g. blocking exceptions"""
+        buf = self._peek_buf(datalen)
+        if buf:
+            self._buffer_in = self._buffer_in[len(buf):]
+        return buf
+
+    def connect_ex(self, sock, addr):
+        """Runs a state machine through connection to authentication to
+        proxy connection request.
+
+        The somewhat strange setup is to facilitate non-intrusive use from
+        BrokerConnection state machine.
+
+        This function is called with a socket in non-blocking mode. Both
+        send and receive calls can return in EWOULDBLOCK/EAGAIN which we
+        specifically avoid handling here. These are handled in main
+        BrokerConnection connection loop, which then would retry calls
+        to this function."""
+        assert sock is self._sock
+        if self._state == ProxyConnectionStates.DISCONNECTED:
+            self._state = ProxyConnectionStates.CONNECTING
+
+        if self._state == ProxyConnectionStates.CONNECTING:
+            _, _, _, _, sockaddr = self._proxy_addr
+            ret = self._sock.connect_ex(sockaddr)
+            if not ret or ret == errno.EISCONN:
+                self._state = ProxyConnectionStates.NEGOTIATE_PROPOSE
+            else:
+                return ret
+
+        if self._state == ProxyConnectionStates.NEGOTIATE_PROPOSE:
+            if self._proxy_url.username and self._proxy_url.password:
+                # Propose username/password
+                self._buffer_out = b"\x05\x01\x02"
+            else:
+                # Propose no auth
+                self._buffer_out = b"\x05\x01\x00"
+            self._state = ProxyConnectionStates.NEGOTIATING
+
+        if self._state == ProxyConnectionStates.NEGOTIATING:
+            self._flush_buf()
+            buf = self._read_buf(2)
+            if buf[0:1] != b"\x05":
+                log.error("Unrecognized SOCKS version")
+                self._state = ProxyConnectionStates.DISCONNECTED
+                self._sock.close()
+                return errno.ECONNREFUSED
+
+            if buf[1:2] == b"\x00":
+                # No authentication required
+                self._state = ProxyConnectionStates.REQUEST_SUBMIT
+            elif buf[1:2] == b"\x02":
+                # Username/password authentication selected
+                userlen = len(self._proxy_url.username)
+                passlen = len(self._proxy_url.password)
+                self._buffer_out = struct.pack(
+                    "!bb{}sb{}s".format(userlen, passlen),
+                    1,  # version
+                    userlen,
+                    self._proxy_url.username.encode(),
+                    passlen,
+                    self._proxy_url.password.encode(),
+                )
+                self._state = ProxyConnectionStates.AUTHENTICATING
+            else:
+                log.error("Unrecognized SOCKS authentication method")
+                self._state = ProxyConnectionStates.DISCONNECTED
+                self._sock.close()
+                return errno.ECONNREFUSED
+
+        if self._state == ProxyConnectionStates.AUTHENTICATING:
+            self._flush_buf()
+            buf = self._read_buf(2)
+            if buf == b"\x01\x00":
+                # Authentication succesful
+                self._state = ProxyConnectionStates.REQUEST_SUBMIT
+            else:
+                log.error("Socks5 proxy authentication failure")
+                self._state = ProxyConnectionStates.DISCONNECTED
+                self._sock.close()
+                return errno.ECONNREFUSED
+
+        if self._state == ProxyConnectionStates.REQUEST_SUBMIT:
+            if self._use_remote_lookup():
+                addr_type = 3
+                addr_len = len(addr[0])
+            elif self._target_afi == socket.AF_INET:
+                addr_type = 1
+                addr_len = 4
+            elif self._target_afi == socket.AF_INET6:
+                addr_type = 4
+                addr_len = 16
+            else:
+                log.error("Unknown address family, %r", self._target_afi)
+                self._state = ProxyConnectionStates.DISCONNECTED
+                self._sock.close()
+                return errno.ECONNREFUSED
+
+            self._buffer_out = struct.pack(
+                "!bbbb",
+                5,  # version
+                1,  # command: connect
+                0,  # reserved
+                addr_type,  # 1 for ipv4, 4 for ipv6 address, 3 for domain name
+            )
+            # Addr format depends on type
+            if addr_type == 3:
+                # len + domain name (no null terminator)
+                self._buffer_out += struct.pack(
+                    "!b{}s".format(addr_len),
+                    addr_len,
+                    addr[0].encode('ascii'),
+                )
+            else:
+                # either 4 (type 1) or 16 (type 4) bytes of actual address
+                self._buffer_out += struct.pack(
+                    "!{}s".format(addr_len),
+                    socket.inet_pton(self._target_afi, addr[0]),
+                )
+            self._buffer_out += struct.pack("!H", addr[1])  # port
+
+            self._state = ProxyConnectionStates.REQUESTING
+
+        if self._state == ProxyConnectionStates.REQUESTING:
+            self._flush_buf()
+            buf = self._read_buf(2)
+            if buf[0:2] == b"\x05\x00":
+                self._state = ProxyConnectionStates.READ_ADDRESS
+            else:
+                log.error("Proxy request failed: %r", buf[1:2])
+                self._state = ProxyConnectionStates.DISCONNECTED
+                self._sock.close()
+                return errno.ECONNREFUSED
+
+        if self._state == ProxyConnectionStates.READ_ADDRESS:
+            # we don't really care about the remote endpoint address, but need to clear the stream
+            buf = self._peek_buf(2)
+            if buf[0:2] == b"\x00\x01":
+                _ = self._read_buf(2 + 4 + 2)  # ipv4 address + port
+            elif buf[0:2] == b"\x00\x05":
+                _ = self._read_buf(2 + 16 + 2)  # ipv6 address + port
+            else:
+                log.error("Unrecognized remote address type %r", buf[1:2])
+                self._state = ProxyConnectionStates.DISCONNECTED
+                self._sock.close()
+                return errno.ECONNREFUSED
+            self._state = ProxyConnectionStates.COMPLETE
+
+        if self._state == ProxyConnectionStates.COMPLETE:
+            return 0
+
+        # not reached;
+        # Send and recv will raise socket error on EWOULDBLOCK/EAGAIN that is assumed to be handled by
+        # the caller. The caller re-enters this state machine from retry logic with timer or via select & family
+        log.error("Internal error, state %r not handled correctly", self._state)
+        self._state = ProxyConnectionStates.DISCONNECTED
+        if self._sock:
+            self._sock.close()
+        return errno.ECONNREFUSED

kafka/net/transport.py

@@ -0,0 +1,415 @@
+from collections import deque
+import logging
+import selectors
+import socket
+import ssl
+import time
+
+import kafka.errors as Errors
+
+
+log = logging.getLogger(__name__)
+
+
+class KafkaTCPTransport:
+    def __init__(self, net, sock, host=None):
+        self._net = net
+        self._sock = sock
+        self.host = host
+        self._closed = False
+        self._write_buffer = deque()
+        self._writing = False
+        self._protocol = None
+        self._read = False
+        self._write = True
+        self.last_write = time.monotonic()
+        self.last_read = time.monotonic()
+
+    @property
+    def last_activity(self):
+        return max(self.last_write, self.last_read)
+
+    # AsyncIO
+    def is_closing(self):
+        """Return True if the transport is closing or closed."""
+        return self._closed
+
+    def close(self):
+        """Close the transport.
+
+        Buffered data will be flushed asynchronously.  No more data
+        will be received.  After all buffered data is flushed, the
+        protocol's connection_lost() method will (eventually) be
+        called with None as its argument.
+        """
+        if not self._closed:
+            log.info('%s: Closing transport', self)
+            self._closed = True
+            self._read = False
+            if not self._write_buffer:
+                self._close()
+
+    def set_protocol(self, protocol):
+        """Set a new protocol."""
+        self._protocol = protocol
+        log.debug('%s: Set protocol %s', self, protocol)
+
+    def get_protocol(self):
+        """Return the current protocol."""
+        return self._protocol
+
+    """Interface for read-only transports."""
+
+    def is_reading(self):
+        """Return True if the transport is receiving."""
+        return self._read
+
+    def pause_reading(self):
+        """Pause the receiving end.
+
+        No data will be passed to the protocol's data_received()
+        method until resume_reading() is called.
+        """
+        self._read = False
+        log.debug('%s: Paused reading', self)
+
+    def resume_reading(self):
+        """Resume the receiving end.
+
+        Data received will once again be passed to the protocol's
+        data_received() method.
+        """
+        if not self._read:
+            self._net.call_soon(self._read_from_sock)
+        self._read = True
+        log.debug('%s: Resumed reading', self)
+
+    async def _read_from_sock(self):
+        while self._read and not self._closed:
+            await self._net.wait_read(self._sock)
+            recvd_data, err = self._sock_recv()
+            if err:
+                return self.abort(error=err)
+            log.debug('%s: received %d bytes', self, len(recvd_data))
+            self.last_read = time.monotonic()
+            if self._protocol and self._protocol._sensors:
+                self._protocol._sensors.bytes_received.record(len(recvd_data))
+            try:
+                self._protocol.data_received(recvd_data)
+            except Errors.KafkaProtocolError as e:
+                return self.abort(error=e)
+
+    def _sock_recv(self):
+        recvd = []
+        err = None
+        while True:
+            try:
+                data = self._sock.recv(4096)
+                # We expect socket.recv to raise an exception if there are no
+                # bytes available to read from the socket in non-blocking mode.
+                # but if the socket is disconnected, we will get empty data
+                # without an exception raised
+                if not data:
+                    log.error('%s: socket disconnected', self)
+                    err = Errors.KafkaConnectionError('socket disconnected')
+                    break
+                else:
+                    recvd.append(data)
+
+            except (BlockingIOError, InterruptedError):
+                break
+            except BaseException as e:
+                log.exception('%s: Error receiving network data'
+                              ' closing socket', self)
+                err = Errors.KafkaConnectionError(e)
+                break
+
+        recvd_data = b''.join(recvd)
+        return recvd_data, err
+
+    """Interface for write-only transports."""
+
+    def set_write_buffer_limits(self, high=None, low=None):
+        """Set the high- and low-water limits for write flow control.
+
+        These two values control when to call the protocol's
+        pause_writing() and resume_writing() methods.  If specified,
+        the low-water limit must be less than or equal to the
+        high-water limit.  Neither value can be negative.
+
+        The defaults are implementation-specific.  If only the
+        high-water limit is given, the low-water limit defaults to an
+        implementation-specific value less than or equal to the
+        high-water limit.  Setting high to zero forces low to zero as
+        well, and causes pause_writing() to be called whenever the
+        buffer becomes non-empty.  Setting low to zero causes
+        resume_writing() to be called only once the buffer is empty.
+        Use of zero for either limit is generally sub-optimal as it
+        reduces opportunities for doing I/O and computation
+        concurrently.
+        """
+        raise NotImplementedError
+
+    def get_write_buffer_size(self):
+        """Return the current size of the write buffer."""
+        raise NotImplementedError
+
+    def get_write_buffer_limits(self):
+        """Get the high and low watermarks for write flow control.
+        Return a tuple (low, high) where low and high are
+        positive number of bytes."""
+        raise NotImplementedError
+
+    def write(self, data):
+        """Write some data bytes to the transport.
+
+        This does not block; it buffers the data and arranges for it
+        to be sent out asynchronously.
+        """
+        if not self._write or self._closed:
+            raise RuntimeError('Transport closed for writes')
+        if not data:
+            raise ValueError('Cant write empty data')
+        self._write_buffer.append(data)
+        if not self._writing:
+            self._writing = True
+            self._net.call_soon(self._write_to_sock)
+
+    def writelines(self, list_of_data):
+        """Write a list (or any iterable) of data bytes to the transport."""
+        if not self._write or self._closed:
+            raise RuntimeError('Transport closed for writes')
+        self._write_buffer.extend(list_of_data)
+        if not self._writing:
+            self._writing = True
+            self._net.call_soon(self._write_to_sock)
+
+    async def _write_to_sock(self):
+        try:
+            while self._write and not self._closed and self._write_buffer:
+                await self._net.wait_write(self._sock)
+                total_bytes, err = self._sock_send()
+                if err:
+                    return self.abort(error=err)
+                log.debug('%s: sent %d bytes', self, total_bytes)
+                self.last_write = time.monotonic()
+                if self._protocol and self._protocol._sensors:
+                    self._protocol._sensors.bytes_sent.record(total_bytes)
+        finally:
+            self._writing = False
+        if not self._write:
+            self._sock.shutdown(socket.SHUT_WR)
+
+    def _sock_send(self):
+        total_bytes = 0
+        while self._write_buffer:
+            next_chunk = self._write_buffer.popleft()
+            # Wrap in memoryview so partial-send slicing is O(1) instead of
+            # copying the unsent tail on every BlockingIOError / short write.
+            if not isinstance(next_chunk, memoryview):
+                next_chunk = memoryview(next_chunk)
+            while next_chunk:
+                try:
+                    sent_bytes = self._sock.send(next_chunk)
+                    total_bytes += sent_bytes
+                    next_chunk = next_chunk[sent_bytes:]
+                except (BlockingIOError, InterruptedError):
+                    self._write_buffer.appendleft(next_chunk)
+                    return total_bytes, None
+                except BaseException as e:
+                    log.exception("%s: Error sending request data: %s", self, e)
+                    return total_bytes, Errors.KafkaConnectionError(e)
+        return total_bytes, None
+
+    def write_eof(self):
+        """Close the write end after flushing buffered data.
+
+        (This is like typing ^D into a UNIX program reading from stdin.)
+
+        Data may still be received.
+        """
+        log.debug('%s: write_eof', self)
+        self._write = False
+        if not self._write_buffer:
+            self._sock.shutdown(socket.SHUT_WR)
+
+    def can_write_eof(self):
+        """Return True if this transport supports write_eof(), False if not."""
+        return True
+
+    def abort(self, error=None):
+        """Close the transport immediately.
+
+        Buffered data will be lost.  No more data will be received.
+        The protocol's connection_lost() method will (eventually) be
+        called with None as its argument.
+        """
+        if not self._closed:
+            log.error('%s: Abort (%s)', self, error)
+            self._closed = True
+            self._write_buffer.clear()
+            self._read = self._write = False
+            self._close(error)
+
+    def _close(self, error=None):
+        # idempotent; no lock
+        sock = self._sock
+        self._sock = None
+        if sock is not None:
+            try:
+                self._net.unregister_event(sock, selectors.EVENT_READ | selectors.EVENT_WRITE)
+            except (KeyError, ValueError):
+                pass
+            sock.close()
+        proto = self._protocol
+        self._protocol = None
+        if proto is not None:
+            proto.connection_lost(error)
+
+  # Twisted
+    def abortConnection(self):
+        """Close the connection abruptly."""
+        return self.abort()
+
+    def getHost(self):
+        """Similar to getPeer, but returns an address describing this side of the connection.
+
+        Returns IPv4Address or IPv6Address.
+        """
+        return self._sock.getsockname()
+
+    def getPeer(self):
+        """Get the remote address of this connection.
+
+        Treat this method with caution. It is the unfortunate result of the CGI and Jabber standards,
+        but should not be considered reliable for the usual host of reasons;
+        port forwarding, proxying, firewalls, IP masquerading, etc.
+
+        Returns IPv4Address or IPv6Address.
+        """
+        return self._sock.getpeername()
+
+    def getTcpKeepAlive(self):
+        """Return if SO_KEEPALIVE is enabled."""
+        return self._sock.getsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE)
+
+    def getTcpNoDelay(self):
+        """Return if TCP_NODELAY is enabled."""
+        return self._sock.getsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY)
+
+    def loseWriteConnection(self):
+        """Half-close the write side of a TCP connection."""
+        return self.write_eof()
+
+    def setTcpKeepAlive(self, enabled):
+        """Enable/disable SO_KEEPALIVE."""
+        return self._sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, enabled)
+
+    def setTcpNoDelay(self, enabled):
+        """Enable/disable TCP_NODELAY."""
+        return self._sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, enabled)
+
+    def loseConnection(self):
+        """Close my connection, after writing all pending data.
+
+        Note that if there is a registered producer on a transport it will not be closed until the producer has been unregistered.
+        """
+        return self.close()
+
+    #def write(self, data):
+    #    """Write some data to the physical connection, in sequence, in a non-blocking fashion.
+    #
+    #    If possible, make sure that it is all written. No data will ever be lost,
+    #    although (obviously) the connection may be closed before it all gets through.
+    #    """
+    #    pass
+
+    def writeSequence(self, data):
+        """Write an iterable of byte strings to the physical connection.
+
+        If possible, make sure that all of the data is written to the socket at once,
+        without first copying it all into a single byte string.
+        """
+        return self.writelines(data)
+
+    async def handshake(self):
+        pass
+
+    def host_port(self):
+        if self._sock is None:
+            return 'none'
+        try:
+            host, port = self._sock.getpeername()[0:2]
+        except (OSError, ValueError):
+            return 'none'
+        try:
+            local_port = self._sock.getsockname()[1]
+        except (OSError, ValueError):
+            return f'{host}:{port}'
+        return f'{host}:{port}<-{local_port}'
+
+    def __str__(self):
+        state = ' (closed)' if self._closed else ''
+        return f"<{self.__class__.__name__} [{self.host_port()}]{state}>"
+
+
+class KafkaSSLTransport(KafkaTCPTransport):
+    def __init__(self, net, sock, ssl_context, host=None, ssl_check_hostname=False):
+        self._ssl_context = ssl_context
+        server_hostname = host if ssl_check_hostname else None
+        sock = ssl_context.wrap_socket(
+            sock, server_hostname=server_hostname, do_handshake_on_connect=False)
+        super().__init__(net, sock, host=host)
+
+    async def handshake(self):
+        while True:
+            try:
+                self._sock.do_handshake()
+                return
+            except ssl.SSLWantReadError:
+                await self._net.wait_read(self._sock)
+            except ssl.SSLWantWriteError:
+                await self._net.wait_write(self._sock)
+
+    def _sock_recv(self):
+        recvd = []
+        err = None
+        while True:
+            try:
+                data = self._sock.recv(4096)
+                if not data:
+                    log.error('%s: socket disconnected', self)
+                    err = Errors.KafkaConnectionError('socket disconnected')
+                    break
+                else:
+                    recvd.append(data)
+            except (BlockingIOError, InterruptedError,
+                    ssl.SSLWantReadError, ssl.SSLWantWriteError):
+                break
+            except BaseException as e:
+                log.exception('%s: Error receiving network data'
+                              ' closing socket', self)
+                err = Errors.KafkaConnectionError(e)
+                break
+        recvd_data = b''.join(recvd)
+        return recvd_data, err
+
+    def _sock_send(self):
+        total_bytes = 0
+        err = None
+        while self._write_buffer:
+            next_chunk = self._write_buffer.popleft()
+            while next_chunk:
+                try:
+                    sent_bytes = self._sock.send(next_chunk)
+                    total_bytes += sent_bytes
+                    next_chunk = next_chunk[sent_bytes:]
+                except (BlockingIOError, InterruptedError,
+                        ssl.SSLWantReadError, ssl.SSLWantWriteError):
+                    self._write_buffer.appendleft(next_chunk)
+                    return total_bytes, err
+                except BaseException as e:
+                    log.exception("%s: Error sending request data: %s", self, e)
+                    err = Errors.KafkaConnectionError(e)
+                    return total_bytes, err
+        return total_bytes, err