From 4f973844556bc6d07881198bad4300427042d2a5 Mon Sep 17 00:00:00 2001
From: yumoqing <yumoqing@gmail.com>
Date: Wed, 9 Oct 2024 11:46:30 +0800
Subject: [PATCH] bugfix

---
 ahserver/websocketProcessor.py | 82 ++++++++++++++++++++++++++--------
 1 file changed, 64 insertions(+), 18 deletions(-)

diff --git a/ahserver/websocketProcessor.py b/ahserver/websocketProcessor.py
index ce47a47..ecf54da 100755
--- a/ahserver/websocketProcessor.py
+++ b/ahserver/websocketProcessor.py
@@ -96,29 +96,67 @@ async def ws_send(ws:web.WebSocketResponse, data):
 		print_exc()
 		return False
 
+class WsSession:
+	def __init__(self, session)
+		self.session = session
+		self.nodes = {}
+	
+	def join(node):
+		self.nodes[node.id] = node
+	
+	def leave(node):
+		self.nodes = {k:v for k,v in self.nodes.items() if k != node.id}
+	
+class WsData:
+	def __init__(self):
+		self.nodes = {}
+		self.sessions = {}
+	
+	def add_node(self, node):
+		self.nodes[node.id] = node
+	
+	def del_node(self, node):
+		self.nodes = {k,v for k,v in self.nodes.items() if k!=node.id}
+	
+	def get_node(self, id):
+		return self.nodes.get(id)
+	
+	def add_session(self, session):
+		self.sessions[session.sessionid] = session
+	
+	def del_session(self, session):
+		self.sessions = {k:v for k,v in self.sessions.items() if k != session.sessionid}
+	
+	def get_session(self, id):
+		return self.sessions.get(id)
+
 class WsPool:
-	def __init__(self, ws, ws_path, app):
+	def __init__(self, ws, ip, ws_path, app):
 		self.app = app
+		self.ip = ip
 		self.id = None
 		self.ws = ws
 		self.ws_path = ws_path
 
 	def get_data(self):
-		return self.app.get_data(self.ws_path)
+		r = self.app.get_data(self.ws_path)
+		if r is None:
+			r = WsData()
+			self.set_data(r)
+		return r
 
 	def set_data(self, data):
 		self.app.set_data(self.ws_path, data)
 
 	def is_online(self, userid):
 		data = self.get_data()
-		ws = data.get(userid)
-		if ws is None:
+		node = data.get_node(userid)
+		if node is None:
 			return False
 		return True
 
 	def register(self, id):
 		iddata = DictObject()
-		iddata.ws = self.ws
 		iddata.id = id
 		self.add_me(iddata)
 
@@ -127,32 +165,40 @@ class WsPool:
 		if data is None:
 			data = {}
 		iddata.ws = self.ws
+		iddata.ip = self.ip
 		self.id = iddata.id
-		data.update({self.id:iddata})
-		info(f'add_me() {data=}')
+		data.add_node(ddata)
 		self.set_data(data)
 
 	def delete_id(self, id):
 		data = self.get_data()
-		if not data.get(id):
-			return
-		data = {k:v for k,v in data.items() if k != id }
+		node = data.get_node(id)
+		if node:
+			data.del_node(node)
 		self.set_data(data)
 
 	def delete_me(self):
 		self.delete_id(self.id)
-		data = self.get_data()
-		if not data.get(self.id):
-			return
-		data = {k:v for k,v in data.items() if k != self.id }
-		self.set_data(data)
 		
+	def add_session(self, session):
+		data = self.get_data()
+		data.add_session(session)
+		self.set_data(data)
+	
+	def del_session(self, session):
+		data = self.get_data()
+		data.del_session(session)
+		self.set_data(data)
+
+	def get_session(self, sessionid):
+		data = self.get_data()
+		return data.get_session(sessionid)
+
 	async def sendto(self, data, id=None):
 		if id is None:
 			return await ws_send(self.ws, data)
 		d = self.get_data()
-		iddata = d.get(id)
-		info(f'{id=} {d=}, {iddata=}')
+		iddata = d.get_node(id)
 		ws = iddata.ws
 		try:
 			return await ws_send(ws, data)
@@ -179,7 +225,7 @@ class WebsocketProcessor(PythonScriptProcessor):
 			print('--------except:', e, type(e))
 			print_exc()
 			raise e
-		ws_pool = WsPool(ws, request.path, request.app)
+		ws_pool = WsPool(ws, request['client_ip'], request.path, request.app)
 		async for msg in ws:
 			if msg.type == aiohttp.WSMsgType.TEXT:
 				if msg.data == 'exit':