From 9c9f764e6dc2eb9c30b1fc14e5cd0f67d0375e3a Mon Sep 17 00:00:00 2001
From: hippoz <10706925-hippoz@users.noreply.gitlab.com>
Date: Thu, 14 Apr 2022 17:17:54 +0300
Subject: [PATCH] backend: add gateway ratelimit

---
 src/errors.ts                     |  1 +
 src/gateway/index.ts              | 10 +++++++++-
 src/types/gatewayclientstate.d.ts |  3 ++-
 3 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/src/errors.ts b/src/errors.ts
index 7bdc59d..1769d5b 100644
--- a/src/errors.ts
+++ b/src/errors.ts
@@ -12,4 +12,5 @@ export const gatewayErrors = {
     BAD_AUTH: { code: 4002, message: "Bad authentication" },
     AUTHENTICATION_TIMEOUT: { code: 4003, message: "Authentication timeout" },
     NO_PING: { code: 4004, message: "No ping" },
+    FLOODING: { code: 4005, message: "Flooding (exceeded maximum messages per batch)" },
 };
diff --git a/src/gateway/index.ts b/src/gateway/index.ts
index 112eeaa..8337dcc 100644
--- a/src/gateway/index.ts
+++ b/src/gateway/index.ts
@@ -9,6 +9,7 @@ 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?
 
 // mapping between a dispatch id and a websocket client
 const dispatchChannels = new Map<string, Set<WebSocket>>();
@@ -128,6 +129,7 @@ export default function(server: Server) {
                 if (!e.state.alive) {
                     return closeWithError(e, gatewayErrors.NO_PING);
                 }
+                e.state.messagesSinceLastCheck = 0;
             }
         });
     }, GATEWAY_BATCH_INTERVAL);
@@ -144,7 +146,8 @@ export default function(server: Server) {
             alive: false,
             ready: false,
             lastAliveCheck: performance.now(),
-            dispatchChannels: new Set()
+            dispatchChannels: new Set(),
+            messagesSinceLastCheck: 0
         };
 
         sendPayload(ws, {
@@ -163,6 +166,11 @@ export default function(server: Server) {
                 return closeWithBadPayload(ws, "Binary messages are not supported");
             }
 
+            ws.state.messagesSinceLastCheck++;
+            if (ws.state.messagesSinceLastCheck > MAX_CLIENT_MESSAGES_PER_BATCH) {
+                return closeWithError(ws, gatewayErrors.FLOODING);
+            }
+
             const payload = ensureFormattedGatewayPayload(parseJsonOrNull(rawData.toString()));
             if (!payload) {
                 return closeWithBadPayload(ws, "Invalid JSON or message does not match schema");
diff --git a/src/types/gatewayclientstate.d.ts b/src/types/gatewayclientstate.d.ts
index 387089a..465f58b 100644
--- a/src/types/gatewayclientstate.d.ts
+++ b/src/types/gatewayclientstate.d.ts
@@ -3,5 +3,6 @@ interface GatewayClientState {
     ready: boolean,
     alive: boolean,
     lastAliveCheck: number,
-    dispatchChannels: Set<string>
+    dispatchChannels: Set<string>,
+    messagesSinceLastCheck: number
 }