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

make tcprelay.py less nested

Browse source
This commit is contained in:
ahxxm 2016-09-06 10:13:10 +08:00
parent 5c11527e1b
commit 9f0c09463b
1 changed files with 67 additions and 62 deletions
Download patch file Download diff file Expand all files Collapse all files

129
shadowsocks/tcprelay.py
View file

@ -190,21 +190,23 @@ class TCPRelayHandler(object):
if self._upstream_status != status: if self._upstream_status != status:
self._upstream_status = status self._upstream_status = status
dirty = True dirty = True
if dirty: if not dirty:
if self._local_sock: return
event = eventloop.POLL_ERR
if self._downstream_status & WAIT_STATUS_WRITING: if self._local_sock:
event |= eventloop.POLL_OUT event = eventloop.POLL_ERR
if self._upstream_status & WAIT_STATUS_READING: if self._downstream_status & WAIT_STATUS_WRITING:
event |= eventloop.POLL_IN event |= eventloop.POLL_OUT
self._loop.modify(self._local_sock, event) if self._upstream_status & WAIT_STATUS_READING:
if self._remote_sock: event |= eventloop.POLL_IN
event = eventloop.POLL_ERR self._loop.modify(self._local_sock, event)
if self._downstream_status & WAIT_STATUS_READING: if self._remote_sock:
event |= eventloop.POLL_IN event = eventloop.POLL_ERR
if self._upstream_status & WAIT_STATUS_WRITING: if self._downstream_status & WAIT_STATUS_READING:
event |= eventloop.POLL_OUT event |= eventloop.POLL_IN
self._loop.modify(self._remote_sock, event) if self._upstream_status & WAIT_STATUS_WRITING:
event |= eventloop.POLL_OUT
self._loop.modify(self._remote_sock, event)
def _write_to_sock(self, data, sock): def _write_to_sock(self, data, sock):
# write data to sock # write data to sock
@ -247,19 +249,20 @@ class TCPRelayHandler(object):
return True return True
def _handle_stage_connecting(self, data): def _handle_stage_connecting(self, data):
if self._is_local: if not 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 self._ota_enable_session: if self._ota_enable_session:
self._ota_chunk_data(data, self._ota_chunk_data(data,
self._data_to_write_to_remote.append) self._data_to_write_to_remote.append)
else: else:
self._data_to_write_to_remote.append(data) self._data_to_write_to_remote.append(data)
if self._is_local and not self._fastopen_connected and \ return
self._config['fast_open']:
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 # for sslocal and fastopen, we basically wait for data and use
# sendto to connect # sendto to connect
try: try:
@ -403,46 +406,48 @@ class TCPRelayHandler(object):
self._log_error(error) self._log_error(error)
self.destroy() self.destroy()
return return
if result and result[1]: if not (result and result[1]):
ip = result[1] self.destroy()
try: return
self._stage = STAGE_CONNECTING
remote_addr = ip
if self._is_local:
remote_port = self._chosen_server[1]
else:
remote_port = self._remote_address[1]
if self._is_local and self._config['fast_open']: ip = result[1]
# for fastopen: try:
# wait for more data arrive and send them in one SYN self._stage = STAGE_CONNECTING
self._stage = STAGE_CONNECTING remote_addr = ip
# we don't have to wait for remote since it's not if self._is_local:
# created remote_port = self._chosen_server[1]
self._update_stream(STREAM_UP, WAIT_STATUS_READING) else:
# TODO when there is already data in this packet remote_port = self._remote_address[1]
else:
# else do connect if self._is_local and self._config['fast_open']:
remote_sock = self._create_remote_socket(remote_addr, # for fastopen:
remote_port) # wait for more data arrive and send them in one SYN
try: self._stage = STAGE_CONNECTING
remote_sock.connect((remote_addr, remote_port)) # we don't have to wait for remote since it's not
except (OSError, IOError) as e: # created
if eventloop.errno_from_exception(e) == \ self._update_stream(STREAM_UP, WAIT_STATUS_READING)
errno.EINPROGRESS: # TODO when there is already data in this packet
pass else:
self._loop.add(remote_sock, # else do connect
eventloop.POLL_ERR | eventloop.POLL_OUT, remote_sock = self._create_remote_socket(remote_addr,
self._server) remote_port)
self._stage = STAGE_CONNECTING try:
self._update_stream(STREAM_UP, WAIT_STATUS_READWRITING) remote_sock.connect((remote_addr, remote_port))
self._update_stream(STREAM_DOWN, WAIT_STATUS_READING) except (OSError, IOError) as e:
return if eventloop.errno_from_exception(e) == \
except Exception as e: errno.EINPROGRESS:
shell.print_exception(e) pass
if self._config['verbose']: self._loop.add(remote_sock,
traceback.print_exc() eventloop.POLL_ERR | eventloop.POLL_OUT,
self.destroy() self._server)
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): def _write_to_sock_remote(self, data):
self._write_to_sock(data, self._remote_sock) self._write_to_sock(data, self._remote_sock)