add async dns

shadowsocks/asyncdns.py

@@ -244,9 +244,7 @@ class DNSResolver(object):
         self._hostname_to_cb = {}
         self._cb_to_hostname = {}
         self._cache = lru_cache.LRUCache(timeout=300)
-        self._sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM,
-                                   socket.SOL_UDP)
-        self._sock.setblocking(False)
+        self._sock = None
         self._parse_config()
 
     def _parse_config(self):
@@ -272,7 +270,12 @@ class DNSResolver(object):
         self._dns_server = ('8.8.8.8', 53)
 
     def add_to_loop(self, loop):
+        if self._loop:
+            raise Exception('already add to loop')
         self._loop = loop
+        self._sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM,
+                                   socket.SOL_UDP)
+        self._sock.setblocking(False)
         loop.add(self._sock, eventloop.POLL_IN)
         loop.add_handler(self.handle_events)
 
@@ -362,6 +365,9 @@ class DNSResolver(object):
             else:
                 arr.append(callback)
 
+    def close(self):
+        self._sock.close()
+
 
 def test():
     logging.getLogger('').handlers = []

shadowsocks/local.py

@@ -29,6 +29,7 @@ import encrypt
 import eventloop
 import tcprelay
 import udprelay
+import asyncdns
 
 
 def main():
@@ -50,14 +51,18 @@ def main():
         logging.info("starting local at %s:%d" %
                      (config['local_address'], config['local_port']))
 
-        tcp_server = tcprelay.TCPRelay(config, True)
-        udp_server = udprelay.UDPRelay(config, True)
+        dns_resolver = asyncdns.DNSResolver()
+        tcp_server = tcprelay.TCPRelay(config, dns_resolver, True)
+        udp_server = udprelay.UDPRelay(config, dns_resolver, True)
         loop = eventloop.EventLoop()
+        dns_resolver.add_to_loop(loop)
         tcp_server.add_to_loop(loop)
         udp_server.add_to_loop(loop)
         loop.run()
     except (KeyboardInterrupt, IOError, OSError) as e:
         logging.error(e)
+        import traceback
+        traceback.print_exc()
         os._exit(0)
 
 if __name__ == '__main__':

shadowsocks/server.py

@@ -29,6 +29,7 @@ import encrypt
 import eventloop
 import tcprelay
 import udprelay
+import asyncdns
 
 
 def main():
@@ -50,18 +51,20 @@ def main():
     encrypt.init_table(config['password'], config['method'])
     tcp_servers = []
     udp_servers = []
+    dns_resolver = asyncdns.DNSResolver()
     for port, password in config['port_password'].items():
         a_config = config.copy()
         a_config['server_port'] = int(port)
         a_config['password'] = password
         logging.info("starting server at %s:%d" %
                      (a_config['server'], int(port)))
-        tcp_servers.append(tcprelay.TCPRelay(a_config, False))
-        udp_servers.append(udprelay.UDPRelay(a_config, False))
+        tcp_servers.append(tcprelay.TCPRelay(a_config, dns_resolver, False))
+        udp_servers.append(udprelay.UDPRelay(a_config, dns_resolver, False))
 
     def run_server():
         try:
             loop = eventloop.EventLoop()
+            dns_resolver.add_to_loop(loop)
             for tcp_server in tcp_servers:
                 tcp_server.add_to_loop(loop)
             for udp_server in udp_servers:
@@ -100,6 +103,7 @@ def main():
                     a_tcp_server.close()
                 for a_udp_server in udp_servers:
                     a_udp_server.close()
+                dns_resolver.close()
 
                 for child in children:
                     os.waitpid(child, 0)

shadowsocks/tcprelay.py

@@ -426,9 +426,10 @@ class TCPRelayHandler(object):
 
 
 class TCPRelay(object):
-    def __init__(self, config, is_local):
+    def __init__(self, config, dns_resolver, is_local):
         self._config = config
         self._is_local = is_local
+        self._dns_resolver = dns_resolver
         self._closed = False
         self._eventloop = None
         self._fd_to_handlers = {}
@@ -466,6 +467,8 @@ class TCPRelay(object):
         self._server_socket = server_socket
 
     def add_to_loop(self, loop):
+        if self._eventloop:
+            raise Exception('already add to loop')
         if self._closed:
             raise Exception('already closed')
         self._eventloop = loop

shadowsocks/udprelay.py

@@ -85,7 +85,7 @@ def client_key(a, b, c, d):
 
 
 class UDPRelay(object):
-    def __init__(self, config, is_local):
+    def __init__(self, config, dns_resolver, is_local):
         if is_local:
             self._listen_addr = config['local_address']
             self._listen_port = config['local_port']
@@ -96,6 +96,7 @@ class UDPRelay(object):
             self._listen_port = config['server_port']
             self._remote_addr = None
             self._remote_port = None
+        self._dns_resolver = dns_resolver
         self._password = config['password']
         self._method = config['method']
         self._timeout = config['timeout']
@@ -220,6 +221,8 @@ class UDPRelay(object):
             pass
 
     def add_to_loop(self, loop):
+        if self._eventloop:
+            raise Exception('already add to loop')
         if self._closed:
             raise Exception('already closed')
         self._eventloop = loop