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
This commit is contained in:
parent
5c11527e1b
commit
5cd9f04948
|
@ -276,10 +276,13 @@ class DNSResolver(object):
|
|||
content = f.readlines()
|
||||
for line in content:
|
||||
line = line.strip()
|
||||
if line:
|
||||
if line.startswith(b'nameserver'):
|
||||
if not (line and 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:
|
||||
|
|
|
@ -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
|
||||
|
||||
|
|
|
@ -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()
|
||||
|
|
|
@ -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:
|
||||
|
|
|
@ -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:
|
||||
|
|
|
@ -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 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 not self._is_local:
|
||||
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 \
|
||||
self._config['fast_open']:
|
||||
return
|
||||
|
||||
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:
|
||||
|
|
Loading…
Reference in New Issue