fix problem when UDP client requesting both IPv4 and IPv6

Author: clowwindy

shadowsocks/udprelay.py

@@ -76,8 +76,9 @@
 BUF_SIZE = 65536
 
 
-def client_key(source_addr, dest_addr):
-    return '%s:%s' % (source_addr[0], source_addr[1])
+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):
@@ -169,27 +170,29 @@ def _handle_server(self):
         else:
             server_addr, server_port = dest_addr, dest_port
 
-        key = client_key(r_addr, (dest_addr, dest_port))
+        addrs = socket.getaddrinfo(server_addr, server_port, 0,
+                                   socket.SOCK_DGRAM, socket.SOL_UDP)
+        if not addrs:
+            # drop
+            return
+
+        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
-            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:
-                # drop
-                return
+            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)
 

tests/jenkins.sh

@@ -24,8 +24,6 @@ function run_test {
     return 0
 }
 
-pip install PySocks
-
 python --version
 coverage erase
 mkdir tmp

tests/test_udp_src.py

@@ -4,6 +4,7 @@
 import socks
 
 if __name__ == '__main__':
+    # Test 1: same source port IPv4
     sock_out = socks.socksocket(socket.AF_INET, socket.SOCK_DGRAM,
                                 socket.SOL_UDP)
     sock_out.set_proxy(socks.SOCKS5, '127.0.0.1', 1081)
@@ -17,10 +18,10 @@
     sock_in1.bind(('127.0.0.1', 9001))
     sock_in2.bind(('127.0.0.1', 9002))
 
-    sock_out.sendto('data', ('127.0.0.1', 9001))
+    sock_out.sendto(b'data', ('127.0.0.1', 9001))
     result1 = sock_in1.recvfrom(8)
 
-    sock_out.sendto('data', ('127.0.0.1', 9002))
+    sock_out.sendto(b'data', ('127.0.0.1', 9002))
     result2 = sock_in2.recvfrom(8)
 
     sock_out.close()
@@ -29,3 +30,49 @@
 
     # make sure they're from the same source port
     assert result1 == result2
+
+    # Test 2: same source port IPv6
+    # try again from the same port but IPv6
+    sock_out = socks.socksocket(socket.AF_INET, socket.SOCK_DGRAM,
+                                socket.SOL_UDP)
+    sock_out.set_proxy(socks.SOCKS5, '127.0.0.1', 1081)
+    sock_out.bind(('127.0.0.1', 9000))
+
+    sock_in1 = socket.socket(socket.AF_INET6, socket.SOCK_DGRAM,
+                             socket.SOL_UDP)
+    sock_in2 = socket.socket(socket.AF_INET6, socket.SOCK_DGRAM,
+                             socket.SOL_UDP)
+
+    sock_in1.bind(('::1', 9001))
+    sock_in2.bind(('::1', 9002))
+
+    sock_out.sendto(b'data', ('::1', 9001))
+    result1 = sock_in1.recvfrom(8)
+
+    sock_out.sendto(b'data', ('::1', 9002))
+    result2 = sock_in2.recvfrom(8)
+
+    sock_out.close()
+    sock_in1.close()
+    sock_in2.close()
+
+    # make sure they're from the same source port
+    assert result1 == result2
+
+    # Test 3: different source ports IPv6
+    sock_out = socks.socksocket(socket.AF_INET, socket.SOCK_DGRAM,
+                                socket.SOL_UDP)
+    sock_out.set_proxy(socks.SOCKS5, '127.0.0.1', 1081)
+    sock_out.bind(('127.0.0.1', 9003))
+
+    sock_in1 = socket.socket(socket.AF_INET6, socket.SOCK_DGRAM,
+                             socket.SOL_UDP)
+    sock_in1.bind(('::1', 9001))
+    sock_out.sendto(b'data', ('::1', 9001))
+    result3 = sock_in1.recvfrom(8)
+
+    # make sure they're from different source ports
+    assert result1 != result3
+
+    sock_out.close()
+    sock_in1.close()

tests/test_udp_src.sh

@@ -4,10 +4,10 @@ PYTHON="coverage run -p -a"
 
 mkdir -p tmp
 
-$PYTHON shadowsocks/local.py -c tests/aes.json &
+$PYTHON shadowsocks/local.py -c tests/aes.json -v &
 LOCAL=$!
 
-$PYTHON shadowsocks/server.py -c tests/aes.json --forbidden-ip "" &
+$PYTHON shadowsocks/server.py -c tests/aes.json --forbidden-ip "" -v &
 SERVER=$!
 
 sleep 3