From 469d9f7bfaa93a07caedd5cd3f981dec34fff3ba Mon Sep 17 00:00:00 2001
From: breakwa11 <mmgac001@gmail.com>
Date: Tue, 14 Jul 2015 19:02:52 +0800
Subject: [PATCH] merge master

---
 shadowsocks/udprelay.py         | 53 +++++++++++++++++----------------
 tests/jenkins.sh                |  2 +-
 utils/fail2ban/shadowsocks.conf |  5 ++++
 3 files changed, 33 insertions(+), 27 deletions(-)
 create mode 100644 utils/fail2ban/shadowsocks.conf

diff --git a/shadowsocks/udprelay.py b/shadowsocks/udprelay.py
index 238bf59..1e142bf 100644
--- a/shadowsocks/udprelay.py
+++ b/shadowsocks/udprelay.py
@@ -76,8 +76,9 @@ from shadowsocks.common import pre_parse_header, parse_header, pack_addr
 BUF_SIZE = 65536
 
 
-def client_key(a, b, c, d):
-    return '%s:%s:%s:%s' % (a, b, c, d)
+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)
 
 
 class UDPRelay(object):
@@ -102,6 +103,7 @@ class UDPRelay(object):
                                          close_callback=self._close_client)
         self._client_fd_to_server_addr = \
             lru_cache.LRUCache(timeout=config['timeout'])
+        self._dns_cache = lru_cache.LRUCache(timeout=300)
         self._eventloop = None
         self._closed = False
         self._last_time = time.time()
@@ -172,37 +174,36 @@ class UDPRelay(object):
 
         if self._is_local:
             server_addr, server_port = self._get_a_server()
-            key = client_key(r_addr[0], r_addr[1], dest_addr, dest_port)
         else:
             server_addr, server_port = dest_addr, dest_port
-            addrs = socket.getaddrinfo(dest_addr, dest_port, 0, socket.SOCK_DGRAM, socket.SOL_UDP)
-            if addrs:
-                af, socktype, proto, canonname, sa = addrs[0]
-                key = client_key(r_addr[0], r_addr[1], af, 0)
-            else:
-                key = None
 
-        client = self._cache.get(key, None)
-        if not client:
-            # TODO async getaddrinfo
-            #logging.info('UDP handle_server %s:%d from %s:%d' % (common.to_str(server_addr), server_port, self._listen_addr, self._listen_port))
+        addrs = self._dns_cache.get(server_addr, None)
+        if addrs is None:
             addrs = socket.getaddrinfo(server_addr, server_port, 0,
                                        socket.SOCK_DGRAM, socket.SOL_UDP)
-            if addrs:
-                af, socktype, proto, canonname, sa = addrs[0]
-                if self._forbidden_iplist:
-                    if common.to_str(sa[0]) in self._forbidden_iplist:
-                        logging.debug('IP %s is in forbidden list, drop' %
-                                      common.to_str(sa[0]))
-                        # drop
-                        return
-                client = socket.socket(af, socktype, proto)
-                client.setblocking(False)
-                self._cache[key] = client
-                self._client_fd_to_server_addr[client.fileno()] = r_addr
-            else:
+            if not addrs:
                 # drop
                 return
+            else:
+                self._dns_cache[server_addr] = addrs
+
+        af, socktype, proto, canonname, sa = addrs[0]
+        key = client_key(r_addr, af)
+        logging.debug(key)
+        client = self._cache.get(key, None)
+        if not client:
+            # TODO async getaddrinfo
+            if self._forbidden_iplist:
+                if common.to_str(sa[0]) in self._forbidden_iplist:
+                    logging.debug('IP %s is in forbidden list, drop' %
+                                  common.to_str(sa[0]))
+                    # drop
+                    return
+            client = socket.socket(af, socktype, proto)
+            client.setblocking(False)
+            self._cache[key] = client
+            self._client_fd_to_server_addr[client.fileno()] = r_addr
+
             self._sockets.add(client.fileno())
             self._eventloop.add(client, eventloop.POLL_IN)
 
diff --git a/tests/jenkins.sh b/tests/jenkins.sh
index 71d5b1c..ea5c163 100755
--- a/tests/jenkins.sh
+++ b/tests/jenkins.sh
@@ -69,7 +69,7 @@ if [ -f /proc/sys/net/ipv4/tcp_fastopen ] ; then
 fi
 
 run_test tests/test_large_file.sh
-
+run_test tests/test_udp_src.sh
 run_test tests/test_command.sh
 
 coverage combine && coverage report --include=shadowsocks/*
diff --git a/utils/fail2ban/shadowsocks.conf b/utils/fail2ban/shadowsocks.conf
new file mode 100644
index 0000000..9b1c7ec
--- /dev/null
+++ b/utils/fail2ban/shadowsocks.conf
@@ -0,0 +1,5 @@
+[Definition]
+
+_daemon = shadowsocks
+
+failregex = ^\s+ERROR\s+can not parse header when handling connection from <HOST>:\d+$