add forbidden support for UDP and add tests

This commit is contained in:
clowwindy 2015-01-12 23:02:59 +08:00
parent eb94bd1cc3
commit 4312eb9e58
4 changed files with 34 additions and 10 deletions

View file

@ -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

View file

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

View file

@ -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

View file

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