From 4312eb9e58f9438b76870101169be5e065695c6c Mon Sep 17 00:00:00 2001 From: clowwindy Date: Mon, 12 Jan 2015 23:02:59 +0800 Subject: [PATCH] add forbidden support for UDP and add tests --- .jenkins.sh | 1 + shadowsocks/tcprelay.py | 4 ++-- shadowsocks/udprelay.py | 11 +++++++++++ tests/test.py | 28 ++++++++++++++++++++-------- 4 files changed, 34 insertions(+), 10 deletions(-) diff --git a/.jenkins.sh b/.jenkins.sh index 7aa3257..bb58e04 100755 --- a/.jenkins.sh +++ b/.jenkins.sh @@ -45,6 +45,7 @@ run_test python tests/test.py --with-coverage -c tests/workers.json run_test python tests/test.py --with-coverage -s tests/ipv6.json -c tests/ipv6-client-side.json run_test python tests/test.py --with-coverage -b "-m rc4-md5 -k testrc4 -s 127.0.0.1 -p 8388 -q" -a "-m rc4-md5 -k testrc4 -s 127.0.0.1 -p 8388 -l 1081 -vv" run_test python tests/test.py --with-coverage -b "-m aes-256-cfb -k testrc4 -s 127.0.0.1 -p 8388 --workers 1" -a "-m aes-256-cfb -k testrc4 -s 127.0.0.1 -p 8388 -l 1081 -t 30 -qq -b 127.0.0.1" +run_test python tests/test.py --with-coverage --should-fail --url="http://127.0.0.1/" -b "-m aes-256-cfb -k testrc4 -s 127.0.0.1 -p 8388 --forbidden-ip=127.0.0.1,::1,8.8.8.8" -a "-m aes-256-cfb -k testrc4 -s 127.0.0.1 -p 8388 -l 1081 -t 30 -b 127.0.0.1" if [ -f /proc/sys/net/ipv4/tcp_fastopen ] ; then if [ 3 -eq `cat /proc/sys/net/ipv4/tcp_fastopen` ] ; then diff --git a/shadowsocks/tcprelay.py b/shadowsocks/tcprelay.py index 6afcad4..c148208 100644 --- a/shadowsocks/tcprelay.py +++ b/shadowsocks/tcprelay.py @@ -123,8 +123,8 @@ class TCPRelayHandler(object): self._downstream_status = WAIT_STATUS_INIT self._client_address = local_sock.getpeername()[:2] self._remote_address = None - if 'forbidden_ip' in self._config: - self._forbidden_iplist = self._config['forbidden_ip'] + if 'forbidden_ip' in config: + self._forbidden_iplist = config['forbidden_ip'] else: self._forbidden_iplist = None if is_local: diff --git a/shadowsocks/udprelay.py b/shadowsocks/udprelay.py index 2b8b12f..ccc0413 100644 --- a/shadowsocks/udprelay.py +++ b/shadowsocks/udprelay.py @@ -112,6 +112,11 @@ class UDPRelay(object): self._closed = False self._last_time = time.time() self._sockets = set() + print(config) + if 'forbidden_ip' in config: + self._forbidden_iplist = config['forbidden_ip'] + else: + self._forbidden_iplist = None addrs = socket.getaddrinfo(self._listen_addr, self._listen_port, 0, socket.SOCK_DGRAM, socket.SOL_UDP) @@ -178,6 +183,12 @@ class UDPRelay(object): 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.warn('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 diff --git a/tests/test.py b/tests/test.py index 0b63a18..f160366 100755 --- a/tests/test.py +++ b/tests/test.py @@ -40,6 +40,9 @@ parser.add_argument('-s', '--server-conf', type=str, default=None) parser.add_argument('-a', '--client-args', type=str, default=None) parser.add_argument('-b', '--server-args', type=str, default=None) parser.add_argument('--with-coverage', action='store_true', default=None) +parser.add_argument('--should-fail', action='store_true', default=None) +parser.add_argument('--url', type=str, default='http://www.example.com/') +parser.add_argument('--dns', type=str, default='8.8.8.8') config = parser.parse_args() @@ -87,6 +90,7 @@ try: for fd in r: line = fd.readline() + sys.stderr.write(line) if not line: if stage == 2 and fd == p3.stdout: stage = 3 @@ -94,7 +98,6 @@ try: stage = 5 if bytes != str: line = str(line, 'utf8') - sys.stdout.write(line) if line.find('starting local') >= 0: local_ready = True if line.find('starting server') >= 0: @@ -103,7 +106,7 @@ try: if stage == 1: time.sleep(2) - p3 = Popen(['curl', 'http://www.example.com/', '-v', '-L', + p3 = Popen(['curl', config.url, '-v', '-L', '--socks5-hostname', '127.0.0.1:1081', '-m', '15', '--connect-timeout', '10'], stdin=PIPE, stdout=PIPE, stderr=PIPE, close_fds=True) @@ -118,9 +121,13 @@ try: fdset.remove(p3.stdout) fdset.remove(p3.stderr) r = p3.wait() - if r != 0: - sys.exit(1) - p4 = Popen(['socksify', 'dig', '@8.8.8.8', 'www.google.com'], + if config.should_fail: + if r == 0: + sys.exit(1) + else: + if r != 0: + sys.exit(1) + p4 = Popen(['socksify', 'dig', '@%s' % config.dns, 'www.google.com'], stdin=PIPE, stdout=PIPE, stderr=PIPE, close_fds=True) if p4 is not None: fdset.append(p4.stdout) @@ -131,9 +138,14 @@ try: if stage == 5: r = p4.wait() - if r != 0: - sys.exit(1) - print('test passed') + if config.should_fail: + if r == 0: + sys.exit(1) + print('test passed (expecting failure)') + else: + if r != 0: + sys.exit(1) + print('test passed') break finally: for p in [p1, p2]: