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 sys
import getopt import getopt
import logging import logging
import traceback
from functools import wraps from functools import wraps
@ -56,20 +57,28 @@ def print_exception(e):
traceback.print_exc() traceback.print_exc()
def exception_handle(self_, err_msg=None, exit_code=None): def exception_handle(self_, err_msg=None, exit_code=None,
""" destroy=False, conn_err=False):
:param self_: if function passes self as first arg # self_: if function passes self as first arg
:param err_msg:
:param exit_code: def process_exception(self, e):
:return:
"""
def process_exception(e):
print_exception(e) print_exception(e)
if err_msg: if err_msg:
logging.error(err_msg) logging.error(err_msg)
if exit_code: if exit_code:
sys.exit(1) 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): def decorator(func):
if self_: if self_:
@wraps(func) @wraps(func)
@ -77,14 +86,14 @@ def exception_handle(self_, err_msg=None, exit_code=None):
try: try:
func(self, *args, **kwargs) func(self, *args, **kwargs)
except Exception as e: except Exception as e:
process_exception(e) process_exception(self, e)
else: else:
@wraps(func) @wraps(func)
def wrapper(*args, **kwargs): def wrapper(*args, **kwargs):
try: try:
func(*args, **kwargs) func(*args, **kwargs)
except Exception as e: except Exception as e:
process_exception(e) process_exception(self, e)
return wrapper return wrapper
return decorator return decorator

21
shadowsocks/tcprelay.py
View file

@ -296,8 +296,8 @@ class TCPRelayHandler(object):
traceback.print_exc() traceback.print_exc()
self.destroy() self.destroy()
@shell.exception_handle(self_=True, destroy=True, conn_err=True)
def _handle_stage_addr(self, data): def _handle_stage_addr(self, data):
try:
if self._is_local: if self._is_local:
cmd = common.ord(data[1]) cmd = common.ord(data[1])
if cmd == CMD_UDP_ASSOCIATE: if cmd == CMD_UDP_ASSOCIATE:
@ -378,11 +378,6 @@ class TCPRelayHandler(object):
# notice here may go into _handle_dns_resolved directly # notice here may go into _handle_dns_resolved directly
self._dns_resolver.resolve(remote_addr, self._dns_resolver.resolve(remote_addr,
self._handle_dns_resolved) 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): def _create_remote_socket(self, ip, port):
addrs = socket.getaddrinfo(ip, port, 0, socket.SOCK_STREAM, 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) remote_sock.setsockopt(socket.SOL_TCP, socket.TCP_NODELAY, 1)
return remote_sock return remote_sock
@shell.exception_handle(self_=True)
def _handle_dns_resolved(self, result, error): def _handle_dns_resolved(self, result, error):
if 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() self.destroy()
return return
if not (result and result[1]): if not (result and result[1]):
@ -411,7 +408,6 @@ class TCPRelayHandler(object):
return return
ip = result[1] ip = result[1]
try:
self._stage = STAGE_CONNECTING self._stage = STAGE_CONNECTING
remote_addr = ip remote_addr = ip
if self._is_local: if self._is_local:
@ -443,11 +439,6 @@ class TCPRelayHandler(object):
self._stage = STAGE_CONNECTING self._stage = STAGE_CONNECTING
self._update_stream(STREAM_UP, WAIT_STATUS_READWRITING) self._update_stream(STREAM_UP, WAIT_STATUS_READWRITING)
self._update_stream(STREAM_DOWN, WAIT_STATUS_READING) 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): def _write_to_sock_remote(self, data):
self._write_to_sock(data, self._remote_sock) self._write_to_sock(data, self._remote_sock)
@ -666,10 +657,6 @@ class TCPRelayHandler(object):
else: else:
logging.warn('unknown socket') 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): def destroy(self):
# destroy the handler and release any resources # destroy the handler and release any resources
# promises: # promises: