add DNS cache

shadowsocks/common.py

@@ -24,6 +24,8 @@ import logging
 import binascii
 import re
 
+from shadowsocks import lru_cache
+
 def compat_ord(s):
     if type(s) == int:
         return s
@@ -340,6 +342,35 @@ class PortRange(object):
     def __ne__(self, other):
         return self.range_str != other.range_str
 
+class UDPAsyncDNSHandler(object):
+    dns_cache = lru_cache.LRUCache(timeout=1800)
+    def __init__(self, params):
+        self.params = params
+        self.remote_addr = None
+        self.call_back = None
+
+    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)
+        else:
+            self.call_back = call_back
+            self.remote_addr = remote_addr
+            dns_resolver.resolve(remote_addr[0], self._handle_dns_resolved)
+            UDPAsyncDNSHandler.dns_cache.sweep()
+
+    def _handle_dns_resolved(self, result, error):
+        if error:
+            logging.error("%s when resolve DNS" % (error,)) #drop
+            return
+        if result:
+            ip = result[1]
+            if ip:
+                if self.call_back:
+                    self.call_back(self.remote_addr, None, ip, True, *self.params)
+                    return
+        logging.warning("can't resolve %s" % (self.remote_addr,))
+
 def test_inet_conv():
     ipv4 = b'8.8.4.4'
     b = inet_pton(socket.AF_INET, ipv4)

shadowsocks/tcprelay.py

@@ -123,29 +123,6 @@ class SpeedTester(object):
             return self.sum_len >= self.max_speed
         return False
 
-class UDPAsyncDNSHandler(object):
-    def __init__(self, params):
-        self.params = params
-        self.remote_addr = None
-        self.call_back = None
-
-    def resolve(self, dns_resolver, remote_addr, call_back):
-        self.call_back = call_back
-        self.remote_addr = remote_addr
-        dns_resolver.resolve(remote_addr[0], self._handle_dns_resolved)
-
-    def _handle_dns_resolved(self, result, error):
-        if error:
-            logging.error("%s when resolve DNS" % (error,)) #drop
-            return
-        if result:
-            ip = result[1]
-            if ip:
-                if self.call_back:
-                    self.call_back(self.params, self.remote_addr, ip)
-                    return
-        logging.warning("can't resolve %s" % (self.remote_addr,))
-
 class TCPRelayHandler(object):
     def __init__(self, server, fd_to_handlers, loop, local_sock, config,
                  dns_resolver, is_local):
@@ -365,10 +342,14 @@ class TCPRelayHandler(object):
                         continue
                     connecttype, addrtype, dest_addr, dest_port, header_length = header_result
                     if (addrtype & 7) == 3:
-                        handler = UDPAsyncDNSHandler(data[header_length:])
+                        af = common.is_ip(dest_addr)
+                        if af == False:
+                            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)
+                    else:
+                        return self._handle_server_dns_resolved(data[header_length:], (dest_addr, dest_port), dest_addr)
 
             except Exception as e:
                 #trace = traceback.format_exc()

shadowsocks/udprelay.py

@@ -123,29 +123,6 @@ RSP_STATE_ERROR = b"\x03"
 RSP_STATE_DISCONNECT = b"\x04"
 RSP_STATE_REDIRECT = b"\x05"
 
-class UDPAsyncDNSHandler(object):
-    def __init__(self, params):
-        self.params = params
-        self.remote_addr = None
-        self.call_back = None
-
-    def resolve(self, dns_resolver, remote_addr, call_back):
-        self.call_back = call_back
-        self.remote_addr = remote_addr
-        dns_resolver.resolve(remote_addr[0], self._handle_dns_resolved)
-
-    def _handle_dns_resolved(self, result, error):
-        if error:
-            logging.error("%s when resolve DNS" % (error,)) #drop
-            return
-        if result:
-            ip = result[1]
-            if ip:
-                if self.call_back:
-                    self.call_back(self.remote_addr, None, ip, True, *self.params)
-                    return
-        logging.warning("can't resolve %s" % (self.remote_addr,))
-
 def client_key(source_addr, server_af):
     # notice this is server af, not dest af
     return '%s:%s:%d' % (source_addr[0], source_addr[1], server_af)
@@ -406,10 +383,14 @@ class UDPRelay(object):
             server_addr, server_port = dest_addr, dest_port
 
         if (addrtype & 7) == 3:
-            handler = UDPAsyncDNSHandler((data, r_addr, uid, header_length))
+            af = common.is_ip(server_addr)
+            if af == False:
+                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)
+        else:
+            self._handle_server_dns_resolved((server_addr, server_port), None, server_addr, False, 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):
         user_id = self._listen_port
@@ -645,7 +626,6 @@ class UDPRelay(object):
         if self._closed:
             self._cache.clear(0)
             self._cache_dns_client.clear(0)
-            #self._dns_cache.sweep()
             if self._eventloop:
                 self._eventloop.remove_periodic(self.handle_periodic)
                 self._eventloop.remove(self._server_socket)
@@ -657,7 +637,6 @@ class UDPRelay(object):
             before_sweep_size = len(self._sockets)
             self._cache.sweep()
             self._cache_dns_client.sweep()
-            #self._dns_cache.sweep()
             if before_sweep_size != len(self._sockets):
                 logging.debug('UDP port %5d sockets %d' % (self._listen_port, len(self._sockets)))
             self._sweep_timeout()