From e329c64eb7e9f905bffc303fd2ed94a67e7b0ea7 Mon Sep 17 00:00:00 2001
From: hippoz <10706925-hippoz@users.noreply.gitlab.com>
Date: Thu, 14 Apr 2022 21:52:42 +0300
Subject: [PATCH] backend/gateway: add session count limit per user id

---
 src/errors.ts        |  1 +
 src/gateway/index.ts | 27 ++++++++++++++++++++++++++-
 2 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/src/errors.ts b/src/errors.ts
index c2b8000..9964d96 100644
--- a/src/errors.ts
+++ b/src/errors.ts
@@ -15,4 +15,5 @@ export const gatewayErrors = {
     FLOODING: { code: 4005, message: "Flooding (exceeded maximum messages per batch)" },
     ALREADY_AUTHENTICATED: { code: 4006, message: "Already authenticated" },
     PAYLOAD_TOO_LARGE: { code: 4007, message: "Payload too large" },
+    TOO_MANY_SESSIONS: { code: 4008, message: "Too many sessions" },
 };
diff --git a/src/gateway/index.ts b/src/gateway/index.ts
index fcc1c3f..892ce7e 100644
--- a/src/gateway/index.ts
+++ b/src/gateway/index.ts
@@ -10,10 +10,14 @@ import { GatewayPayloadType } from "./gatewaypayloadtype";
 const GATEWAY_BATCH_INTERVAL = 50000;
 const GATEWAY_PING_INTERVAL = 40000;
 const MAX_CLIENT_MESSAGES_PER_BATCH = 6; // TODO: how well does this work for weak connections?
+const MAX_GATEWAY_SESSIONS_PER_USER = 5;
 
 // mapping between a dispatch id and a websocket client
 const dispatchChannels = new Map<string, Set<WebSocket>>();
 
+// mapping between a user id and the websocket sessions it has
+const sessionsByUserId = new Map<number, Set<WebSocket>>();
+
 function clientSubscribe(ws: WebSocket, dispatchChannel: string) {
     ws.state.dispatchChannels.add(dispatchChannel);
     if (!dispatchChannels.get(dispatchChannel)) {
@@ -159,6 +163,15 @@ export default function(server: Server) {
 
         ws.on("close", () => {
             clientUnsubscribeAll(ws);
+            if (ws.state.user && ws.state.user.id) {
+                const sessions = sessionsByUserId.get(ws.state.user.id);
+                if (sessions) {
+                    sessions.delete(ws);
+                    if (sessions.size < 1) {
+                        sessionsByUserId.delete(ws.state.user.id);
+                    }
+                }
+            }
         });
 
         ws.on("message", async (rawData, isBinary) => {
@@ -195,7 +208,19 @@ export default function(server: Server) {
                     if (!user) {
                         return closeWithError(ws, gatewayErrors.BAD_AUTH);
                     }
-                    // each user should have their own list of channels that they join
+
+                    let sessions = sessionsByUserId.get(user.id);
+                    if (sessions) {
+                        if ((sessions.size + 1) > MAX_GATEWAY_SESSIONS_PER_USER) {
+                            return closeWithError(ws, gatewayErrors.TOO_MANY_SESSIONS);
+                        }
+                    } else {
+                        sessions = new Set();
+                        sessionsByUserId.set(user.id, sessions);
+                    }
+                    sessions.add(ws);
+
+                    // TODO: each user should have their own list of channels that they join
                     const channels = await query("SELECT id, name, owner_id FROM channels");
 
                     clientSubscribe(ws, "*");