From 8e97a55331c355368116b517f4a2ee947ad4edea Mon Sep 17 00:00:00 2001
From: Archangel <jesse@archangel.wtf>
Date: Mon, 30 Dec 2024 16:27:46 +0100
Subject: [PATCH] Extend WebSocket to include sendAsync

---
 project/src/servers/WebSocketServer.ts        |  4 ++--
 project/src/servers/ws/SPTWebsocket.ts        | 16 +++++++++++++
 .../ws/SptWebSocketConnectionHandler.ts       | 23 +++++--------------
 3 files changed, 24 insertions(+), 19 deletions(-)
 create mode 100644 project/src/servers/ws/SPTWebsocket.ts

diff --git a/project/src/servers/WebSocketServer.ts b/project/src/servers/WebSocketServer.ts
index 4a142a7c4..b1fe655b7 100644
--- a/project/src/servers/WebSocketServer.ts
+++ b/project/src/servers/WebSocketServer.ts
@@ -57,7 +57,7 @@ export class WebSocketServer {
         if ((socketHandlers?.length ?? 0) === 0) {
             const message = `Socket connection received for url ${req.url}, but there is not websocket handler configured for it`;
             this.logger.warning(message);
-            await this.sendAsync(ws, this.jsonUtil.serialize({ error: message }));
+            await WebSocketServer.sendAsync(ws, this.jsonUtil.serialize({ error: message }));
             ws.close();
             return;
         }
@@ -69,7 +69,7 @@ export class WebSocketServer {
     }
 
     // biome-ignore lint/suspicious/noExplicitAny: Any is required here, I dont see any other way considering it will complain if we use BufferLike
-    protected sendAsync(ws: WebSocket, data: any): Promise<void> {
+    public static sendAsync(ws: WebSocket, data: any): Promise<void> {
         return new Promise((resolve, reject) => {
             ws.send(data, (error) => {
                 if (error) {
diff --git a/project/src/servers/ws/SPTWebsocket.ts b/project/src/servers/ws/SPTWebsocket.ts
new file mode 100644
index 000000000..8e543c49e
--- /dev/null
+++ b/project/src/servers/ws/SPTWebsocket.ts
@@ -0,0 +1,16 @@
+import WebSocket from "ws";
+
+class SPTWebSocket extends WebSocket {
+    // biome-ignore lint/suspicious/noExplicitAny: Any is required here, I dont see any other way considering it will complain if we use BufferLike
+    public sendAsync(ws: WebSocket, data: any): Promise<void> {
+        return new Promise((resolve, reject) => {
+            ws.send(data, (error) => {
+                if (error) {
+                    reject(error);
+                } else {
+                    resolve();
+                }
+            });
+        });
+    }
+}
\ No newline at end of file
diff --git a/project/src/servers/ws/SptWebSocketConnectionHandler.ts b/project/src/servers/ws/SptWebSocketConnectionHandler.ts
index ca1e72025..7d9f92a80 100644
--- a/project/src/servers/ws/SptWebSocketConnectionHandler.ts
+++ b/project/src/servers/ws/SptWebSocketConnectionHandler.ts
@@ -12,11 +12,13 @@ import { LocalisationService } from "@spt/services/LocalisationService";
 import { JsonUtil } from "@spt/utils/JsonUtil";
 import { inject, injectAll, injectable } from "tsyringe";
 import { WebSocket } from "ws";
+import { WebSocketServer } from "../WebSocketServer";
+import { SPTWebSocket } from "./SPTWebsocket";
 
 @injectable()
 export class SptWebSocketConnectionHandler implements IWebSocketConnectionHandler {
     protected httpConfig: IHttpConfig;
-    protected webSockets: Map<string, WebSocket> = new Map<string, WebSocket>();
+    protected webSockets: Map<string, SPTWebSocket> = new Map<string, SPTWebSocket>();
     protected defaultNotification: IWsNotificationEvent = { type: NotificationEventType.PING, eventId: "ping" };
 
     protected websocketPingHandler: NodeJS.Timeout | undefined;
@@ -39,7 +41,7 @@ export class SptWebSocketConnectionHandler implements IWebSocketConnectionHandle
         return "/notifierServer/getwebsocket/";
     }
 
-    public async onConnection(ws: WebSocket, req: IncomingMessage): Promise<void> {
+    public async onConnection(ws: SPTWebSocket, req: IncomingMessage): Promise<void> {
         // Strip request and break it into sections
         const splitUrl = req.url.substring(0, req.url.indexOf("?")).split("/");
         const sessionID = splitUrl.pop();
@@ -67,7 +69,7 @@ export class SptWebSocketConnectionHandler implements IWebSocketConnectionHandle
             this.logger.debug(this.localisationService.getText("websocket-pinging_player", sessionID));
 
             if (ws.readyState === WebSocket.OPEN) {
-                await this.sendAsync(ws, this.jsonUtil.serialize(this.defaultNotification));
+                await ws.sendAsync(ws, this.jsonUtil.serialize(this.defaultNotification));
             } else {
                 this.logger.debug(this.localisationService.getText("websocket-socket_lost_deleting_handle"));
                 clearInterval(this.websocketPingHandler);
@@ -81,7 +83,7 @@ export class SptWebSocketConnectionHandler implements IWebSocketConnectionHandle
             if (this.isConnectionWebSocket(sessionID)) {
                 const ws = this.webSockets.get(sessionID);
 
-                await this.sendAsync(this.webSockets.get(sessionID), this.jsonUtil.serialize(output));
+                await ws.sendAsync(this.webSockets.get(sessionID), this.jsonUtil.serialize(output));
                 this.logger.debug(this.localisationService.getText("websocket-message_sent"));
             } else {
                 this.logger.debug(this.localisationService.getText("websocket-not_ready_message_not_sent", sessionID));
@@ -98,17 +100,4 @@ export class SptWebSocketConnectionHandler implements IWebSocketConnectionHandle
     public getSessionWebSocket(sessionID: string): WebSocket {
         return this.webSockets[sessionID];
     }
-
-    // biome-ignore lint/suspicious/noExplicitAny: Any is required here, I dont see any other way considering it will complain if we use BufferLike
-    public sendAsync(ws: WebSocket, data: any): Promise<void> {
-        return new Promise((resolve, reject) => {
-            ws.send(data, (error) => {
-                if (error) {
-                    reject(error);
-                } else {
-                    resolve();
-                }
-            });
-        });
-    }
 }