From 6e5a753d41f54668d5ca98495bd409c2d0be3ac0 Mon Sep 17 00:00:00 2001 From: clowwindy Date: Thu, 22 May 2014 12:01:13 +0800 Subject: [PATCH] add timeout in TCP server --- shadowsocks/local.py | 6 +++--- shadowsocks/server.py | 25 +++++++++++++++++++------ shadowsocks/utils.py | 15 +++++++++------ 3 files changed, 31 insertions(+), 15 deletions(-) diff --git a/shadowsocks/local.py b/shadowsocks/local.py index 1bb9cec..5159a76 100755 --- a/shadowsocks/local.py +++ b/shadowsocks/local.py @@ -282,7 +282,7 @@ def main(): config_path = utils.find_config() try: - optlist, args = getopt.getopt(sys.argv[1:], 's:b:p:k:l:m:c:', + optlist, args = getopt.getopt(sys.argv[1:], 's:b:p:k:l:m:c:t:', ['fast-open']) for key, value in optlist: if key == '-c': @@ -300,7 +300,7 @@ def main(): else: config = {} - optlist, args = getopt.getopt(sys.argv[1:], 's:b:p:k:l:m:c:', + optlist, args = getopt.getopt(sys.argv[1:], 's:b:p:k:l:m:c:t:', ['fast-open']) for key, value in optlist: if key == '-p': @@ -328,7 +328,7 @@ def main(): config_password = config['password'] config_method = config.get('method', None) config_local_address = config.get('local_address', '127.0.0.1') - config_timeout = config.get('timeout', 600) + config_timeout = int(config.get('timeout', 300)) config_fast_open = config.get('fast_open', False) if not config_password and not config_path: diff --git a/shadowsocks/server.py b/shadowsocks/server.py index 180243a..67dfb55 100755 --- a/shadowsocks/server.py +++ b/shadowsocks/server.py @@ -74,6 +74,11 @@ class ThreadingTCPServer(SocketServer.ThreadingMixIn, SocketServer.TCPServer): logging.error('warning: fast open is not available') self.socket.listen(self.request_queue_size) + def get_request(self): + connection = self.socket.accept() + connection[0].settimeout(config_timeout) + return connection + class Socks5Server(SocketServer.StreamRequestHandler): def handle_tcp(self, sock, remote): @@ -81,7 +86,10 @@ class Socks5Server(SocketServer.StreamRequestHandler): fdset = [sock, remote] while True: should_break = False - r, w, e = select.select(fdset, [], []) + r, w, e = select.select(fdset, [], [], config_timeout) + if not r: + logging.warn('read time out') + break if sock in r: data = self.decrypt(sock.recv(4096)) if len(data) <= 0: @@ -147,7 +155,9 @@ class Socks5Server(SocketServer.StreamRequestHandler): port = struct.unpack('>H', self.decrypt(self.rfile.read(2))) try: logging.info('connecting %s:%d' % (addr, port[0])) - remote = socket.create_connection((addr, port[0])) + remote = socket.create_connection((addr, port[0]), + timeout=config_timeout) + remote.settimeout(config_timeout) remote.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) except socket.error, e: # Connection refused @@ -159,7 +169,8 @@ class Socks5Server(SocketServer.StreamRequestHandler): def main(): - global config_server, config_server_port, config_method, config_fast_open + global config_server, config_server_port, config_method, config_fast_open, \ + config_timeout logging.basicConfig(level=logging.DEBUG, format='%(asctime)s %(levelname)-8s %(message)s', @@ -176,7 +187,7 @@ def main(): config_path = utils.find_config() try: - optlist, args = getopt.getopt(sys.argv[1:], 's:p:k:m:c:', + optlist, args = getopt.getopt(sys.argv[1:], 's:p:k:m:c:t:', ['fast-open', 'workers:']) for key, value in optlist: if key == '-c': @@ -194,7 +205,7 @@ def main(): else: config = {} - optlist, args = getopt.getopt(sys.argv[1:], 's:p:k:m:c:', + optlist, args = getopt.getopt(sys.argv[1:], 's:p:k:m:c:t:', ['fast-open', 'workers=']) for key, value in optlist: if key == '-p': @@ -205,6 +216,8 @@ def main(): config['server'] = value elif key == '-m': config['method'] = value + elif key == '-t': + config['timeout'] = value elif key == '--fast-open': config['fast_open'] = True elif key == '--workers': @@ -218,7 +231,7 @@ def main(): config_key = config['password'] config_method = config.get('method', None) config_port_password = config.get('port_password', None) - config_timeout = config.get('timeout', 600) + config_timeout = int(config.get('timeout', 300)) config_fast_open = config.get('fast_open', False) config_workers = config.get('workers', 1) diff --git a/shadowsocks/utils.py b/shadowsocks/utils.py index eba8200..f37eed9 100644 --- a/shadowsocks/utils.py +++ b/shadowsocks/utils.py @@ -88,17 +88,18 @@ def check_config(config): if (config.get('method', '') or '').lower() == 'rc4': logging.warn('warning: RC4 is not safe; please use a safer cipher, ' 'like AES-256-CFB') - if (config.get('timeout', 600) or 600) < 100: + if (int(config.get('timeout', 300)) or 300) < 100: logging.warn('warning: your timeout %d seems too short' % - config.get('timeout')) - if (config.get('timeout', 600) or 600) > 600: + int(config.get('timeout'))) + if (int(config.get('timeout', 300)) or 300) > 600: logging.warn('warning: your timeout %d seems too long' % - config.get('timeout')) + int(config.get('timeout'))) def print_local_help(): print '''usage: sslocal [-h] -s SERVER_ADDR -p SERVER_PORT [-b LOCAL_ADDR] - -l LOCAL_PORT -k PASSWORD -m METHOD [-c config] [--fast-open] + -l LOCAL_PORT -k PASSWORD -m METHOD [-t TIMEOUT] [-c CONFIG] + [--fast-open] optional arguments: -h, --help show this help message and exit @@ -108,6 +109,7 @@ optional arguments: -l LOCAL_PORT local port -k PASSWORD password -m METHOD encryption method, for example, aes-256-cfb + -t TIMEOUT timeout in seconds -c CONFIG path to config file --fast-open use TCP_FASTOPEN, requires Linux 3.7+ ''' @@ -115,7 +117,7 @@ optional arguments: def print_server_help(): print '''usage: ssserver [-h] -s SERVER_ADDR -p SERVER_PORT -k PASSWORD - -m METHOD [-c config] [--fast-open] + -m METHOD [-t TIMEOUT] [-c CONFIG] [--fast-open] optional arguments: -h, --help show this help message and exit @@ -123,6 +125,7 @@ optional arguments: -p SERVER_PORT server port -k PASSWORD password -m METHOD encryption method, for example, aes-256-cfb + -t TIMEOUT timeout in seconds -c CONFIG path to config file --fast-open use TCP_FASTOPEN, requires Linux 3.7+ --workers WORKERS number of workers, available on Unix/Linux