add forbidden support for UDP and add tests

.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

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:

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

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]: