From a624cad2cac928a713b6c22c67395728a4756019 Mon Sep 17 00:00:00 2001
From: Fleny <Fleny113@outlook.com>
Date: Tue, 25 Feb 2025 08:53:12 +0100
Subject: [PATCH] fix(gateway): Keep state to avoid unwanted disconnections
 (#4138)

* Keep state to avoid unwanted disconnections

This keeps a state to know that we wanted to end the connection with a
code of 1000 or 1001 to avoid unwanted disconnections in case discord /
someone in the middle of the TCP connection (such as cloudflare) sends a
close code of 1000 or 1001.

* Apply suggestions from code review

Co-authored-by: Awesome Stickz <awesome@stickz.dev>

---------

Co-authored-by: Awesome Stickz <awesome@stickz.dev>
---
 packages/gateway/src/Shard.ts | 29 +++++++++++++++++++++++++++--
 1 file changed, 27 insertions(+), 2 deletions(-)

diff --git a/packages/gateway/src/Shard.ts b/packages/gateway/src/Shard.ts
index 6b8b90d6d..4db743aa5 100644
--- a/packages/gateway/src/Shard.ts
+++ b/packages/gateway/src/Shard.ts
@@ -55,6 +55,16 @@ export class DiscordenoShard {
   bucket: LeakyBucket
   /** Logger for the bucket. */
   logger: Pick<typeof logger, 'debug' | 'info' | 'warn' | 'error' | 'fatal'>
+  /**
+   * Is the shard going offline?
+   *
+   * @remarks
+   * This will be true if the close method has been called with either 1000 or 1001
+   *
+   * @internal
+   * This is for internal purposes only, and subject to breaking changes.
+   */
+  goingOffline = false
   /** Text decoder used for compressed payloads. */
   textDecoder = new TextDecoder()
   /** ZLib Inflate instance for ZLib-stream transport payloads. */
@@ -134,6 +144,8 @@ export class DiscordenoShard {
       return
     }
 
+    this.goingOffline = code === GatewayCloseEventCodes.NormalClosure || code === GatewayCloseEventCodes.GoingAway
+
     // This has to be created before the actual call to socket.close as for example Bun calls socket.onclose immediately on the .close() call instead of waiting for the connection to end
     const promise = new Promise((resolve) => {
       this.resolveAfterClose = resolve
@@ -405,8 +417,6 @@ export class DiscordenoShard {
         return
       }
       // On these codes a manual start will be done.
-      case GatewayCloseEventCodes.NormalClosure:
-      case GatewayCloseEventCodes.GoingAway:
       case ShardSocketCloseCodes.Shutdown:
       case ShardSocketCloseCodes.ReIdentifying:
       case ShardSocketCloseCodes.Resharded:
@@ -440,6 +450,21 @@ export class DiscordenoShard {
         await this.identify()
         return
       }
+      // NOTE: This case must always be right above the cases that runs with default case because of how switch works when you don't break / return, more info below.
+      case GatewayCloseEventCodes.NormalClosure:
+      case GatewayCloseEventCodes.GoingAway: {
+        // If the shard is marked as goingOffline, it stays disconnected.
+        if (this.goingOffline) {
+          this.state = ShardState.Disconnected
+          this.events.disconnected?.(this)
+
+        this.goingOffline = false
+
+          return
+        }
+
+        // Otherwise, we want the shard to go through the default case where it gets resumed, as it might be an unexpected closure from Discord or Cloudflare for example, so we don't use break / return here.
+      }
       // Gateway connection closes on which a resume is allowed.
       case GatewayCloseEventCodes.UnknownError:
       case GatewayCloseEventCodes.UnknownOpcode: