use_remote_lookup(proxy_url) classmethod; add debug statement for proxy init; consolidate buffer_out; fix port H encoding

Author: dpkp

kafka/conn.py

@@ -326,7 +326,7 @@ def _dns_lookup(self):
         return True
 
     def _next_afi_sockaddr(self):
-        if self._socks5_proxy and self._socks5_proxy.use_remote_lookup():
+        if self.config["socks5_proxy"] and Socks5Wrapper.use_remote_lookup(self.config["socks5_proxy"]):
             return (socket.AF_UNSPEC, (self.host, self.port))
 
         if not self._gai:
@@ -382,6 +382,7 @@ def connect(self):
                 self._sock_afi, self._sock_addr = next_lookup
                 try:
                     if self.config["socks5_proxy"] is not None:
+                        log.debug('%s: initializing Socks5 proxy at %s', self, self.config["socks5_proxy"])
                         self._socks5_proxy = Socks5Wrapper(self.config["socks5_proxy"], self.afi)
                         self._sock = self._socks5_proxy.socket(self._sock_afi, socket.SOCK_STREAM)
                     else:
@@ -867,7 +868,7 @@ def connection_delay(self):
         if self.disconnected() or self.connecting():
             if len(self._gai) > 0:
                 return 0
-            elif self._socks5_proxy and self._socks5_proxy.use_remote_lookup():
+            elif Socks5Wrapper.use_remote_lookup(self.config["socks5_proxy"]):
                 return 0
             else:
                 time_waited = time.time() - self.last_attempt

kafka/socks5_wrapper.py

@@ -64,7 +64,11 @@ 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):
+    @classmethod
+    def use_remote_lookup(cls, proxy_url):
+        return urlparse(proxy_url).scheme == 'socks5h'
+
+    def _use_remote_lookup(self):
         return self._proxy_url.scheme == 'socks5h'
 
     def socket(self, family, sock_type):
@@ -190,7 +194,7 @@ def connect_ex(self, addr):
                 return errno.ECONNREFUSED
 
         if self._state == ProxyConnectionStates.REQUEST_SUBMIT:
-            if self.use_remote_lookup():
+            if self._use_remote_lookup():
                 addr_type = 3
                 addr_len = len(addr[0])
             elif self._target_afi == socket.AF_INET:
@@ -205,27 +209,29 @@ def connect_ex(self, addr):
                 self._sock.close()
                 return errno.ECONNREFUSED
 
-            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
+            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'),  # host
-                    addr[1],  # port
+                    addr[0].encode('ascii'),
                 )
             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
+                # 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: