rename things
This commit is contained in:
parent
d69272767f
commit
bbad408ace
2 changed files with 26 additions and 27 deletions
|
@ -194,7 +194,6 @@ class EventLoop(object):
|
||||||
|
|
||||||
def run(self):
|
def run(self):
|
||||||
while not self.stopping:
|
while not self.stopping:
|
||||||
events = None
|
|
||||||
try:
|
try:
|
||||||
events = self.poll(1)
|
events = self.poll(1)
|
||||||
except (OSError, IOError) as e:
|
except (OSError, IOError) as e:
|
||||||
|
|
|
@ -51,11 +51,11 @@ STAGE_STREAM = 5
|
||||||
STREAM_UP = 0
|
STREAM_UP = 0
|
||||||
STREAM_DOWN = 1
|
STREAM_DOWN = 1
|
||||||
|
|
||||||
# stream status
|
# stream wait status
|
||||||
STATUS_WAIT_INIT = 0
|
WAIT_STATUS_INIT = 0
|
||||||
STATUS_WAIT_READING = 1
|
WAIT_STATUS_READING = 1
|
||||||
STATUS_WAIT_WRITING = 2
|
WAIT_STATUS_WRITING = 2
|
||||||
STATUS_WAIT_READWRITING = STATUS_WAIT_READING | STATUS_WAIT_WRITING
|
WAIT_STATUS_READWRITING = WAIT_STATUS_READING | WAIT_STATUS_WRITING
|
||||||
|
|
||||||
BUF_SIZE = 8 * 1024
|
BUF_SIZE = 8 * 1024
|
||||||
|
|
||||||
|
@ -73,8 +73,8 @@ class TCPRelayHandler(object):
|
||||||
config['method'])
|
config['method'])
|
||||||
self._data_to_write_to_local = []
|
self._data_to_write_to_local = []
|
||||||
self._data_to_write_to_remote = []
|
self._data_to_write_to_remote = []
|
||||||
self._upstream_status = STATUS_WAIT_READING
|
self._upstream_status = WAIT_STATUS_READING
|
||||||
self._downstream_status = STATUS_WAIT_INIT
|
self._downstream_status = WAIT_STATUS_INIT
|
||||||
fd_to_handlers[local_sock.fileno()] = self
|
fd_to_handlers[local_sock.fileno()] = self
|
||||||
local_sock.setblocking(False)
|
local_sock.setblocking(False)
|
||||||
local_sock.setsockopt(socket.SOL_TCP, socket.TCP_NODELAY, 1)
|
local_sock.setsockopt(socket.SOL_TCP, socket.TCP_NODELAY, 1)
|
||||||
|
@ -93,20 +93,20 @@ class TCPRelayHandler(object):
|
||||||
if dirty:
|
if dirty:
|
||||||
if self._local_sock:
|
if self._local_sock:
|
||||||
event = eventloop.POLL_ERR
|
event = eventloop.POLL_ERR
|
||||||
if self._downstream_status & STATUS_WAIT_WRITING:
|
if self._downstream_status & WAIT_STATUS_WRITING:
|
||||||
event |= eventloop.POLL_OUT
|
event |= eventloop.POLL_OUT
|
||||||
if self._upstream_status & STATUS_WAIT_READING:
|
if self._upstream_status & WAIT_STATUS_READING:
|
||||||
event |= eventloop.POLL_IN
|
event |= eventloop.POLL_IN
|
||||||
self._loop.modify(self._local_sock, event)
|
self._loop.modify(self._local_sock, event)
|
||||||
if self._remote_sock:
|
if self._remote_sock:
|
||||||
event = eventloop.POLL_ERR
|
event = eventloop.POLL_ERR
|
||||||
if self._downstream_status & STATUS_WAIT_READING:
|
if self._downstream_status & WAIT_STATUS_READING:
|
||||||
event |= eventloop.POLL_IN
|
event |= eventloop.POLL_IN
|
||||||
if self._upstream_status & STATUS_WAIT_WRITING:
|
if self._upstream_status & WAIT_STATUS_WRITING:
|
||||||
event |= eventloop.POLL_OUT
|
event |= eventloop.POLL_OUT
|
||||||
self._loop.modify(self._remote_sock, event)
|
self._loop.modify(self._remote_sock, event)
|
||||||
|
|
||||||
def write_all_to_sock(self, data, sock):
|
def write_to_sock(self, data, sock):
|
||||||
if not data or not sock:
|
if not data or not sock:
|
||||||
return
|
return
|
||||||
uncomplete = False
|
uncomplete = False
|
||||||
|
@ -126,17 +126,17 @@ class TCPRelayHandler(object):
|
||||||
if uncomplete:
|
if uncomplete:
|
||||||
if sock == self._local_sock:
|
if sock == self._local_sock:
|
||||||
self._data_to_write_to_local.append(data)
|
self._data_to_write_to_local.append(data)
|
||||||
self.update_stream(STREAM_DOWN, STATUS_WAIT_WRITING)
|
self.update_stream(STREAM_DOWN, WAIT_STATUS_WRITING)
|
||||||
elif sock == self._remote_sock:
|
elif sock == self._remote_sock:
|
||||||
self._data_to_write_to_remote.append(data)
|
self._data_to_write_to_remote.append(data)
|
||||||
self.update_stream(STREAM_UP, STATUS_WAIT_WRITING)
|
self.update_stream(STREAM_UP, WAIT_STATUS_WRITING)
|
||||||
else:
|
else:
|
||||||
logging.error('write_all_to_sock:unknown socket')
|
logging.error('write_all_to_sock:unknown socket')
|
||||||
else:
|
else:
|
||||||
if sock == self._local_sock:
|
if sock == self._local_sock:
|
||||||
self.update_stream(STREAM_DOWN, STATUS_WAIT_READING)
|
self.update_stream(STREAM_DOWN, WAIT_STATUS_READING)
|
||||||
elif sock == self._remote_sock:
|
elif sock == self._remote_sock:
|
||||||
self.update_stream(STREAM_UP, STATUS_WAIT_READING)
|
self.update_stream(STREAM_UP, WAIT_STATUS_READING)
|
||||||
else:
|
else:
|
||||||
logging.error('write_all_to_sock:unknown socket')
|
logging.error('write_all_to_sock:unknown socket')
|
||||||
|
|
||||||
|
@ -162,11 +162,11 @@ class TCPRelayHandler(object):
|
||||||
if self._stage == STAGE_STREAM:
|
if self._stage == STAGE_STREAM:
|
||||||
if self._is_local:
|
if self._is_local:
|
||||||
data = self._encryptor.encrypt(data)
|
data = self._encryptor.encrypt(data)
|
||||||
self.write_all_to_sock(data, self._remote_sock)
|
self.write_to_sock(data, self._remote_sock)
|
||||||
return
|
return
|
||||||
elif is_local and self._stage == STAGE_INIT:
|
elif is_local and self._stage == STAGE_INIT:
|
||||||
# TODO check auth method
|
# TODO check auth method
|
||||||
self.write_all_to_sock('\x05\00', self._local_sock)
|
self.write_to_sock('\x05\00', self._local_sock)
|
||||||
self._stage = STAGE_HELLO
|
self._stage = STAGE_HELLO
|
||||||
return
|
return
|
||||||
elif self._stage == STAGE_REPLY:
|
elif self._stage == STAGE_REPLY:
|
||||||
|
@ -190,7 +190,7 @@ class TCPRelayHandler(object):
|
||||||
logging.debug('connecting %s:%d' % (remote_addr, remote_port))
|
logging.debug('connecting %s:%d' % (remote_addr, remote_port))
|
||||||
if is_local:
|
if is_local:
|
||||||
# forward address to remote
|
# forward address to remote
|
||||||
self.write_all_to_sock('\x05\x00\x00\x01' +
|
self.write_to_sock('\x05\x00\x00\x01' +
|
||||||
'\x00\x00\x00\x00\x10\x10',
|
'\x00\x00\x00\x00\x10\x10',
|
||||||
self._local_sock)
|
self._local_sock)
|
||||||
data_to_send = self._encryptor.encrypt(data)
|
data_to_send = self._encryptor.encrypt(data)
|
||||||
|
@ -224,8 +224,8 @@ class TCPRelayHandler(object):
|
||||||
eventloop.POLL_ERR | eventloop.POLL_OUT)
|
eventloop.POLL_ERR | eventloop.POLL_OUT)
|
||||||
|
|
||||||
self._stage = STAGE_REPLY
|
self._stage = STAGE_REPLY
|
||||||
self.update_stream(STREAM_UP, STATUS_WAIT_READWRITING)
|
self.update_stream(STREAM_UP, WAIT_STATUS_READWRITING)
|
||||||
self.update_stream(STREAM_DOWN, STATUS_WAIT_READING)
|
self.update_stream(STREAM_DOWN, WAIT_STATUS_READING)
|
||||||
return
|
return
|
||||||
except Exception:
|
except Exception:
|
||||||
import traceback
|
import traceback
|
||||||
|
@ -250,7 +250,7 @@ class TCPRelayHandler(object):
|
||||||
else:
|
else:
|
||||||
data = self._encryptor.encrypt(data)
|
data = self._encryptor.encrypt(data)
|
||||||
try:
|
try:
|
||||||
self.write_all_to_sock(data, self._local_sock)
|
self.write_to_sock(data, self._local_sock)
|
||||||
except Exception:
|
except Exception:
|
||||||
import traceback
|
import traceback
|
||||||
traceback.print_exc()
|
traceback.print_exc()
|
||||||
|
@ -261,18 +261,18 @@ class TCPRelayHandler(object):
|
||||||
if self._data_to_write_to_local:
|
if self._data_to_write_to_local:
|
||||||
data = ''.join(self._data_to_write_to_local)
|
data = ''.join(self._data_to_write_to_local)
|
||||||
self._data_to_write_to_local = []
|
self._data_to_write_to_local = []
|
||||||
self.write_all_to_sock(data, self._local_sock)
|
self.write_to_sock(data, self._local_sock)
|
||||||
else:
|
else:
|
||||||
self.update_stream(STREAM_DOWN, STATUS_WAIT_READING)
|
self.update_stream(STREAM_DOWN, WAIT_STATUS_READING)
|
||||||
|
|
||||||
def on_remote_write(self):
|
def on_remote_write(self):
|
||||||
self._stage = STAGE_STREAM
|
self._stage = STAGE_STREAM
|
||||||
if self._data_to_write_to_remote:
|
if self._data_to_write_to_remote:
|
||||||
data = ''.join(self._data_to_write_to_remote)
|
data = ''.join(self._data_to_write_to_remote)
|
||||||
self._data_to_write_to_remote = []
|
self._data_to_write_to_remote = []
|
||||||
self.write_all_to_sock(data, self._remote_sock)
|
self.write_to_sock(data, self._remote_sock)
|
||||||
else:
|
else:
|
||||||
self.update_stream(STREAM_UP, STATUS_WAIT_READING)
|
self.update_stream(STREAM_UP, WAIT_STATUS_READING)
|
||||||
|
|
||||||
def on_local_error(self):
|
def on_local_error(self):
|
||||||
logging.error(eventloop.get_sock_error(self._local_sock))
|
logging.error(eventloop.get_sock_error(self._local_sock))
|
||||||
|
|
Loading…
Reference in a new issue