apply exception_handle decorator to tcprelay
This commit is contained in:
parent
52556d5e48
commit
bcb528901d
2 changed files with 130 additions and 134 deletions
|
@ -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
|
||||
|
|
|
@ -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:
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue