Refactor (#615)

* make tcprelay.py less nested

* import traceback module at the top

* make loops in DNSResolver less nested

* make manager.py less nested

* introduce exception_handle decorator

make try/except block more clean

* apply exception_handle decorator to tcprelay

* quote condition judgement

* pep8 fix

shadowsocks/asyncdns.py

@@ -276,10 +276,13 @@ class DNSResolver(object):
                 content = f.readlines()
                 for line in content:
                     line = line.strip()
-                    if line:
+                    if not (line and line.startswith(b'nameserver')):
-                        if line.startswith(b'nameserver'):
+                        continue
+
                     parts = line.split()
-                            if len(parts) >= 2:
+                    if len(parts) < 2:
+                        continue
+
                     server = parts[1]
                     if common.is_ip(server) == socket.AF_INET:
                         if type(server) != str:
@@ -299,9 +302,13 @@ class DNSResolver(object):
                 for line in f.readlines():
                     line = line.strip()
                     parts = line.split()
-                    if len(parts) >= 2:
+                    if len(parts) < 2:
+                        continue
+
                     ip = parts[0]
-                        if common.is_ip(ip):
+                    if not common.is_ip(ip):
+                        continue
+
                     for i in range(1, len(parts)):
                         hostname = parts[i]
                         if hostname:

shadowsocks/eventloop.py

@@ -25,6 +25,7 @@ import os
 import time
 import socket
 import select
+import traceback
 import errno
 import logging
 from collections import defaultdict
@@ -204,7 +205,6 @@ class EventLoop(object):
                     logging.debug('poll:%s', e)
                 else:
                     logging.error('poll:%s', e)
-                    import traceback
                     traceback.print_exc()
                     continue
 

shadowsocks/local.py

@@ -27,6 +27,7 @@ sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../'))
 from shadowsocks import shell, daemon, eventloop, tcprelay, udprelay, asyncdns
 
 
+@shell.exception_handle(self_=False, exit_code=1)
 def main():
     shell.check_python()
 
@@ -37,10 +38,8 @@ def main():
         os.chdir(p)
 
     config = shell.get_config(True)
-
     daemon.daemon_exec(config)
 
-    try:
     logging.info("starting local at %s:%d" %
                  (config['local_address'], config['local_port']))
 
@@ -64,9 +63,6 @@ def main():
 
     daemon.set_user(config.get('user', None))
     loop.run()
-    except Exception as e:
-        shell.print_exception(e)
-        sys.exit(1)
 
 if __name__ == '__main__':
     main()

shadowsocks/manager.py

@@ -173,7 +173,9 @@ class Manager(object):
         self._statistics.clear()
 
     def _send_control_data(self, data):
-        if self._control_client_addr:
+        if not self._control_client_addr:
+            return
+
         try:
             self._control_socket.sendto(data, self._control_client_addr)
         except (socket.error, OSError, IOError) as e:

shadowsocks/shell.py

@@ -23,6 +23,10 @@ import json
 import sys
 import getopt
 import logging
+import traceback
+
+from functools import wraps
+
 from shadowsocks.common import to_bytes, to_str, IPNetwork
 from shadowsocks import encrypt
 
@@ -53,6 +57,49 @@ def print_exception(e):
         traceback.print_exc()
 
 
+def exception_handle(self_, err_msg=None, exit_code=None,
+                     destroy=False, conn_err=False):
+    # self_: if function passes self as first arg
+
+    def process_exception(e, self=None):
+        print_exception(e)
+        if err_msg:
+            logging.error(err_msg)
+        if exit_code:
+            sys.exit(1)
+
+        if not self_:
+            return
+
+        if conn_err:
+            addr, port = self._client_address[0], self._client_address[1]
+            logging.error('%s when handling connection from %s:%d' %
+                          (e, addr, port))
+        if self._config['verbose']:
+            traceback.print_exc()
+        if destroy:
+            self.destroy()
+
+    def decorator(func):
+        if self_:
+            @wraps(func)
+            def wrapper(self, *args, **kwargs):
+                try:
+                    func(self, *args, **kwargs)
+                except Exception as e:
+                    process_exception(e, self)
+        else:
+            @wraps(func)
+            def wrapper(*args, **kwargs):
+                try:
+                    func(*args, **kwargs)
+                except Exception as e:
+                    process_exception(e)
+
+        return wrapper
+    return decorator
+
+
 def print_shadowsocks():
     version = ''
     try:

shadowsocks/tcprelay.py

@@ -190,7 +190,9 @@ class TCPRelayHandler(object):
             if self._upstream_status != status:
                 self._upstream_status = status
                 dirty = True
-        if dirty:
+        if not dirty:
+            return
+
         if self._local_sock:
             event = eventloop.POLL_ERR
             if self._downstream_status & WAIT_STATUS_WRITING:
@@ -247,19 +249,20 @@ class TCPRelayHandler(object):
         return True
 
     def _handle_stage_connecting(self, data):
-        if self._is_local:
+        if not self._is_local:
-            if self._ota_enable_session:
-                data = self._ota_chunk_data_gen(data)
-            data = self._encryptor.encrypt(data)
-            self._data_to_write_to_remote.append(data)
-        else:
             if self._ota_enable_session:
                 self._ota_chunk_data(data,
                                      self._data_to_write_to_remote.append)
             else:
                 self._data_to_write_to_remote.append(data)
-        if self._is_local and not self._fastopen_connected and \
+            return
-                self._config['fast_open']:
+
+        if self._ota_enable_session:
+            data = self._ota_chunk_data_gen(data)
+        data = self._encryptor.encrypt(data)
+        self._data_to_write_to_remote.append(data)
+
+        if self._config['fast_open'] and not self._fastopen_connected:
             # for sslocal and fastopen, we basically wait for data and use
             # sendto to connect
             try:
@@ -293,8 +296,8 @@ class TCPRelayHandler(object):
                         traceback.print_exc()
                     self.destroy()
 
+    @shell.exception_handle(self_=True, destroy=True, conn_err=True)
     def _handle_stage_addr(self, data):
-        try:
         if self._is_local:
             cmd = common.ord(data[1])
             if cmd == CMD_UDP_ASSOCIATE:
@@ -375,11 +378,6 @@ class TCPRelayHandler(object):
             # notice here may go into _handle_dns_resolved directly
             self._dns_resolver.resolve(remote_addr,
                                        self._handle_dns_resolved)
-        except Exception as e:
-            self._log_error(e)
-            if self._config['verbose']:
-                traceback.print_exc()
-            self.destroy()
 
     def _create_remote_socket(self, ip, port):
         addrs = socket.getaddrinfo(ip, port, 0, socket.SOCK_STREAM,
@@ -398,14 +396,19 @@ class TCPRelayHandler(object):
         remote_sock.setsockopt(socket.SOL_TCP, socket.TCP_NODELAY, 1)
         return remote_sock
 
+    @shell.exception_handle(self_=True)
     def _handle_dns_resolved(self, result, error):
         if error:
-            self._log_error(error)
+            addr, port = self._client_address[0], self._client_address[1]
+            logging.error('%s when handling connection from %s:%d' %
+                          (error, addr, port))
             self.destroy()
             return
-        if result and result[1]:
+        if not (result and result[1]):
+            self.destroy()
+            return
+
         ip = result[1]
-            try:
         self._stage = STAGE_CONNECTING
         remote_addr = ip
         if self._is_local:
@@ -437,12 +440,6 @@ class TCPRelayHandler(object):
             self._stage = STAGE_CONNECTING
             self._update_stream(STREAM_UP, WAIT_STATUS_READWRITING)
             self._update_stream(STREAM_DOWN, WAIT_STATUS_READING)
-                return
-            except Exception as e:
-                shell.print_exception(e)
-                if self._config['verbose']:
-                    traceback.print_exc()
-        self.destroy()
 
     def _write_to_sock_remote(self, data):
         self._write_to_sock(data, self._remote_sock)
@@ -661,10 +658,6 @@ class TCPRelayHandler(object):
         else:
             logging.warn('unknown socket')
 
-    def _log_error(self, e):
-        logging.error('%s when handling connection from %s:%d' %
-                      (e, self._client_address[0], self._client_address[1]))
-
     def destroy(self):
         # destroy the handler and release any resources
         # promises: