From 8c52abce7fb107c92846c623b1ae9f09f5a25fe6 Mon Sep 17 00:00:00 2001
From: ITOH <72305210+itohatweb@users.noreply.github.com>
Date: Thu, 15 Apr 2021 16:46:44 +0200
Subject: [PATCH] add closeWS function

---
 src/ws/close_ws.ts                 | 6 ++++++
 src/ws/heartbeat.ts                | 3 ++-
 src/ws/identify.ts                 | 5 ++---
 src/ws/resume.ts                   | 8 +++-----
 src/ws/tell_cluster_to_identify.ts | 5 +++--
 tests/ws/ws_close.ts               | 3 ++-
 6 files changed, 18 insertions(+), 12 deletions(-)
 create mode 100644 src/ws/close_ws.ts

diff --git a/src/ws/close_ws.ts b/src/ws/close_ws.ts
new file mode 100644
index 000000000..399174439
--- /dev/null
+++ b/src/ws/close_ws.ts
@@ -0,0 +1,6 @@
+/** Use this function to close a ws connection properly */
+export function closeWS(ws: WebSocket, code?: number, reason?: string) {
+  if (ws.readyState !== WebSocket.OPEN) return;
+
+  ws.close(code, reason);
+}
diff --git a/src/ws/heartbeat.ts b/src/ws/heartbeat.ts
index a32d964ee..29576ac4c 100644
--- a/src/ws/heartbeat.ts
+++ b/src/ws/heartbeat.ts
@@ -1,5 +1,6 @@
 import { DiscordGatewayOpcodes } from "../types/codes/gateway_opcodes.ts";
 import { delay } from "../util/utils.ts";
+import { closeWS } from "./close_ws.ts";
 import { identify } from "./identify.ts";
 import { ws } from "./ws.ts";
 
@@ -43,7 +44,7 @@ export async function heartbeat(shardId: number, interval: number) {
     }
 
     if (!currentShard.heartbeat.acknowledged) {
-      currentShard.ws.close(3065, "Did not receive an ACK in time.");
+      closeWS(currentShard.ws, 3066, "Did not receive an ACK in time.");
       return identify(shardId, ws.maxShards);
     }
 
diff --git a/src/ws/identify.ts b/src/ws/identify.ts
index 19e5a92ee..0671e8d9a 100644
--- a/src/ws/identify.ts
+++ b/src/ws/identify.ts
@@ -1,4 +1,5 @@
 import { DiscordGatewayOpcodes } from "../types/codes/gateway_opcodes.ts";
+import { closeWS } from "./close_ws.ts";
 import { ws } from "./ws.ts";
 
 export async function identify(shardId: number, maxShards: number) {
@@ -7,9 +8,7 @@ export async function identify(shardId: number, maxShards: number) {
   // Need to clear the old heartbeat interval
   const oldShard = ws.shards.get(shardId);
   if (oldShard) {
-    if (oldShard.ws.readyState === WebSocket.OPEN) {
-      oldShard.ws.close(3065, "Reidentifying closure of old shard");
-    }
+    closeWS(oldShard.ws, 3065, "Reidentifying closure of old shard");
     clearInterval(oldShard.heartbeat.intervalId);
   }
 
diff --git a/src/ws/resume.ts b/src/ws/resume.ts
index f54439905..db8c4135d 100644
--- a/src/ws/resume.ts
+++ b/src/ws/resume.ts
@@ -1,4 +1,5 @@
 import { DiscordGatewayOpcodes } from "../types/codes/gateway_opcodes.ts";
+import { closeWS } from "./close_ws.ts";
 import { identify } from "./identify.ts";
 import { ws } from "./ws.ts";
 
@@ -13,11 +14,8 @@ export async function resume(shardId: number) {
   // CREATE A SHARD
   const socket = await ws.createShard(shardId);
 
-  // ONLY CLOSE IF SHARD SOCKET IS STILL CONNECTED
-  if (oldShard.ws.readyState === WebSocket.OPEN) {
-    // HOW TO CLOSE OLD SHARD SOCKET!!!
-    oldShard.ws.close(3065, "Resuming the shard, closing old shard.");
-  }
+  // HOW TO CLOSE OLD SHARD SOCKET!!!
+  closeWS(oldShard.ws, 3064, "Resuming the shard, closing old shard.");
   // STOP OLD HEARTBEAT
   clearInterval(oldShard.heartbeat.intervalId);
 
diff --git a/src/ws/tell_cluster_to_identify.ts b/src/ws/tell_cluster_to_identify.ts
index 99a5951eb..6a3974584 100644
--- a/src/ws/tell_cluster_to_identify.ts
+++ b/src/ws/tell_cluster_to_identify.ts
@@ -1,3 +1,4 @@
+import { closeWS } from "./close_ws.ts";
 import { ws } from "./ws.ts";
 
 /** Allows users to hook in and change to communicate to different clusters across different servers or anything they like. For example using redis pubsub to talk to other servers. */
@@ -10,7 +11,7 @@ export async function tellClusterToIdentify(
   const oldShard = ws.shards.get(shardId);
   await ws.identify(shardId, ws.maxShards);
 
-  if (oldShard?.ws.readyState === WebSocket.OPEN) {
-    oldShard.ws.close(3065, "Resharded!");
+  if (oldShard) {
+    closeWS(oldShard.ws, 3063, "Resharded!");
   }
 }
diff --git a/tests/ws/ws_close.ts b/tests/ws/ws_close.ts
index 3b313fc99..1a1ce6111 100644
--- a/tests/ws/ws_close.ts
+++ b/tests/ws/ws_close.ts
@@ -1,4 +1,5 @@
 import { delay } from "../../src/util/utils.ts";
+import { closeWS } from "../../src/ws/close_ws.ts";
 import { ws } from "../../src/ws/ws.ts";
 import { defaultTestOptions } from "./start_bot.ts";
 
@@ -8,7 +9,7 @@ Deno.test({
   async fn() {
     ws.shards.forEach((shard) => {
       clearInterval(shard.heartbeat.intervalId);
-      shard.ws.close(3064, "Discordeno Testing Finished! Do Not RESUME!");
+      closeWS(shard.ws, 3061, "Discordeno Testing Finished! Do Not RESUME!");
     });
 
     await delay(3000);