From 0950241d3be13702b8f5456d3b6b97fe18ac8e80 Mon Sep 17 00:00:00 2001
From: hippoz <10706925-hippoz@users.noreply.gitlab.com>
Date: Thu, 3 Feb 2022 20:41:23 +0200
Subject: [PATCH] hopefully make the DiscordClient connection more robust

---
 DiscordClient.js | 111 +++++++++++++++++++++++++++++++++++------------
 common.js        |   7 +--
 2 files changed, 84 insertions(+), 34 deletions(-)

diff --git a/DiscordClient.js b/DiscordClient.js
index 70de801..342ceb7 100644
--- a/DiscordClient.js
+++ b/DiscordClient.js
@@ -7,10 +7,16 @@ const opcodes = {
     EVENT: 0,
     CLIENT_HEARTBEAT: 1,
     IDENTIFY: 2,
+    RESUME: 6,
+    INVALID_SESSION: 9,
     HELLO: 10,
     HEARTBEAT_ACK: 11,
 };
 
+const reconnectOnCloseCodes = [
+    1000, 1001, 4000, 4001, 4002, 4003, 4005, 4007, 4008, 4009
+];
+
 class DiscordClient extends EventEmitter {
     constructor(token, { intents, baseDomain="discord.com", gatewayUrl="wss://gateway.discord.gg/?v=9&encoding=json&compress=zlib-stream", apiBase="https://discord.com/api/v9" } = {}) {
         super();
@@ -50,6 +56,37 @@ class DiscordClient extends EventEmitter {
         }, this._heartbeatIntervalTime);
     }
 
+    _getIdentifyPayload() {
+        return {
+            token: this.token,
+            intents: this.intents,
+            properties: {
+                "$os": "linux",
+                "$browser": "generic",
+                "$device": "generic"
+            },
+            presence: {
+                since: Date.now(),
+                activities: [
+                    {
+                        type: 2, // LISTENING
+                        name: "the voices"
+                    }
+                ],
+                status: "online",
+                afk: false
+            }
+        };
+    }
+
+    _getResumePayload() {
+        return {
+            token: this.token,
+            session_id: this.sessionId,
+            seq: this.seq
+        };
+    }
+
     _handleGatewayMessage(ws, message) {
         try {
             message = JSON.parse(message);
@@ -68,29 +105,18 @@ class DiscordClient extends EventEmitter {
             case opcodes.HELLO: {
                 this._setHeartbeat(payload.heartbeat_interval);
 
-                ws.send(JSON.stringify({
-                    op: opcodes.IDENTIFY,
-                    d: {
-                        token: this.token,
-                        intents: this.intents,
-                        properties: {
-                            "$os": "linux",
-                            "$browser": "generic",
-                            "$device": "generic"
-                        },
-                        presence: {
-                            since: Date.now(),
-                            activities: [
-                                {
-                                    type: 2, // LISTENING
-                                    name: "the voices"
-                                }
-                            ],
-                            status: "online",
-                            afk: false
-                        }
-                    }
-                }));
+                if (this.resuming) {
+                    console.warn("DiscordClient: resuming...");
+                    ws.send(JSON.stringify({
+                        op: opcodes.RESUME,
+                        d: this._getResumePayload()
+                    }));
+                } else {
+                    ws.send(JSON.stringify({
+                        op: opcodes.IDENTIFY,
+                        d: this._getIdentifyPayload()
+                    }));
+                }
                 break;
             }
 
@@ -183,6 +209,23 @@ class DiscordClient extends EventEmitter {
                 break;
             }
 
+            case opcodes.INVALID_SESSION: {
+                if (message.d) {
+                    // connection is resumable, we are going to resume the connection
+                    this.resuming = true;
+                    this.connect();
+                } else {
+                    // connection is not resumable, wait some time and then send a new IDENTIFY payload
+                    setTimeout(() => {
+                        ws.send(JSON.stringify({
+                            op: opcodes.IDENTIFY,
+                            d: this._getIdentifyPayload()
+                        }));
+                    }, 3500);
+                }
+                break;
+            }
+
             default: {
                 console.warn(`warn: DiscordClient: got unhandled opcode "${message.op}"`);
                 break;
@@ -191,6 +234,12 @@ class DiscordClient extends EventEmitter {
     }
 
     connect() {
+        if (this.ws) {
+            this.ws.removeAllListeners();
+            this.ws.close();
+            this.ws = null;
+        }
+
         const ws = new WebSocket(this.gatewayUrl);
 
         this.ws = ws;
@@ -208,16 +257,22 @@ class DiscordClient extends EventEmitter {
 
         ws.on("close", (code, reason) => {
             reason = reason.toString();
-            this.emit("close", code, reason);
             console.error(`DiscordClient: on \`close\`: disconnected from gateway: code \`${code}\`, reason \`${reason}\``);
+
+            this.emit("close", code, reason);
+            this._setHeartbeat(-1);
+            if (reconnectOnCloseCodes.includes(code)) {
+                this.resuming = true;
+                this.connect();
+            }
         });
 
         ws.on("error", (e) => {
             console.error("DiscordClient: websocket error:", e);
-            console.log("DiscordClient: reconnecting in a couple of miliseconds");
-            setTimeout(() => {
-                this.connect();
-            }, 400);
+            console.log("DiscordClient: reconnecting?");
+            this._setHeartbeat(-1);
+            this.resuming = true;
+            this.connect();
         });
     }
 
diff --git a/common.js b/common.js
index 222640f..1b62a2d 100644
--- a/common.js
+++ b/common.js
@@ -13,7 +13,7 @@ export function wait(time, shouldReject=false) {
     });
 }
 
-bot.on("READY", () => {
+bot.once("READY", () => {
     watchedGuildIds.forEach(id => {
         const watchedGuild = new WatchedGuild();
         watchedGuild.upstreamGuildId = id;
@@ -21,8 +21,3 @@ bot.on("READY", () => {
         guildMap.set(id, watchedGuild);
     });
 });
-
-bot.on("close", (code, reason) => {
-    console.log("bot: connection closed, reconnecting...");
-    bot.connect();
-});