fix async DNS

breakwa11 · breakwa11 · commit 8826f515b635 · 2017-06-03T21:47:17.000+08:00
diff --git a/shadowsocks/common.py b/shadowsocks/common.py
@@ -352,7 +352,7 @@ def __init__(self, params):
     def resolve(self, dns_resolver, remote_addr, call_back):
         if remote_addr in UDPAsyncDNSHandler.dns_cache:
             if call_back:
-                call_back(remote_addr, None, UDPAsyncDNSHandler.dns_cache[remote_addr], True, *self.params)
+                call_back(remote_addr, UDPAsyncDNSHandler.dns_cache[remote_addr], self.params)
         else:
             self.call_back = call_back
             self.remote_addr = remote_addr
@@ -367,7 +367,7 @@ def _handle_dns_resolved(self, result, error):
             ip = result[1]
             if ip:
                 if self.call_back:
-                    self.call_back(self.remote_addr, None, ip, True, *self.params)
+                    self.call_back(self.remote_addr, ip, self.params)
                     return
         logging.warning("can't resolve %s" % (self.remote_addr,))
 
diff --git a/shadowsocks/tcprelay.py b/shadowsocks/tcprelay.py
@@ -347,9 +347,9 @@ def _write_to_sock(self, data, sock):
                             handler = common.UDPAsyncDNSHandler(data[header_length:])
                             handler.resolve(self._dns_resolver, (dest_addr, dest_port), self._handle_server_dns_resolved)
                         else:
-                            return self._handle_server_dns_resolved(data[header_length:], (dest_addr, dest_port), dest_addr)
+                            return self._handle_server_dns_resolved((dest_addr, dest_port), dest_addr, data[header_length:])
                     else:
-                        return self._handle_server_dns_resolved(data[header_length:], (dest_addr, dest_port), dest_addr)
+                        return self._handle_server_dns_resolved((dest_addr, dest_port), dest_addr, data[header_length:])
 
             except Exception as e:
                 #trace = traceback.format_exc()
@@ -412,7 +412,7 @@ def _write_to_sock(self, data, sock):
                 logging.error('write_all_to_sock:unknown socket from %s:%d' % (self._client_address[0], self._client_address[1]))
         return True
 
-    def _handle_server_dns_resolved(self, data, remote_addr, server_addr):
+    def _handle_server_dns_resolved(self, remote_addr, server_addr, data):
         try:
             addrs = socket.getaddrinfo(server_addr, remote_addr[1], 0, socket.SOCK_DGRAM, socket.SOL_UDP)
             if not addrs: # drop
diff --git a/shadowsocks/udprelay.py b/shadowsocks/udprelay.py
@@ -388,17 +388,17 @@ def _handle_server(self):
                 handler = common.UDPAsyncDNSHandler((data, r_addr, uid, header_length))
                 handler.resolve(self._dns_resolver, (server_addr, server_port), self._handle_server_dns_resolved)
             else:
-                self._handle_server_dns_resolved((server_addr, server_port), None, server_addr, False, data, r_addr, uid, header_length)
+                self._handle_server_dns_resolved((server_addr, server_port), server_addr, (data, r_addr, uid, header_length))
         else:
-            self._handle_server_dns_resolved((server_addr, server_port), None, server_addr, False, data, r_addr, uid, header_length)
+            self._handle_server_dns_resolved((server_addr, server_port), server_addr, (data, r_addr, uid, header_length))
 
-    def _handle_server_dns_resolved(self, remote_addr, addrs, server_addr, dns_resolved, data, r_addr, uid, header_length):
+    def _handle_server_dns_resolved(self, remote_addr, server_addr, params):
+        data, r_addr, uid, header_length = params
         user_id = self._listen_port
         try:
             server_port = remote_addr[1]
-            if addrs is None:
-                addrs = socket.getaddrinfo(server_addr, server_port, 0,
-                                            socket.SOCK_DGRAM, socket.SOL_UDP)
+            addrs = socket.getaddrinfo(server_addr, server_port, 0,
+                                        socket.SOCK_DGRAM, socket.SOL_UDP)
             if not addrs: # drop
                 return
             af, socktype, proto, canonname, sa = addrs[0]
