Allow an application to connect to websocket multiple times.

Author: EvieePy

api/routes/members.py

@@ -64,20 +64,23 @@ async def post_dpy_modlog(self, request: Request) -> Response:
         }
 
         count = 0
+        total = 0
         for subscriber in self.app.subscription_sockets[core.WebsocketSubscriptions.DPY_MOD_LOG]:
-            websocket: WebSocket = self.app.sockets[subscriber]
+            websockets: list[WebSocket] = list(self.app.sockets[subscriber].values())
 
+            total += len(websockets)
             payload['user_id'] = subscriber
 
-            try:
-                await websocket.send_json(data=payload)
-            except Exception as e:
-                logger.debug(f'Failed to send payload to websocket "{subscriber}": {e}')
-            else:
-                count += 1
+            for websocket in websockets:
+                try:
+                    await websocket.send_json(data=payload)
+                except Exception as e:
+                    logger.debug(f'Failed to send payload to a websocket for "{subscriber}": {e}')
+                else:
+                    count += 1
 
         to_send: dict[str, int] = {
-            'subscribers': len(self.app.subscription_sockets[core.WebsocketSubscriptions.DPY_MOD_LOG]),
+            'subscribers': total,
             'successful': count
         }
         return JSONResponse(to_send, status_code=200)

api/server.py

@@ -20,6 +20,7 @@
 OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 SOFTWARE.
 """
+import secrets
 from typing import Any
 
 import aiohttp
@@ -50,7 +51,7 @@ def __init__(self, *, session: aiohttp.ClientSession, database: core.Database) -
             Middleware(AuthenticationMiddleware, backend=AuthBackend(self)),
         ]
 
-        self.sockets: dict[int, WebSocket] = {}
+        self.sockets: dict[int, dict[str, WebSocket]] = {}
         self.subscription_sockets: dict[str, set[int]] = {
             core.WebsocketSubscriptions.DPY_MOD_LOG: set()
         }
@@ -71,7 +72,12 @@ async def websocket_connector(self, websocket: WebSocket) -> None:
         uid: int | None = core.id_from_token(token)
 
         assert uid
-        self.sockets[uid] = websocket
+
+        hash_ = secrets.token_urlsafe(8)
+        try:
+            self.sockets[uid][hash_] = websocket
+        except KeyError:
+            self.sockets[uid] = {hash_: websocket}
 
         # Filter out bad subscriptions...
         valid: list[str] = list(self.subscription_sockets.keys())
@@ -116,12 +122,8 @@ async def websocket_connector(self, websocket: WebSocket) -> None:
                 }
                 await websocket.send_json(data=response)
 
-        # Remove the websocket and it's subscriptions...
-        del self.sockets[uid]
-
-        subscribed: list[str] = [sub for sub in self.subscription_sockets if uid in self.subscription_sockets[sub]]
-        for sub in subscribed:
-            self.subscription_sockets[sub].remove(uid)
+        # Remove the websocket...
+        del self.sockets[uid][hash_]
 
     def websocket_subscribe(self, *, uid: int, message: dict[str, Any]) -> dict[str, Any]:
         subs: list[str] = message.get('subscriptions', [])
@@ -151,17 +153,24 @@ def websocket_unsubscribe(self, *, uid: int, message: dict[str, Any]) -> dict[st
         # Filter out bad subscriptions...
         valid: list[str] = list(self.subscription_sockets.keys())
         subscriptions: list[str] = [sub for sub in subs if sub in valid]
+        removed: list[str] = []
 
         for sub in subscriptions:
-            self.subscription_sockets[sub].remove(uid)
+
+            try:
+                self.subscription_sockets[sub].remove(uid)
+            except KeyError:
+                pass
+            else:
+                removed.append(sub)
 
         subscribed: list[str] = [sub for sub in self.subscription_sockets if uid in self.subscription_sockets[sub]]
 
         data: dict[str, Any] = {
             'op': core.WebsocketOPCodes.NOTIFICATION,
             'type': core.WebsocketNotificationTypes.SUBSCRIPTION_REMOVED,
             'user_id': uid,
-            'removed': subscriptions,
+            'removed': removed,
             'subscriptions': subscribed
         }