vbatts/shadowsocks
1
0
Fork 0
Code Issues Pull requests Releases Wiki Activity

apply exception_handle decorator to tcprelay

Browse source
This commit is contained in:
ahxxm 2016-09-06 11:08:35 +08:00
parent 52556d5e48
commit bcb528901d
2 changed files with 130 additions and 134 deletions
Download patch file Download diff file Expand all files Collapse all files

29
shadowsocks/shell.py
View file

@ -23,6 +23,7 @@ import json
import sys
import getopt
import logging
import traceback
from functools import wraps
@ -56,20 +57,28 @@ 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):
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(self, e):
print_exception(e)
if err_msg:
logging.error(err_msg)
if exit_code:
sys.exit(1)
if not self_:
return
if conn_err:
logging.error('%s when handling connection from %s:%d' %
(e, self._client_address[0], self._client_address[1]))
if self._config['verbose']:
traceback.print_exc()
if destroy:
self.destroy()
def decorator(func):
if self_:
@wraps(func)
@ -77,14 +86,14 @@ def exception_handle(self_, err_msg=None, exit_code=None):
try:
func(self, *args, **kwargs)
except Exception as e:
process_exception(e)
process_exception(self, e)
else:
@wraps(func)
def wrapper(*args, **kwargs):
try:
func(*args, **kwargs)
except Exception as e:
process_exception(e)
process_exception(self, e)
return wrapper
return decorator

21
shadowsocks/tcprelay.py
View file

@ -296,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:
@ -378,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,
@ -401,9 +396,11 @@ 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)
logging.error('%s when handling connection from %s:%d' %
(error, self._client_address[0], self._client_address[1]))
self.destroy()
return
if not (result and result[1]):
@ -411,7 +408,6 @@ class TCPRelayHandler(object):
return
ip = result[1]
try:
self._stage = STAGE_CONNECTING
remote_addr = ip
if self._is_local:
@ -443,11 +439,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()
def _write_to_sock_remote(self, data):
self._write_to_sock(data, self._remote_sock)
@ -666,10 +657,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: