diff --git a/shadowsocks/udprelay.py b/shadowsocks/udprelay.py index ed6937f..888b036 100644 --- a/shadowsocks/udprelay.py +++ b/shadowsocks/udprelay.py @@ -76,8 +76,9 @@ from shadowsocks.common import parse_header, pack_addr 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 @@ class UDPRelay(object): 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) diff --git a/tests/jenkins.sh b/tests/jenkins.sh index 5e4e61f..ea5c163 100755 --- a/tests/jenkins.sh +++ b/tests/jenkins.sh @@ -24,8 +24,6 @@ function run_test { return 0 } -pip install PySocks - python --version coverage erase mkdir tmp diff --git a/tests/test_udp_src.py b/tests/test_udp_src.py index 83c615d..840b8f6 100644 --- a/tests/test_udp_src.py +++ b/tests/test_udp_src.py @@ -4,6 +4,7 @@ import socket 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 @@ if __name__ == '__main__': 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 @@ if __name__ == '__main__': # 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() diff --git a/tests/test_udp_src.sh b/tests/test_udp_src.sh index 876f242..d356581 100755 --- a/tests/test_udp_src.sh +++ b/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