fix a mem leak

breakwa11 · breakwa11 · commit 47c6e1683353 · 2017-06-03T23:06:12.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, UDPAsyncDNSHandler.dns_cache[remote_addr], self.params)
+                call_back("", remote_addr, UDPAsyncDNSHandler.dns_cache[remote_addr], self.params)
         else:
             self.call_back = call_back
             self.remote_addr = remote_addr
@@ -362,14 +362,13 @@ def resolve(self, dns_resolver, remote_addr, call_back):
     def _handle_dns_resolved(self, result, error):
         if error:
             logging.error("%s when resolve DNS" % (error,)) #drop
-            return
+            return self.call_back(error, self.remote_addr, None, self.params)
         if result:
             ip = result[1]
             if ip:
-                if self.call_back:
-                    self.call_back(self.remote_addr, ip, self.params)
-                    return
+                return self.call_back("", self.remote_addr, ip, self.params)
         logging.warning("can't resolve %s" % (self.remote_addr,))
+        return self.call_back("fail to resolve", self.remote_addr, None, self.params)
 
 def test_inet_conv():
     ipv4 = b'8.8.4.4'
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((dest_addr, dest_port), dest_addr, data[header_length:])
+                            return self._handle_server_dns_resolved("", (dest_addr, dest_port), dest_addr, data[header_length:])
                     else:
-                        return self._handle_server_dns_resolved((dest_addr, dest_port), dest_addr, data[header_length:])
+                        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,10 @@ 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, remote_addr, server_addr, data):
+    def _handle_server_dns_resolved(self, error, remote_addr, server_addr, data):
+        if error:
+            self.destroy()
+            return
         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,11 +388,13 @@ 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), server_addr, (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), server_addr, (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, server_addr, params):
+    def _handle_server_dns_resolved(self, error, remote_addr, server_addr, params):
+        if error:
+            return
         data, r_addr, uid, header_length = params
         user_id = self._listen_port
         try:
