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

add event names

Browse source
This commit is contained in:
clowwindy 2014-06-01 16:21:33 +08:00
parent c721a1c02f
commit f49d086d6a
3 changed files with 60 additions and 27 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

19
shadowsocks/eventloop.py
View file

@ -32,7 +32,7 @@ from collections import defaultdict
__all__ = ['EventLoop', 'POLL_NULL', 'POLL_IN', 'POLL_OUT', 'POLL_ERR', __all__ = ['EventLoop', 'POLL_NULL', 'POLL_IN', 'POLL_OUT', 'POLL_ERR',
'POLL_HUP', 'POLL_NVAL'] 'POLL_HUP', 'POLL_NVAL', 'EVENT_NAMES']
POLL_NULL = 0x00 POLL_NULL = 0x00
POLL_IN = 0x01 POLL_IN = 0x01
@ -42,6 +42,16 @@ POLL_HUP = 0x10
POLL_NVAL = 0x20 POLL_NVAL = 0x20
EVENT_NAMES = {
POLL_NULL: 'POLL_NULL',
POLL_IN: 'POLL_IN',
POLL_OUT: 'POLL_OUT',
POLL_ERR: 'POLL_ERR',
POLL_HUP: 'POLL_HUP',
POLL_NVAL: 'POLL_NVAL',
}
class EpollLoop(object): class EpollLoop(object):
def __init__(self): def __init__(self):
@ -144,15 +154,22 @@ class EventLoop(object):
def __init__(self): def __init__(self):
if hasattr(select, 'epoll'): if hasattr(select, 'epoll'):
self._impl = EpollLoop() self._impl = EpollLoop()
self._model = 'epoll'
elif hasattr(select, 'kqueue'): elif hasattr(select, 'kqueue'):
self._impl = KqueueLoop() self._impl = KqueueLoop()
self._model = 'kqueue'
elif hasattr(select, 'select'): elif hasattr(select, 'select'):
self._impl = SelectLoop() self._impl = SelectLoop()
self._model = 'select'
else: else:
raise Exception('can not find any available functions in select ' raise Exception('can not find any available functions in select '
'package') 'package')
self._fd_to_f = {} self._fd_to_f = {}
@property
def model(self):
return self._model
def poll(self, timeout=None): def poll(self, timeout=None):
events = self._impl.poll(timeout) events = self._impl.poll(timeout)
return ((self._fd_to_f[fd], event) for fd, event in events) return ((self._fd_to_f[fd], event) for fd, event in events)

2
shadowsocks/local.py
View file

@ -22,7 +22,6 @@
# SOFTWARE. # SOFTWARE.
import sys import sys
import socket
import os import os
import logging import logging
import encrypt import encrypt
@ -50,6 +49,7 @@ def main():
logging.info("starting local at %s:%d" % logging.info("starting local at %s:%d" %
(config['local_address'], config['local_port'])) (config['local_address'], config['local_port']))
# TODO combine the two threads into one loop on a single thread
udprelay.UDPRelay(config, True).start() udprelay.UDPRelay(config, True).start()
tcprelay.TCPRelay(config, True).start() tcprelay.TCPRelay(config, True).start()
while sys.stdin.read(): while sys.stdin.read():

66
shadowsocks/tcprelay.py
View file

@ -56,6 +56,7 @@ STREAM_DOWN = 1
STATUS_WAIT_INIT = 0 STATUS_WAIT_INIT = 0
STATUS_WAIT_READING = 1 STATUS_WAIT_READING = 1
STATUS_WAIT_WRITING = 2 STATUS_WAIT_WRITING = 2
STATUS_WAIT_READWRITING = STATUS_WAIT_READING | STATUS_WAIT_WRITING
BUF_SIZE = 8 * 1024 BUF_SIZE = 8 * 1024
@ -73,10 +74,11 @@ 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_INIT self._upstream_status = STATUS_WAIT_READING
self._downstream_status = STATUS_WAIT_INIT self._downstream_status = STATUS_WAIT_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)
loop.add(local_sock, eventloop.POLL_IN | eventloop.POLL_ERR) loop.add(local_sock, eventloop.POLL_IN | eventloop.POLL_ERR)
def update_stream(self, stream, status): def update_stream(self, stream, status):
@ -92,16 +94,16 @@ 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 & STATUS_WAIT_WRITING:
event |= eventloop.POLL_OUT event |= eventloop.POLL_OUT
if self._upstream_status == STATUS_WAIT_READING: if self._upstream_status & STATUS_WAIT_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 & STATUS_WAIT_READING:
event |= eventloop.POLL_IN event |= eventloop.POLL_IN
if self._upstream_status == STATUS_WAIT_WRITING: if self._upstream_status & STATUS_WAIT_WRITING:
event |= eventloop.POLL_OUT event |= eventloop.POLL_OUT
self._loop.modify(self._remote_sock, event) self._loop.modify(self._remote_sock, event)
@ -131,6 +133,13 @@ class TCPRelayHandler(object):
self.update_stream(STREAM_UP, STATUS_WAIT_WRITING) self.update_stream(STREAM_UP, STATUS_WAIT_WRITING)
else: else:
logging.error('write_all_to_sock:unknown socket') logging.error('write_all_to_sock:unknown socket')
else:
if sock == self._local_sock:
self.update_stream(STREAM_DOWN, STATUS_WAIT_READING)
elif sock == self._remote_sock:
self.update_stream(STREAM_UP, STATUS_WAIT_READING)
else:
logging.error('write_all_to_sock:unknown socket')
def on_local_read(self): def on_local_read(self):
# TODO update timeout # TODO update timeout
@ -156,14 +165,16 @@ class TCPRelayHandler(object):
data = self._encryptor.encrypt(data) data = self._encryptor.encrypt(data)
self.write_all_to_sock(data, self._remote_sock) self.write_all_to_sock(data, self._remote_sock)
return return
if 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_all_to_sock('\x05\00', self._local_sock)
self._stage = STAGE_HELLO self._stage = STAGE_HELLO
return return
if self._stage == STAGE_REPLY: elif self._stage == STAGE_REPLY:
if is_local:
data = self._encryptor.encrypt(data)
self._data_to_write_to_remote.append(data) self._data_to_write_to_remote.append(data)
if (is_local and self._stage == STAGE_HELLO) or \ elif (is_local and self._stage == STAGE_HELLO) or \
(not is_local and self._stage == STAGE_INIT): (not is_local and self._stage == STAGE_INIT):
try: try:
if is_local: if is_local:
@ -199,20 +210,22 @@ class TCPRelayHandler(object):
raise Exception("can't get addrinfo for %s:%d" % raise Exception("can't get addrinfo for %s:%d" %
(remote_addr, remote_port)) (remote_addr, remote_port))
af, socktype, proto, canonname, sa = addrs[0] af, socktype, proto, canonname, sa = addrs[0]
self._remote_sock = socket.socket(af, socktype, proto) remote_sock = socket.socket(af, socktype, proto)
self._fd_to_handlers[self._remote_sock.fileno()] = self self._remote_sock = remote_sock
self._remote_sock.setblocking(False) self._fd_to_handlers[remote_sock.fileno()] = self
remote_sock.setblocking(False)
remote_sock.setsockopt(socket.SOL_TCP, socket.TCP_NODELAY, 1)
# TODO support TCP fast open # TODO support TCP fast open
try: try:
self._remote_sock.connect(sa) remote_sock.connect(sa)
except (OSError, IOError) as e: except (OSError, IOError) as e:
if eventloop.errno_from_exception(e) == errno.EINPROGRESS: if eventloop.errno_from_exception(e) == errno.EINPROGRESS:
pass pass
self._loop.add(self._remote_sock, self._loop.add(remote_sock,
eventloop.POLL_ERR | eventloop.POLL_OUT) eventloop.POLL_ERR | eventloop.POLL_OUT)
self._stage = 4 self._stage = STAGE_REPLY
self.update_stream(STREAM_UP, STATUS_WAIT_WRITING) self.update_stream(STREAM_UP, STATUS_WAIT_READWRITING)
self.update_stream(STREAM_DOWN, STATUS_WAIT_READING) self.update_stream(STREAM_DOWN, STATUS_WAIT_READING)
return return
except Exception: except Exception:
@ -221,9 +234,6 @@ class TCPRelayHandler(object):
# TODO use logging when debug completed # TODO use logging when debug completed
self.destroy() self.destroy()
elif self._stage == STAGE_REPLY:
self._data_to_write_to_remote.append(data)
def on_remote_read(self): def on_remote_read(self):
# TODO update timeout # TODO update timeout
data = None data = None
@ -334,6 +344,7 @@ class TCPRelay(object):
def _run(self): def _run(self):
server_socket = self._server_socket server_socket = self._server_socket
self._eventloop = eventloop.EventLoop() self._eventloop = eventloop.EventLoop()
logging.debug('using event model: %s', self._eventloop.model)
self._eventloop.add(server_socket, self._eventloop.add(server_socket,
eventloop.POLL_IN | eventloop.POLL_ERR) eventloop.POLL_IN | eventloop.POLL_ERR)
last_time = time.time() last_time = time.time()
@ -348,7 +359,9 @@ class TCPRelay(object):
logging.error(e) logging.error(e)
continue continue
for sock, event in events: for sock, event in events:
logging.debug('%s %d', sock, event) if sock:
logging.debug('fd %d %s', sock.fileno(),
eventloop.EVENT_NAMES[event])
if sock == self._server_socket: if sock == self._server_socket:
if event & eventloop.POLL_ERR: if event & eventloop.POLL_ERR:
# TODO # TODO
@ -364,13 +377,16 @@ class TCPRelay(object):
else: else:
logging.error(e) logging.error(e)
else: else:
handler = self._fd_to_handlers.get(sock.fileno(), None) if sock:
if handler: handler = self._fd_to_handlers.get(sock.fileno(), None)
handler.handle_event(sock, event) if handler:
handler.handle_event(sock, event)
else:
logging.warn('can not find handler for fd %d',
sock.fileno())
self._eventloop.remove(sock)
else: else:
logging.warn('can not find handler for fd %d', logging.warn('poll removed fd')
sock.fileno())
self._eventloop.remove(sock)
now = time.time() now = time.time()
if now - last_time > 5: if now - last_time > 5:
# TODO sweep timeouts # TODO sweep timeouts