From 52556d5e48eb6aadb6ef8267e68b5145bd598884 Mon Sep 17 00:00:00 2001
From: ahxxm <i@ahxxm.com>
Date: Tue, 6 Sep 2016 10:51:23 +0800
Subject: [PATCH] introduce exception_handle decorator

make try/except block more clean
---
 shadowsocks/local.py | 44 ++++++++++++++++++++------------------------
 shadowsocks/shell.py | 37 +++++++++++++++++++++++++++++++++++++
 2 files changed, 57 insertions(+), 24 deletions(-)

diff --git a/shadowsocks/local.py b/shadowsocks/local.py
index 4255a2e..dfc8032 100755
--- a/shadowsocks/local.py
+++ b/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,36 +38,31 @@ 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']))
+    logging.info("starting local at %s:%d" %
+                 (config['local_address'], config['local_port']))
 
-        dns_resolver = asyncdns.DNSResolver()
-        tcp_server = tcprelay.TCPRelay(config, dns_resolver, True)
-        udp_server = udprelay.UDPRelay(config, dns_resolver, True)
-        loop = eventloop.EventLoop()
-        dns_resolver.add_to_loop(loop)
-        tcp_server.add_to_loop(loop)
-        udp_server.add_to_loop(loop)
+    dns_resolver = asyncdns.DNSResolver()
+    tcp_server = tcprelay.TCPRelay(config, dns_resolver, True)
+    udp_server = udprelay.UDPRelay(config, dns_resolver, True)
+    loop = eventloop.EventLoop()
+    dns_resolver.add_to_loop(loop)
+    tcp_server.add_to_loop(loop)
+    udp_server.add_to_loop(loop)
 
-        def handler(signum, _):
-            logging.warn('received SIGQUIT, doing graceful shutting down..')
-            tcp_server.close(next_tick=True)
-            udp_server.close(next_tick=True)
-        signal.signal(getattr(signal, 'SIGQUIT', signal.SIGTERM), handler)
+    def handler(signum, _):
+        logging.warn('received SIGQUIT, doing graceful shutting down..')
+        tcp_server.close(next_tick=True)
+        udp_server.close(next_tick=True)
+    signal.signal(getattr(signal, 'SIGQUIT', signal.SIGTERM), handler)
 
-        def int_handler(signum, _):
-            sys.exit(1)
-        signal.signal(signal.SIGINT, int_handler)
-
-        daemon.set_user(config.get('user', None))
-        loop.run()
-    except Exception as e:
-        shell.print_exception(e)
+    def int_handler(signum, _):
         sys.exit(1)
+    signal.signal(signal.SIGINT, int_handler)
+
+    daemon.set_user(config.get('user', None))
+    loop.run()
 
 if __name__ == '__main__':
     main()
diff --git a/shadowsocks/shell.py b/shadowsocks/shell.py
index 6cca837..3ec830a 100644
--- a/shadowsocks/shell.py
+++ b/shadowsocks/shell.py
@@ -23,6 +23,9 @@ import json
 import sys
 import getopt
 import logging
+
+from functools import wraps
+
 from shadowsocks.common import to_bytes, to_str, IPNetwork
 from shadowsocks import encrypt
 
@@ -53,6 +56,40 @@ def print_exception(e):
         traceback.print_exc()
 
 
+def exception_handle(self_, err_msg=None, exit_code=None):
+    """
+    :param self_: if function passes self as first arg
+    :param err_msg:
+    :param exit_code:
+    :return:
+    """
+    def process_exception(e):
+        print_exception(e)
+        if err_msg:
+            logging.error(err_msg)
+        if exit_code:
+            sys.exit(1)
+
+    def decorator(func):
+        if self_:
+            @wraps(func)
+            def wrapper(self, *args, **kwargs):
+                try:
+                    func(self, *args, **kwargs)
+                except Exception as e:
+                    process_exception(e)
+        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: