add timeout in TCP server

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:

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)
 

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