From 71c3759c126ebbc7cfc1ea749e08913f4c116474 Mon Sep 17 00:00:00 2001 From: clowwindy Date: Sun, 8 Jun 2014 15:58:59 +0800 Subject: [PATCH] add async dns --- shadowsocks/asyncdns.py | 12 +++++++++--- shadowsocks/local.py | 9 +++++++-- shadowsocks/server.py | 8 ++++++-- shadowsocks/tcprelay.py | 5 ++++- shadowsocks/udprelay.py | 5 ++++- 5 files changed, 30 insertions(+), 9 deletions(-) diff --git a/shadowsocks/asyncdns.py b/shadowsocks/asyncdns.py index 7c52b16..2272ea2 100644 --- a/shadowsocks/asyncdns.py +++ b/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 = [] diff --git a/shadowsocks/local.py b/shadowsocks/local.py index 0b4b17d..0a00136 100755 --- a/shadowsocks/local.py +++ b/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__': diff --git a/shadowsocks/server.py b/shadowsocks/server.py index 13c2452..9145ce4 100755 --- a/shadowsocks/server.py +++ b/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) diff --git a/shadowsocks/tcprelay.py b/shadowsocks/tcprelay.py index c93b6fc..381a809 100644 --- a/shadowsocks/tcprelay.py +++ b/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 diff --git a/shadowsocks/udprelay.py b/shadowsocks/udprelay.py index 2aeb069..5218feb 100644 --- a/shadowsocks/udprelay.py +++ b/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