From 2cf5f1982e3c007fcf3e21eb775372a8f55c86e8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=A0=B4=E5=A8=83=E9=85=B1?= Date: Thu, 1 Jun 2017 00:11:15 +0800 Subject: [PATCH] add DNS cache --- shadowsocks/common.py | 31 +++++++++++++++++++++++++++++++ shadowsocks/tcprelay.py | 31 ++++++------------------------- shadowsocks/udprelay.py | 33 ++++++--------------------------- 3 files changed, 43 insertions(+), 52 deletions(-) diff --git a/shadowsocks/common.py b/shadowsocks/common.py index f40b5fa..f5666d3 100644 --- a/shadowsocks/common.py +++ b/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) diff --git a/shadowsocks/tcprelay.py b/shadowsocks/tcprelay.py index f401276..328b48d 100644 --- a/shadowsocks/tcprelay.py +++ b/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,8 +342,12 @@ class TCPRelayHandler(object): continue connecttype, addrtype, dest_addr, dest_port, header_length = header_result if (addrtype & 7) == 3: - handler = UDPAsyncDNSHandler(data[header_length:]) - handler.resolve(self._dns_resolver, (dest_addr, dest_port), self._handle_server_dns_resolved) + 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) diff --git a/shadowsocks/udprelay.py b/shadowsocks/udprelay.py index c69b6ed..23eb6d9 100644 --- a/shadowsocks/udprelay.py +++ b/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,8 +383,12 @@ class UDPRelay(object): server_addr, server_port = dest_addr, dest_port if (addrtype & 7) == 3: - handler = UDPAsyncDNSHandler((data, r_addr, uid, header_length)) - handler.resolve(self._dns_resolver, (server_addr, server_port), self._handle_server_dns_resolved) + 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) @@ -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()