def main() -> None:
args = parse_args()
config_path = args.config
delattr(args, 'config')
config = Config()
if config_path is not None:
config.load_file(config_path)
config.from_object(args)
try:
config.validate([
'tunnel',
'dest',
'secretkey',
'max_spare_count',
])
except ConfigMissing as e:
logger.error(e)
exit()
logger.setLevel(name2level(args.level))
Protocol.set_secret_key(args.secretkey)
Protocol.recalc_crc32()
pkgbuilder = PKGBuilder(Protocol)
local_ = Local(pkgbuilder, config)
logger.debug('PID: {}'.format(os.getpid()))
logger.info('init successful, running as local')
local_.run_forever()
def transfer_from_dest(self, r_conn: socket_t, mask: int,
buf: List[deque]) -> None:
"""receive data from dest and store in buffer"""
w_conn = self.working_pool.inv.get(r_conn)
if w_conn is None:
self._sel.unregister(r_conn)
r_conn.close()
return
data = b''
need_close = False
try:
data = r_conn.recv(BUF_SIZE)
except ConnectionError:
need_close = True
if data == b'' or need_close:
try:
peer = r_conn.getpeername()
logger.info(
f'closing dest connection from {format_addr(peer)}')
except OSError as e:
logger.warn(e)
self._sel.unregister(r_conn)
self._sel.modify(w_conn, selectors.EVENT_WRITE,
partial(self.send_to_tunnel, buf=buf))
r_conn.close()
buf[NEG].append(sentinel)
del self.working_pool.inv[r_conn]
return
buf[NEG].append(data)
def transfer_from_tunnel(self, r_conn: socket_t, mask: int,
buf: List[deque]) -> None:
"""receive data from tunnel and store in buffer"""
w_conn = self.work_pool.inv.get(r_conn)
if w_conn is None:
self._sel.unregister(r_conn)
r_conn.close()
return
data = b''
need_close = False
try:
data = r_conn.recv(BUF_SIZE)
except ConnectionError:
need_close = True
if data == b'' or need_close:
try:
peer = r_conn.getpeername()
logger.info(
f'closing tunnel connection from {format_addr(peer)}'
) # noqa
except OSError as e:
logger.warn(e)
self._sel.unregister(r_conn)
r_conn.close()
buf[NEG].append(sentinel)
del self.work_pool.inv[r_conn]
return
buf[NEG].append(data)
def accept_expose(self, expose_sock: socket_t, mask: int) -> None:
"""accept user connection"""
conn, addr = expose_sock.accept()
conn.setblocking(False)
self._sel.register(conn, selectors.EVENT_READ, self.prepare_transfer)
logger.info(f'accept user connection from {format_addr(addr)}')
def accept_tunnel(self, tunnel_sock: socket_t, mask: int) -> None:
"""accept tunnel connection"""
conn, addr = tunnel_sock.accept()
conn.setblocking(False)
self.tunnel_pool.append(conn)
logger.info(f'accept tunnel connection from {format_addr(addr)}, '
f'poolsize is {len(self.tunnel_pool)}')
def run_forever(self) -> None:
while not self._stopping:
self.manage_tunnel()
events = self._sel.select(timeout=self._timeout)
for key, mask in events:
callback = key.data
callback(key.fileobj, mask)
logger.info('stopping now ...')
self.exit()
def _handshake(self, conn: socket_t) -> bool:
buff = conn.recv(self._pkgbuilder.PACKAGE_SIZE)
if (buff == b'' or not self._pkgbuilder.decode_verify(
buff, self._pkgbuilder.PTYPE_HS_M2S)):
logger.info('handshake failed')
conn.close()
return False
logger.debug('handshake successful')
conn.setblocking(True) # TODO
conn.send(self._pkgbuilder.pbuild_hs_s2m())
conn.setblocking(False)
return True
def send_to_tunnel(self, w_conn: socket_t, mask: int,
buf: List[deque]) -> None:
if not len(buf[NEG]):
return
try:
data = buf[NEG].popleft()
if data is sentinel:
self._sel.unregister(w_conn)
w_conn.close()
return
byte = w_conn.send(data)
except socket.error as e:
if e.args[0] == errno.EWOULDBLOCK:
logger.info('EWOULDBLOCK occur in send to tunnel')
buf[NEG].appendleft(data[byte:])
def run_forever(self) -> None:
"""main loop"""
while not self._stopping:
events = self._sel.select(timeout=self._timeout)
for job in self._ready: # TODO heartbeat
job()
self._ready.clear()
# from pprint import pprint
# pprint(events)
for key, mask in events:
callback = key.data
callback(key.fileobj, mask)
logger.info('stopping now ...')
self.exit()
def _connect_tunnel(self) -> bool:
"""establish tunnel connection"""
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
try:
sock.connect(self.tunnel_addr)
except ConnectionError:
self.next_timeout()
logger.warn(f'master is unreachable, '
f'retry after {self._timeout}sec...') # noqa
return False
self.reset_timeout()
self._sel.register(sock, selectors.EVENT_READ, self.prepare_transfer)
self.tunnel_pool.append(sock)
logger.info(f'connect to tunnel {format_addr(self.tunnel_addr)}, '
f'poolsize is {len(self.tunnel_pool)}')
return True
def find_available_tunnel(self) -> Optional[socket_t]:
while True:
try:
conn = self.tunnel_pool.popleft()
except IndexError:
# no available tunnel connection, just return
# do not need to wait in a loop, because we work in LT Mode
self.next_timeout()
logger.info('no available tunnel connection, waiting')
return None
else:
self.reset_timeout()
if not self._handshake(conn): # handshake first
conn.close()
continue
return conn
def _connect_dest(self) -> Optional[socket_t]:
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock.settimeout(2)
try:
sock.connect(self.dest_addr)
except socket.timeout:
logger.info('connect dest timeout')
return None
except ConnectionRefusedError:
logger.info('dest refused connection')
return None
logger.info('connected to dest {} at {}'.format(
format_addr(sock.getpeername()),
format_addr(sock.getsockname()),
))
return sock
def handle_signal(self, expose_sock: socket_t, mask: int) -> None:
sig = self._wake_fds[0].recv(1)
logger.info('recving signal {}'.format(sig))
self._stopping = True