wehub-resource-sync
728 lines
24 KiB
Python
728 lines
24 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
# ruff: noqa: F841
|
|
"""RPC proxy, allows both client/server to connect and match connection.
|
|
|
|
In normal RPC, client directly connect to server's IP address.
|
|
Sometimes this cannot be done when server do not have a static address.
|
|
RPCProxy allows both client and server connect to the proxy server,
|
|
the proxy server will forward the message between the client and server.
|
|
"""
|
|
|
|
# pylint: disable=unused-variable, unused-argument
|
|
import asyncio
|
|
import errno
|
|
import logging
|
|
import os
|
|
import socket
|
|
import struct
|
|
import threading
|
|
import time
|
|
|
|
try:
|
|
import tornado
|
|
from tornado import gen, ioloop, websocket
|
|
|
|
from . import tornado_util
|
|
except ImportError as error_msg:
|
|
raise ImportError(
|
|
f"RPCProxy module requires tornado package {error_msg}. Try 'pip install tornado'."
|
|
)
|
|
|
|
from tvm.support.popen_pool import PopenWorker
|
|
|
|
from . import _ffi_api, base
|
|
from .base import TrackerCode
|
|
from .server import _server_env
|
|
|
|
|
|
class ForwardHandler:
|
|
"""Forward handler to forward the message."""
|
|
|
|
def _init_handler(self):
|
|
"""Initialize handler."""
|
|
self._init_message = b""
|
|
self._init_req_nbytes = 4
|
|
self._magic = None
|
|
self.timeout = None
|
|
self._rpc_key_length = None
|
|
self._done = False
|
|
self._proxy = ProxyServerHandler.current
|
|
assert self._proxy
|
|
self.rpc_key = None
|
|
self.match_key = None
|
|
self.forward_proxy = None
|
|
self.alloc_time = None
|
|
|
|
def __del__(self):
|
|
logging.info("Delete %s...", self.name())
|
|
|
|
def name(self):
|
|
"""Name of this connection."""
|
|
return "RPCConnection"
|
|
|
|
def _init_step(self, message):
|
|
if self._magic is None:
|
|
assert len(message) == 4
|
|
self._magic = struct.unpack("<i", message)[0]
|
|
if self._magic != base.RPC_MAGIC:
|
|
logging.info("Invalid RPC magic from %s", self.name())
|
|
self.close()
|
|
self._init_req_nbytes = 4
|
|
elif self._rpc_key_length is None:
|
|
assert len(message) == 4
|
|
self._rpc_key_length = struct.unpack("<i", message)[0]
|
|
self._init_req_nbytes = self._rpc_key_length
|
|
elif self.rpc_key is None:
|
|
assert len(message) == self._rpc_key_length
|
|
self.rpc_key = message.decode("utf-8")
|
|
# match key is used to do the matching
|
|
self.match_key = self.rpc_key[7:].split()[0]
|
|
self.on_start()
|
|
else:
|
|
assert False
|
|
|
|
def on_start(self):
|
|
"""Event when the initialization is completed"""
|
|
self._proxy.handler_ready(self)
|
|
|
|
def on_data(self, message):
|
|
"""on data"""
|
|
assert isinstance(message, bytes)
|
|
if self.forward_proxy:
|
|
self.forward_proxy.send_data(message)
|
|
else:
|
|
while message and self._init_req_nbytes > len(self._init_message):
|
|
nbytes = self._init_req_nbytes - len(self._init_message)
|
|
self._init_message += message[:nbytes]
|
|
message = message[nbytes:]
|
|
if self._init_req_nbytes == len(self._init_message):
|
|
temp = self._init_message
|
|
self._init_req_nbytes = 0
|
|
self._init_message = b""
|
|
self._init_step(temp)
|
|
if message:
|
|
logging.info("Invalid RPC protocol, too many bytes %s", self.name())
|
|
self.close()
|
|
|
|
def on_error(self, err):
|
|
logging.info("%s: Error in RPC %s", self.name(), err)
|
|
self.close_pair()
|
|
|
|
def close_pair(self):
|
|
if self.forward_proxy:
|
|
self.forward_proxy.signal_close()
|
|
self.forward_proxy = None
|
|
self.close()
|
|
|
|
def on_close_event(self):
|
|
"""on close event"""
|
|
assert not self._done
|
|
logging.info("RPCProxy:on_close_event %s ...", self.name())
|
|
if self.match_key:
|
|
key = self.match_key
|
|
if self._proxy._client_pool.get(key, None) == self:
|
|
self._proxy._client_pool.pop(key)
|
|
if self._proxy._server_pool.get(key, None) == self:
|
|
self._proxy._server_pool.pop(key)
|
|
self._done = True
|
|
self.forward_proxy = None
|
|
|
|
|
|
class TCPHandler(tornado_util.TCPHandler, ForwardHandler):
|
|
"""Event driven TCP handler."""
|
|
|
|
def __init__(self, sock, addr):
|
|
super().__init__(sock)
|
|
self._init_handler()
|
|
self.addr = addr
|
|
|
|
def name(self):
|
|
return f"TCPSocketProxy:{self.addr[0]!s}:{self.rpc_key}"
|
|
|
|
def send_data(self, message, binary=True):
|
|
self.write_message(message, True)
|
|
|
|
def on_message(self, message):
|
|
self.on_data(message)
|
|
|
|
def on_close(self):
|
|
logging.info("RPCProxy: on_close %s ...", self.name())
|
|
self._close_process = True
|
|
|
|
if self.forward_proxy:
|
|
self.forward_proxy.signal_close()
|
|
self.forward_proxy = None
|
|
self.on_close_event()
|
|
|
|
|
|
class WebSocketHandler(websocket.WebSocketHandler, ForwardHandler):
|
|
"""Handler for websockets."""
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self._init_handler()
|
|
|
|
def name(self):
|
|
return f"WebSocketProxy:{self.rpc_key}"
|
|
|
|
def on_message(self, message):
|
|
self.on_data(message)
|
|
|
|
def data_received(self, _):
|
|
raise NotImplementedError()
|
|
|
|
def send_data(self, message):
|
|
try:
|
|
self.write_message(message, True)
|
|
except websocket.WebSocketClosedError as err:
|
|
self.on_error(err)
|
|
|
|
def on_close(self):
|
|
logging.info("RPCProxy: on_close %s ...", self.name())
|
|
if self.forward_proxy:
|
|
self.forward_proxy.signal_close()
|
|
self.forward_proxy = None
|
|
self.on_close_event()
|
|
|
|
def signal_close(self):
|
|
self.close()
|
|
|
|
|
|
MIME_MAP = {
|
|
"js": "text/javascript",
|
|
"css": "text/css",
|
|
"wasm": "application/wasm",
|
|
"json": "application/json",
|
|
}
|
|
|
|
|
|
class RequestHandler(tornado.web.RequestHandler):
|
|
"""Handles html request."""
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
file_path = kwargs.pop("file_path")
|
|
self.format = file_path.split(".")[-1]
|
|
|
|
if file_path.endswith("html"):
|
|
self.page = open(file_path).read()
|
|
web_port = kwargs.pop("rpc_web_port", None)
|
|
if web_port:
|
|
self.page = self.page.replace(
|
|
"ws://localhost:9190/ws", f"ws://localhost:{web_port}/ws"
|
|
)
|
|
else:
|
|
self.page = open(file_path, "rb").read()
|
|
|
|
super().__init__(*args, **kwargs)
|
|
|
|
def data_received(self, _):
|
|
pass
|
|
|
|
def get(self, *args, **kwargs):
|
|
if self.format in MIME_MAP:
|
|
self.set_header("Content-Type", MIME_MAP[self.format])
|
|
self.write(self.page)
|
|
|
|
|
|
class ProxyServerHandler:
|
|
"""Internal proxy server handler class."""
|
|
|
|
current = None
|
|
|
|
def __init__(
|
|
self,
|
|
sock,
|
|
listen_port,
|
|
web_port,
|
|
timeout_client,
|
|
timeout_server,
|
|
tracker_addr,
|
|
index_page=None,
|
|
resource_files=None,
|
|
):
|
|
assert ProxyServerHandler.current is None
|
|
ProxyServerHandler.current = self
|
|
if web_port:
|
|
handlers = [(r"/ws", WebSocketHandler)]
|
|
if index_page:
|
|
handlers.append(
|
|
(r"/", RequestHandler, {"file_path": index_page, "rpc_web_port": web_port})
|
|
)
|
|
logging.info("Serving RPC index html page at http://localhost:%d", web_port)
|
|
resource_files = resource_files if resource_files else []
|
|
for item in resource_files:
|
|
prefix, fname = item
|
|
if not prefix.endswith("/"):
|
|
prefix += "/"
|
|
if not prefix.startswith("/"):
|
|
prefix = "/" + prefix
|
|
basename = os.path.basename(fname)
|
|
pair = (rf"{prefix}{basename}", RequestHandler, {"file_path": fname})
|
|
handlers.append(pair)
|
|
logging.info(pair)
|
|
self.app = tornado.web.Application(handlers)
|
|
self.app.listen(web_port)
|
|
|
|
self.sock = sock
|
|
self.sock.setblocking(0)
|
|
self.loop = ioloop.IOLoop.current()
|
|
|
|
def event_handler(_, events):
|
|
self._on_event(events)
|
|
|
|
self.loop.add_handler(self.sock.fileno(), event_handler, self.loop.READ)
|
|
self._client_pool = {}
|
|
self._server_pool = {}
|
|
self.timeout_alloc = 5
|
|
self.timeout_client = timeout_client
|
|
self.timeout_server = timeout_server
|
|
# tracker information
|
|
self._listen_port = listen_port
|
|
self._tracker_addr = tracker_addr
|
|
self._tracker_conn = None
|
|
self._tracker_pending_puts = []
|
|
self._key_set = set()
|
|
self.update_tracker_period = 2
|
|
if tracker_addr:
|
|
logging.info("Tracker address:%s", str(tracker_addr))
|
|
|
|
def _callback():
|
|
self._update_tracker(True)
|
|
|
|
self.loop.call_later(self.update_tracker_period, _callback)
|
|
logging.info("RPCProxy: Websock port bind to %d", web_port)
|
|
|
|
def _on_event(self, _):
|
|
while True:
|
|
try:
|
|
conn, addr = self.sock.accept()
|
|
TCPHandler(conn, addr)
|
|
except OSError as err:
|
|
if err.args[0] in (errno.EAGAIN, errno.EWOULDBLOCK):
|
|
break
|
|
|
|
def _pair_up(self, lhs, rhs):
|
|
lhs.forward_proxy = rhs
|
|
rhs.forward_proxy = lhs
|
|
|
|
lhs.send_data(struct.pack("<i", base.RPC_CODE_SUCCESS))
|
|
lhs.send_data(struct.pack("<i", len(rhs.rpc_key)))
|
|
lhs.send_data(rhs.rpc_key.encode("utf-8"))
|
|
|
|
rhs.send_data(struct.pack("<i", base.RPC_CODE_SUCCESS))
|
|
rhs.send_data(struct.pack("<i", len(lhs.rpc_key)))
|
|
rhs.send_data(lhs.rpc_key.encode("utf-8"))
|
|
logging.info("Pairup connect %s and %s", lhs.name(), rhs.name())
|
|
|
|
def _regenerate_server_keys(self, keys):
|
|
"""Regenerate keys for server pool"""
|
|
keyset = set(self._server_pool.keys())
|
|
new_keys = []
|
|
# re-generate the server match key, so old information is invalidated.
|
|
for key in tuple(keys):
|
|
rpc_key, _ = base.split_random_key(key)
|
|
handle = self._server_pool[key]
|
|
del self._server_pool[key]
|
|
new_key = base.random_key(rpc_key, keyset)
|
|
self._server_pool[new_key] = handle
|
|
keyset.add(new_key)
|
|
new_keys.append(new_key)
|
|
return new_keys
|
|
|
|
def _update_tracker(self, period_update=False):
|
|
"""Update information on tracker."""
|
|
try:
|
|
if self._tracker_conn is None:
|
|
self._tracker_conn = socket.socket(
|
|
base.get_addr_family(self._tracker_addr), socket.SOCK_STREAM
|
|
)
|
|
self._tracker_conn.connect(self._tracker_addr)
|
|
self._tracker_conn.sendall(struct.pack("<i", base.RPC_TRACKER_MAGIC))
|
|
magic = struct.unpack("<i", base.recvall(self._tracker_conn, 4))[0]
|
|
if magic != base.RPC_TRACKER_MAGIC:
|
|
self.loop.stop()
|
|
raise RuntimeError(f"{self._tracker_addr} is not RPC Tracker")
|
|
# just connect to tracker, need to update all keys
|
|
self._tracker_pending_puts = self._server_pool.keys()
|
|
|
|
if self._tracker_conn and period_update:
|
|
# periodically update tracker information
|
|
# regenerate key if the key is not in tracker anymore
|
|
# and there is no in-coming connection after timeout_alloc
|
|
base.sendjson(self._tracker_conn, [TrackerCode.GET_PENDING_MATCHKEYS])
|
|
pending_keys = set(base.recvjson(self._tracker_conn))
|
|
update_keys = []
|
|
for k, v in self._server_pool.items():
|
|
if k not in pending_keys:
|
|
if v.alloc_time is None:
|
|
v.alloc_time = time.time()
|
|
elif time.time() - v.alloc_time > self.timeout_alloc:
|
|
update_keys.append(k)
|
|
v.alloc_time = None
|
|
if update_keys:
|
|
logging.info(
|
|
"RPCProxy: No incoming conn on %s, regenerate keys...", str(update_keys)
|
|
)
|
|
new_keys = self._regenerate_server_keys(update_keys)
|
|
self._tracker_pending_puts += new_keys
|
|
|
|
need_update_info = False
|
|
# report new connections
|
|
for key in self._tracker_pending_puts:
|
|
rpc_key, _ = base.split_random_key(key)
|
|
base.sendjson(
|
|
self._tracker_conn, [TrackerCode.PUT, rpc_key, (self._listen_port, key), None]
|
|
)
|
|
assert base.recvjson(self._tracker_conn) == TrackerCode.SUCCESS
|
|
if rpc_key not in self._key_set:
|
|
self._key_set.add(rpc_key)
|
|
need_update_info = True
|
|
|
|
if need_update_info:
|
|
keylist = "[" + ",".join(self._key_set) + "]"
|
|
cinfo = {"key": "server:proxy" + keylist, "addr": [None, self._listen_port]}
|
|
base.sendjson(self._tracker_conn, [TrackerCode.UPDATE_INFO, cinfo])
|
|
assert base.recvjson(self._tracker_conn) == TrackerCode.SUCCESS
|
|
self._tracker_pending_puts = []
|
|
except OSError as err:
|
|
logging.info(
|
|
"Lost tracker connection: %s, try reconnect in %g sec",
|
|
str(err),
|
|
self.update_tracker_period,
|
|
)
|
|
self._tracker_conn.close()
|
|
self._tracker_conn = None
|
|
self._regenerate_server_keys(self._server_pool.keys())
|
|
|
|
if period_update:
|
|
|
|
def _callback():
|
|
self._update_tracker(True)
|
|
|
|
self.loop.call_later(self.update_tracker_period, _callback)
|
|
|
|
def _handler_ready_tracker_mode(self, handler):
|
|
"""tracker mode to handle handler ready."""
|
|
if handler.rpc_key.startswith("server:"):
|
|
key = base.random_key(handler.match_key, cmap=self._server_pool)
|
|
handler.match_key = key
|
|
self._server_pool[key] = handler
|
|
self._tracker_pending_puts.append(key)
|
|
self._update_tracker()
|
|
else:
|
|
if handler.match_key in self._server_pool:
|
|
self._pair_up(self._server_pool.pop(handler.match_key), handler)
|
|
else:
|
|
handler.send_data(struct.pack("<i", base.RPC_CODE_MISMATCH))
|
|
handler.signal_close()
|
|
|
|
def _handler_ready_proxy_mode(self, handler):
|
|
"""Normal proxy mode when handler is ready."""
|
|
if handler.rpc_key.startswith("server:"):
|
|
pool_src, pool_dst = self._client_pool, self._server_pool
|
|
timeout = self.timeout_server
|
|
else:
|
|
pool_src, pool_dst = self._server_pool, self._client_pool
|
|
timeout = self.timeout_client
|
|
|
|
key = handler.match_key
|
|
if key in pool_src:
|
|
self._pair_up(pool_src.pop(key), handler)
|
|
return
|
|
if key not in pool_dst:
|
|
pool_dst[key] = handler
|
|
|
|
def cleanup():
|
|
"""Cleanup client connection if timeout"""
|
|
if pool_dst.get(key, None) == handler:
|
|
logging.info(
|
|
"Timeout client connection %s, cannot find match key=%s",
|
|
handler.name(),
|
|
key,
|
|
)
|
|
pool_dst.pop(key)
|
|
handler.send_data(struct.pack("<i", base.RPC_CODE_MISMATCH))
|
|
handler.signal_close()
|
|
|
|
self.loop.call_later(timeout, cleanup)
|
|
else:
|
|
logging.info("Duplicate connection with same key=%s", key)
|
|
handler.send_data(struct.pack("<i", base.RPC_CODE_DUPLICATE))
|
|
handler.signal_close()
|
|
|
|
def handler_ready(self, handler):
|
|
"""Report handler to be ready."""
|
|
logging.info("Handler ready %s", handler.name())
|
|
if self._tracker_addr:
|
|
self._handler_ready_tracker_mode(handler)
|
|
else:
|
|
self._handler_ready_proxy_mode(handler)
|
|
|
|
def run(self):
|
|
"""Run the proxy server"""
|
|
ioloop.IOLoop.current().start()
|
|
|
|
|
|
def _proxy_server(
|
|
listen_sock,
|
|
listen_port,
|
|
web_port,
|
|
timeout_client,
|
|
timeout_server,
|
|
tracker_addr,
|
|
index_page,
|
|
resource_files,
|
|
):
|
|
asyncio.set_event_loop(asyncio.new_event_loop())
|
|
handler = ProxyServerHandler(
|
|
listen_sock,
|
|
listen_port,
|
|
web_port,
|
|
timeout_client,
|
|
timeout_server,
|
|
tracker_addr,
|
|
index_page,
|
|
resource_files,
|
|
)
|
|
handler.run()
|
|
|
|
|
|
class PopenProxyServerState:
|
|
"""Internal PopenProxy State for Popen"""
|
|
|
|
current = None
|
|
|
|
def __init__(
|
|
self,
|
|
host,
|
|
port=9091,
|
|
port_end=9199,
|
|
web_port=0,
|
|
timeout_client=600,
|
|
timeout_server=600,
|
|
tracker_addr=None,
|
|
index_page=None,
|
|
resource_files=None,
|
|
):
|
|
sock = socket.socket(base.get_addr_family((host, port)), socket.SOCK_STREAM)
|
|
self.port = None
|
|
for my_port in range(port, port_end):
|
|
try:
|
|
sock.bind((host, my_port))
|
|
self.port = my_port
|
|
break
|
|
except OSError as sock_err:
|
|
if sock_err.errno in [errno.EADDRINUSE]:
|
|
continue
|
|
raise sock_err
|
|
if not self.port:
|
|
raise ValueError(f"cannot bind to any port in [{port}, {port_end})")
|
|
logging.info("RPCProxy: client port bind to %s:%d", host, self.port)
|
|
sock.listen(1)
|
|
self.thread = threading.Thread(
|
|
target=_proxy_server,
|
|
args=(
|
|
sock,
|
|
self.port,
|
|
web_port,
|
|
timeout_client,
|
|
timeout_server,
|
|
tracker_addr,
|
|
index_page,
|
|
resource_files,
|
|
),
|
|
)
|
|
# start the server in a different thread
|
|
# so we can return the port directly
|
|
self.thread.start()
|
|
|
|
|
|
def _popen_start_proxy_server(
|
|
host,
|
|
port=9091,
|
|
port_end=9199,
|
|
web_port=0,
|
|
timeout_client=600,
|
|
timeout_server=600,
|
|
tracker_addr=None,
|
|
index_page=None,
|
|
resource_files=None,
|
|
):
|
|
# This is a function that will be sent to the
|
|
# Popen worker to run on a separate process.
|
|
# Create and start the server in a different thread
|
|
state = PopenProxyServerState(
|
|
host,
|
|
port,
|
|
port_end,
|
|
web_port,
|
|
timeout_client,
|
|
timeout_server,
|
|
tracker_addr,
|
|
index_page,
|
|
resource_files,
|
|
)
|
|
PopenProxyServerState.current = state
|
|
# returns the port so that the main can get the port number.
|
|
return state.port
|
|
|
|
|
|
class Proxy:
|
|
"""Start RPC proxy server on a separate process.
|
|
|
|
Python implementation based on PopenWorker.
|
|
|
|
Parameters
|
|
----------
|
|
host : str
|
|
The host url of the server.
|
|
|
|
port : int
|
|
The TCP port to be bind to
|
|
|
|
port_end : int, optional
|
|
The end TCP port to search
|
|
|
|
web_port : int, optional
|
|
The http/websocket port of the server.
|
|
|
|
timeout_client : float, optional
|
|
Timeout of client until it sees a matching connection.
|
|
|
|
timeout_server : float, optional
|
|
Timeout of server until it sees a matching connection.
|
|
|
|
tracker_addr: Tuple (str, int) , optional
|
|
The address of RPC Tracker in tuple (host, ip) format.
|
|
If is not None, the server will register itself to the tracker.
|
|
|
|
index_page : str, optional
|
|
Path to an index page that can be used to display at proxy index.
|
|
|
|
resource_files : str, optional
|
|
Path to local resources that can be included in the http request
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
host,
|
|
port=9091,
|
|
port_end=9199,
|
|
web_port=0,
|
|
timeout_client=600,
|
|
timeout_server=600,
|
|
tracker_addr=None,
|
|
index_page=None,
|
|
resource_files=None,
|
|
):
|
|
self.proc = PopenWorker()
|
|
# send the function
|
|
self.proc.send(
|
|
_popen_start_proxy_server,
|
|
[
|
|
host,
|
|
port,
|
|
port_end,
|
|
web_port,
|
|
timeout_client,
|
|
timeout_server,
|
|
tracker_addr,
|
|
index_page,
|
|
resource_files,
|
|
],
|
|
)
|
|
# receive the port
|
|
self.port = self.proc.recv()
|
|
self.host = host
|
|
|
|
def terminate(self):
|
|
"""Terminate the server process"""
|
|
if self.proc:
|
|
logging.info("Terminating Proxy Server...")
|
|
self.proc.kill()
|
|
self.proc = None
|
|
|
|
def __del__(self):
|
|
try:
|
|
self.terminate()
|
|
except ImportError:
|
|
pass
|
|
|
|
|
|
def websocket_proxy_server(url, key=""):
|
|
"""Create a RPC server that uses an websocket that connects to a proxy.
|
|
|
|
Parameters
|
|
----------
|
|
url : str
|
|
The url to be connected.
|
|
|
|
key : str
|
|
The key to identify the server.
|
|
"""
|
|
|
|
def create_on_message(conn):
|
|
def _fsend(data):
|
|
data = bytes(data)
|
|
conn.write_message(data, binary=True)
|
|
return len(data)
|
|
|
|
on_message = _ffi_api.CreateEventDrivenServer(_fsend, "WebSocketProxyServer", "%toinit")
|
|
return on_message
|
|
|
|
@gen.coroutine
|
|
def _connect(key):
|
|
conn = yield websocket.websocket_connect(url)
|
|
on_message = create_on_message(conn)
|
|
temp = _server_env(None)
|
|
# Start connecton
|
|
conn.write_message(struct.pack("<i", base.RPC_MAGIC), binary=True)
|
|
key = "server:" + key
|
|
conn.write_message(struct.pack("<i", len(key)), binary=True)
|
|
conn.write_message(key.encode("utf-8"), binary=True)
|
|
msg = yield conn.read_message()
|
|
assert len(msg) >= 4
|
|
magic = struct.unpack("<i", msg[:4])[0]
|
|
if magic == base.RPC_CODE_DUPLICATE:
|
|
raise RuntimeError(f"key: {key} has already been used in proxy")
|
|
if magic == base.RPC_CODE_MISMATCH:
|
|
logging.info("RPCProxy do not have matching client key %s", key)
|
|
elif magic != base.RPC_CODE_SUCCESS:
|
|
raise RuntimeError(f"{url} is not RPC Proxy")
|
|
msg = msg[4:]
|
|
|
|
logging.info("Connection established with remote")
|
|
|
|
if msg:
|
|
on_message(bytearray(msg), 3)
|
|
|
|
while True:
|
|
try:
|
|
msg = yield conn.read_message()
|
|
if msg is None:
|
|
break
|
|
on_message(bytearray(msg), 3)
|
|
except websocket.WebSocketClosedError as err:
|
|
break
|
|
logging.info("WebSocketProxyServer closed...")
|
|
temp.remove()
|
|
ioloop.IOLoop.current().stop()
|
|
|
|
ioloop.IOLoop.current().spawn_callback(_connect, key)
|
|
ioloop.IOLoop.current().start()
|