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:
def prepare(self, request):
# self.app = self.settings["app"]
child, parent = socket.socketpair()
if request.headers.getlist("X-Forwarded-For"):
remote_ip = request.headers.getlist("X-Forwarded-For")[0]
else:
......@@ -36,10 +35,6 @@ class BaseWebSocketHandler:
self.clients[request.sid]["request"].meta = {"width": self.clients[request.sid]["cols"],
"height": self.clients[request.sid]["rows"]}
# 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):
return True
......@@ -64,9 +59,10 @@ class SSHws(Namespace, BaseWebSocketHandler):
"cols": int(request.cookies.get('cols', 80)),
"rows": int(request.cookies.get('rows', 24)),
"room": room,
"chan": None,
"proxy": None,
"client": None,
# "chan": dict(),
"proxy": dict(),
"client": dict(),
"forwarder": dict(),
"request": None,
}
self.rooms[room] = {
......@@ -80,18 +76,31 @@ class SSHws(Namespace, BaseWebSocketHandler):
self.prepare(request)
def on_data(self, message):
if self.clients[request.sid]["proxy"]:
self.clients[request.sid]["proxy"].send({"data": message})
if message['room'] and self.clients[request.sid]["proxy"][message['room']]:
self.clients[request.sid]["proxy"][message['room']].send({"data": message['data']})
def on_host(self, message):
# 此处获取主机的信息
uuid = message.get('uuid', None)
connection = str(uuid.uuid4())
assetID = message.get('uuid', None)
userid = message.get('userid', None)
if uuid and userid:
asset = self.app.service.get_asset(uuid)
self.emit('room', {'room': connection, 'secret': message['secret']})
if assetID and userid:
asset = self.app.service.get_asset(assetID)
system_user = self.app.service.get_system_user(userid)
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)
else:
self.on_disconnect()
......@@ -125,13 +134,21 @@ class SSHws(Namespace, BaseWebSocketHandler):
def on_disconnect(self):
self.on_leave(self.clients[request.sid]["room"])
try:
# todo: there maybe have bug
self.clients[request.sid]["proxy"].close()
del self.clients[request.sid]
except:
pass
# self.ssh.close()
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:
# prepare may be rewrite it
......
......@@ -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 child: sock child pair
......@@ -196,6 +196,7 @@ class WSProxy:
self.stop_event = threading.Event()
self.room = room
self.auto_forward()
self.connection = connection
def send(self, msg):
"""
......@@ -215,7 +216,7 @@ class WSProxy:
data = self.child.recv(BUF_SIZE)
if len(data) == 0:
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):
thread = threading.Thread(target=self.forward, args=())
......@@ -225,4 +226,3 @@ class WSProxy:
def close(self):
self.stop_event.set()
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