diff --git a/kafka/conn.py b/kafka/conn.py old mode 100644 new mode 100755 index 64445fab0..469e19c1e --- a/kafka/conn.py +++ b/kafka/conn.py @@ -326,6 +326,8 @@ def _dns_lookup(self): return True def _next_afi_sockaddr(self): + if self._socks5_proxy.use_remote_lookup(): + return (socket.AF_UNSPEC, (self.host, self.port)) if not self._gai: if not self._dns_lookup(): return @@ -366,6 +368,9 @@ def connect_blocking(self, timeout=float('inf')): def connect(self): """Attempt to connect and return ConnectionState""" + if self.config["socks5_proxy"] is not None and self._socks5_proxy is None: + self._socks5_proxy = Socks5Wrapper(self.config["socks5_proxy"], self.afi) + if self.state is ConnectionStates.DISCONNECTED and not self.blacked_out(): self.state = ConnectionStates.CONNECTING self.last_attempt = time.time() @@ -379,7 +384,6 @@ def connect(self): self._sock_afi, self._sock_addr = next_lookup try: if self.config["socks5_proxy"] is not None: - self._socks5_proxy = Socks5Wrapper(self.config["socks5_proxy"], self.afi) self._sock = self._socks5_proxy.socket(self._sock_afi, socket.SOCK_STREAM) else: self._sock = socket.socket(self._sock_afi, socket.SOCK_STREAM) @@ -862,7 +866,7 @@ def connection_delay(self): large number to handle slow/stalled connections. """ if self.disconnected() or self.connecting(): - if len(self._gai) > 0: + if len(self._gai) > 0 or (self._socks5_proxy is not None and self._socks5_proxy.use_remote_lookup()): return 0 else: time_waited = time.time() - self.last_attempt diff --git a/kafka/socks5_wrapper.py b/kafka/socks5_wrapper.py old mode 100644 new mode 100755 index 18bea7c8d..cdd600b7a --- a/kafka/socks5_wrapper.py +++ b/kafka/socks5_wrapper.py @@ -64,6 +64,9 @@ def dns_lookup(cls, host, port, afi=socket.AF_UNSPEC): log.warning("DNS lookup failed for proxy %s:%d, %r", host, port, ex) return [] + def use_remote_lookup(self): + return self._proxy_url.scheme == 'socks5h' + def socket(self, family, sock_type): """Open and record a socket. @@ -187,7 +190,10 @@ def connect_ex(self, addr): return errno.ECONNREFUSED if self._state == ProxyConnectionStates.REQUEST_SUBMIT: - if self._target_afi == socket.AF_INET: + 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: @@ -199,15 +205,27 @@ def connect_ex(self, addr): self._sock.close() return errno.ECONNREFUSED - self._buffer_out = struct.pack( - "!bbbb{}sh".format(addr_len), - 5, # version - 1, # command: connect - 0, # reserved - addr_type, # 1 for ipv4, 4 for ipv6 address - socket.inet_pton(self._target_afi, addr[0]), # either 4 or 16 bytes of actual address - addr[1], # port - ) + if self.use_remote_lookup(): + self._buffer_out = struct.pack( + "!bbbbb{}sh".format(addr_len), + 5, # version + 1, # command: connect + 0, # reserved + addr_type, # 1 for ipv4, 4 for ipv6 address, 3 for domain name + addr_len, + addr[0].encode('ascii'), # host + addr[1], # port + ) + else: + self._buffer_out = struct.pack( + "!bbbb{}sh".format(addr_len), + 5, # version + 1, # command: connect + 0, # reserved + addr_type, # 1 for ipv4, 4 for ipv6 address, 3 for domain name + socket.inet_pton(self._target_afi, addr[0]), # either 4 or 16 bytes of actual address + addr[1], # port + ) self._state = ProxyConnectionStates.REQUESTING if self._state == ProxyConnectionStates.REQUESTING: