kafka-python 3.0.0__py3-none-any.whl

kafka/net/http_connect.py

@@ -0,0 +1,144 @@
+import base64
+import errno
+import logging
+import random
+import socket
+from urllib.parse import urlparse
+
+from kafka.errors import KafkaConnectionError
+from kafka.net.inet import KafkaNetSocket
+
+
+log = logging.getLogger(__name__)
+
+_WOULD_BLOCK = {errno.EWOULDBLOCK, errno.EAGAIN}
+_MAX_RESPONSE_SIZE = 65536
+
+
+class _States:
+    DISCONNECTED = '<disconnected>'
+    CONNECTING = '<connecting>'
+    SENDING = '<sending>'
+    READING = '<reading>'
+    COMPLETE = '<complete>'
+
+
+class HttpConnectProxy(KafkaNetSocket):
+    """Tunnels broker connections through an HTTP CONNECT proxy (RFC 7231 s4.3.6).
+
+    Registered for the ``http`` scheme -- pass ``proxy_url='http://host:port'``
+    to KafkaConsumer/KafkaProducer/KafkaAdminClient.
+
+    Basic proxy auth is supported via URL credentials: ``http://user:pass@host:8080``.
+    Broker hostnames are always forwarded unresolved so the proxy handles DNS.
+    """
+
+    SCHEMES = ('http',)
+
+    def __init__(self, proxy_url):
+        self._proxy_url = urlparse(proxy_url)
+        self._sock = None
+        self._state = _States.DISCONNECTED
+        self._send_buf = b''
+        self._recv_buf = b''
+        self._proxy_addr = self._get_proxy_addr()
+
+    def _get_proxy_addr(self):
+        addrs = self.dns_lookup(self._proxy_url.hostname, self._proxy_url.port, proxy=True)
+        if not addrs:
+            raise KafkaConnectionError('Unable to resolve proxy_url via dns')
+        return random.choice(addrs)
+
+    def dns_lookup(self, host, port, proxy=False):
+        if proxy:
+            return super().dns_lookup(host, port, raise_error=True)
+        # Always forward broker hostname unresolved; the proxy handles DNS
+        return [(socket.AF_UNSPEC, socket.SOCK_STREAM, socket.IPPROTO_TCP, '', (host, port))]
+
+    def socket(self, family=socket.AF_UNSPEC, sock_type=socket.SOCK_STREAM, proto=socket.IPPROTO_TCP):
+        self._target_afi = family
+        proxy_family, _, _, _, _ = self._proxy_addr
+        self._sock = socket.socket(proxy_family, sock_type, proto)
+        return self._sock
+
+    def connect_ex(self, sock, addr):
+        assert sock is self._sock
+
+        if self._state == _States.DISCONNECTED:
+            self._state = _States.CONNECTING
+
+        if self._state == _States.CONNECTING:
+            ret = self._do_connecting(addr)
+            if ret is not None:
+                return ret
+
+        if self._state == _States.SENDING:
+            ret = self._do_sending()
+            if ret is not None:
+                return ret
+
+        if self._state == _States.READING:
+            ret = self._do_reading()
+            if ret is not None:
+                return ret
+
+        if self._state == _States.COMPLETE:
+            return 0
+
+        return errno.ECONNREFUSED
+
+    def _do_connecting(self, addr):
+        _, _, _, _, proxy_sockaddr = self._proxy_addr
+        ret = self._sock.connect_ex(proxy_sockaddr)
+        if ret and ret != errno.EISCONN:
+            return ret
+        host, port = addr[0], addr[1]
+        headers = 'CONNECT {0}:{1} HTTP/1.1\r\nHost: {0}:{1}\r\n'.format(host, port)
+        if self._proxy_url.username and self._proxy_url.password:
+            credentials = base64.b64encode(
+                '{0}:{1}'.format(self._proxy_url.username, self._proxy_url.password).encode()
+            ).decode()
+            headers += 'Proxy-Authorization: Basic {}\r\n'.format(credentials)
+        self._send_buf = (headers + '\r\n').encode()
+        self._state = _States.SENDING
+        return None
+
+    def _do_sending(self):
+        while self._send_buf:
+            try:
+                sent = self._sock.send(self._send_buf)
+                if sent == 0:
+                    log.error('Proxy closed connection while sending CONNECT request')
+                    return errno.ECONNREFUSED
+                self._send_buf = self._send_buf[sent:]
+            except OSError as exc:
+                if exc.errno in _WOULD_BLOCK:
+                    return errno.EWOULDBLOCK
+                raise
+        self._state = _States.READING
+        return None
+
+    def _do_reading(self):
+        while b'\r\n\r\n' not in self._recv_buf:
+            try:
+                chunk = self._sock.recv(4096)
+                if not chunk:
+                    log.error('Proxy closed connection during CONNECT handshake')
+                    self._sock.close()
+                    return errno.ECONNREFUSED
+                self._recv_buf += chunk
+                if len(self._recv_buf) > _MAX_RESPONSE_SIZE:
+                    log.error('Proxy response exceeded %d bytes without end-of-headers', _MAX_RESPONSE_SIZE)
+                    self._sock.close()
+                    return errno.ECONNREFUSED
+            except OSError as exc:
+                if exc.errno in _WOULD_BLOCK:
+                    return errno.EWOULDBLOCK
+                raise
+        first_line = self._recv_buf.split(b'\r\n')[0]
+        if b' 200 ' in first_line or first_line.endswith(b' 200'):
+            self._state = _States.COMPLETE
+            return None
+        log.error('HTTP CONNECT to proxy failed: %r', first_line)
+        self._sock.close()
+        return errno.ECONNREFUSED

kafka/net/inet.py

@@ -0,0 +1,122 @@
+import errno
+import logging
+import socket
+import time
+from urllib.parse import urlparse
+
+import kafka.errors as Errors
+
+
+log = logging.getLogger(__name__)
+
+
+async def create_connection(net, host, port, socket_options=(), proxy_url=None, timeout_at=None):
+    """Connect to host:port; raises KafkaConnectionError on failure"""
+    socket_factory = KafkaNetSocket(proxy_url)
+    addrs = socket_factory.dns_lookup(host, port)
+    exceptions = [Errors.KafkaConnectionError('DNS Resolution failure')]
+    for res in addrs:
+        try:
+            log.debug('%s: Attempting to connect to %s (options: %s)', socket_factory, res, socket_options)
+            sock = await socket_factory.connect(net, res, socket_options, timeout_at=timeout_at)
+        except (socket.error, OSError) as e:
+            exceptions.append(Errors.KafkaConnectionError('unable to connect: %s' % (e,)))
+            continue
+        except Errors.KafkaTimeoutError:
+            raise Errors.KafkaConnectionError('Connection timed out')
+        except Errors.KafkaConnectionError as e:
+            exceptions.append(e)
+            continue
+        else:
+            return sock
+    raise exceptions[-1]
+
+
+class KafkaNetSocket:
+    # scheme => handling class
+    _registry = {}
+
+    @classmethod
+    def register_class(cls, klass):
+        for scheme in klass.SCHEMES:
+            cls._registry[scheme] = klass
+
+    def __init_subclass__(cls, **kw):
+        super().__init_subclass__(**kw)
+        KafkaNetSocket.register_class(cls)
+
+    def __new__(cls, proxy_url=None):
+        if proxy_url is None:
+            return super().__new__(cls)
+        try:
+            parsed = urlparse(proxy_url)
+        except Exception:
+            raise ValueError('Unable to parse proxy_url: %s' % (proxy_url,))
+        if not parsed.scheme:
+            raise ValueError('proxy_url requires scheme:// (%s)' % (proxy_url,))
+        try:
+            klass = KafkaNetSocket._registry[parsed.scheme]
+        except KeyError:
+            raise ValueError('Unsupported proxy url scheme: %s' % (parsed.scheme))
+        return super().__new__(klass)
+
+    def __init__(self, proxy_url=None):
+        pass
+
+    # simple sockets / no proxy
+    def dns_lookup(self, host, port, raise_error=False):
+        # XXX: all DNS functions in Python are blocking. If we really
+        # want to be non-blocking here, we need to use a 3rd-party
+        # library like python-adns, or move resolution onto its
+        # own thread. This will be subject to the default libc
+        # name resolution timeout (5s on most Linux boxes)
+        try:
+            return socket.getaddrinfo(host, port, socket.AF_UNSPEC, socket.SOCK_STREAM)
+        except socket.gaierror as ex:
+            err_str = "DNS lookup failed for %s:%d, %r" % (host, port, ex)
+            if not raise_error:
+                log.warning(err_str)
+                return []
+            raise Errors.KafkaConnectionError(err_str)
+
+    def socket(self, family=socket.AF_UNSPEC, sock_type=socket.SOCK_STREAM, proto=socket.IPPROTO_TCP):
+        return socket.socket(family, sock_type, proto)
+
+    async def connect(self, net, addrinfo, socket_options=(), timeout_at=None):
+        """Create non-blocking socket (with options) and connect to addrinfo tuple"""
+        family, sock_type, proto, _canonname, sockaddr = addrinfo
+        sock = self.socket(family, sock_type, proto)
+        sock.setblocking(False)
+        for option in socket_options:
+            sock.setsockopt(*option)
+        return await self.sock_connect(net, sock, sockaddr, timeout_at=timeout_at)
+
+    async def sock_connect(self, net, sock, sockaddr, timeout_at=None):
+        while timeout_at is None or time.monotonic() < timeout_at:
+            ret = None
+            try:
+                ret = self.connect_ex(sock, sockaddr)
+            except BlockingIOError:
+                ret = errno.EWOULDBLOCK
+            except socket.error as err:
+                ret = err.errno
+
+            # Connection succeeded
+            if not ret or ret == errno.EISCONN:
+                log.debug('Connected: %s', sock)
+                return sock
+
+            # Needs retry
+            # WSAEINVAL == 10022, but errno.WSAEINVAL is not available on non-win systems
+            elif ret in (errno.EINPROGRESS, errno.EALREADY, errno.EWOULDBLOCK, 10022):
+                await net.wait_write(sock, timeout_at=timeout_at)
+
+            # Connection failed
+            else:
+                errstr = errno.errorcode.get(ret, 'UNKNOWN')
+                raise Errors.KafkaConnectionError('{} {}'.format(ret, errstr))
+        else:
+            raise Errors.KafkaTimeoutError('Connection timed out')
+
+    def connect_ex(self, sock, sockaddr):
+        return sock.connect_ex(sockaddr)