feat: one ws, fit for luna

this is fix for browser can't hold more than 5 ws connections
parent 801a1e84
...@@ -26,7 +26,6 @@ class BaseWebSocketHandler: ...@@ -26,7 +26,6 @@ class BaseWebSocketHandler:
def prepare(self, request): def prepare(self, request):
# self.app = self.settings["app"] # self.app = self.settings["app"]
child, parent = socket.socketpair()
if request.headers.getlist("X-Forwarded-For"): if request.headers.getlist("X-Forwarded-For"):
remote_ip = request.headers.getlist("X-Forwarded-For")[0] remote_ip = request.headers.getlist("X-Forwarded-For")[0]
else: else:
...@@ -36,10 +35,6 @@ class BaseWebSocketHandler: ...@@ -36,10 +35,6 @@ class BaseWebSocketHandler:
self.clients[request.sid]["request"].meta = {"width": self.clients[request.sid]["cols"], self.clients[request.sid]["request"].meta = {"width": self.clients[request.sid]["cols"],
"height": self.clients[request.sid]["rows"]} "height": self.clients[request.sid]["rows"]}
# self.request.__dict__.update(request.__dict__) # self.request.__dict__.update(request.__dict__)
self.clients[request.sid]["client"] = Client(parent, self.clients[request.sid]["request"])
self.clients[request.sid]["proxy"] = WSProxy(self, child, self.clients[request.sid]["room"])
self.app.clients.append(self.clients[request.sid]["client"])
self.clients[request.sid]["forwarder"] = ProxyServer(self.app, self.clients[request.sid]["client"])
def check_origin(self, origin): def check_origin(self, origin):
return True return True
...@@ -64,9 +59,10 @@ class SSHws(Namespace, BaseWebSocketHandler): ...@@ -64,9 +59,10 @@ class SSHws(Namespace, BaseWebSocketHandler):
"cols": int(request.cookies.get('cols', 80)), "cols": int(request.cookies.get('cols', 80)),
"rows": int(request.cookies.get('rows', 24)), "rows": int(request.cookies.get('rows', 24)),
"room": room, "room": room,
"chan": None, # "chan": dict(),
"proxy": None, "proxy": dict(),
"client": None, "client": dict(),
"forwarder": dict(),
"request": None, "request": None,
} }
self.rooms[room] = { self.rooms[room] = {
...@@ -80,18 +76,31 @@ class SSHws(Namespace, BaseWebSocketHandler): ...@@ -80,18 +76,31 @@ class SSHws(Namespace, BaseWebSocketHandler):
self.prepare(request) self.prepare(request)
def on_data(self, message): def on_data(self, message):
if self.clients[request.sid]["proxy"]: if message['room'] and self.clients[request.sid]["proxy"][message['room']]:
self.clients[request.sid]["proxy"].send({"data": message}) self.clients[request.sid]["proxy"][message['room']].send({"data": message['data']})
def on_host(self, message): def on_host(self, message):
# 此处获取主机的信息 # 此处获取主机的信息
uuid = message.get('uuid', None) connection = str(uuid.uuid4())
assetID = message.get('uuid', None)
userid = message.get('userid', None) userid = message.get('userid', None)
if uuid and userid: self.emit('room', {'room': connection, 'secret': message['secret']})
asset = self.app.service.get_asset(uuid)
if assetID and userid:
asset = self.app.service.get_asset(assetID)
system_user = self.app.service.get_system_user(userid) system_user = self.app.service.get_system_user(userid)
if system_user: if system_user:
self.socketio.start_background_task(self.clients[request.sid]["forwarder"].proxy, asset, system_user)
child, parent = socket.socketpair()
self.clients[request.sid]["client"][connection] = Client(parent, self.clients[request.sid]["request"])
self.clients[request.sid]["proxy"][connection] = WSProxy(self, child, self.clients[request.sid]["room"],
connection)
self.app.clients.append(self.clients[request.sid]["client"][connection])
self.clients[request.sid]["forwarder"][connection] = ProxyServer(self.app,
self.clients[request.sid]["client"][connection])
self.socketio.start_background_task(self.clients[request.sid]["forwarder"][connection].proxy, asset,
system_user)
# self.forwarder.proxy(self.asset, system_user) # self.forwarder.proxy(self.asset, system_user)
else: else:
self.on_disconnect() self.on_disconnect()
...@@ -125,13 +134,21 @@ class SSHws(Namespace, BaseWebSocketHandler): ...@@ -125,13 +134,21 @@ class SSHws(Namespace, BaseWebSocketHandler):
def on_disconnect(self): def on_disconnect(self):
self.on_leave(self.clients[request.sid]["room"]) self.on_leave(self.clients[request.sid]["room"])
try: try:
# todo: there maybe have bug del self.clients[request.sid]
self.clients[request.sid]["proxy"].close()
except: except:
pass pass
# self.ssh.close() # self.ssh.close()
pass pass
def on_logout(self, connection):
print("logout", connection)
if connection:
self.clients[request.sid]["proxy"][connection].close()
del self.clients[request.sid]["proxy"][connection]
del self.clients[request.sid]["forwarder"][connection]
self.clients[request.sid]["client"][connection].close()
del self.clients[request.sid]["client"][connection]
class HttpServer: class HttpServer:
# prepare may be rewrite it # prepare may be rewrite it
......
...@@ -186,7 +186,7 @@ class WSProxy: ...@@ -186,7 +186,7 @@ class WSProxy:
``` ```
""" """
def __init__(self, ws, child, room): def __init__(self, ws, child, room, connection):
""" """
:param ws: websocket instance or handler, have write_message method :param ws: websocket instance or handler, have write_message method
:param child: sock child pair :param child: sock child pair
...@@ -196,6 +196,7 @@ class WSProxy: ...@@ -196,6 +196,7 @@ class WSProxy:
self.stop_event = threading.Event() self.stop_event = threading.Event()
self.room = room self.room = room
self.auto_forward() self.auto_forward()
self.connection = connection
def send(self, msg): def send(self, msg):
""" """
...@@ -215,7 +216,7 @@ class WSProxy: ...@@ -215,7 +216,7 @@ class WSProxy:
data = self.child.recv(BUF_SIZE) data = self.child.recv(BUF_SIZE)
if len(data) == 0: if len(data) == 0:
self.close() self.close()
self.ws.emit("data", data.decode("utf-8"), room=self.room) self.ws.emit("data", {'data': data.decode("utf-8"), 'room': self.connection}, room=self.room)
def auto_forward(self): def auto_forward(self):
thread = threading.Thread(target=self.forward, args=()) thread = threading.Thread(target=self.forward, args=())
...@@ -225,4 +226,3 @@ class WSProxy: ...@@ -225,4 +226,3 @@ class WSProxy:
def close(self): def close(self):
self.stop_event.set() self.stop_event.set()
self.child.close() self.child.close()
self.ws.on_disconnect()
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment