From d113803777b14f317188dbfa6bd4e49c54dce9fb Mon Sep 17 00:00:00 2001
From: dave caruso <me@paperdave.net>
Date: Thu, 28 Mar 2024 18:17:55 -0700
Subject: [PATCH] add server ws.terminate (#9693)

* add server ws.terminate

* [autofix.ci] apply automated fixes

* it has to be blazing fast

---------

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
Co-authored-by: Jarred Sumner <jarred@jarredsumner.com>
---
 src/js/thirdparty/ws.js           | 17 +++++++++++++++++
 test/js/first_party/ws/ws.test.ts | 19 +++++++++++++++++++
 2 files changed, 36 insertions(+)

diff --git a/src/js/thirdparty/ws.js b/src/js/thirdparty/ws.js
index 25db515009..731d381c14 100644
--- a/src/js/thirdparty/ws.js
+++ b/src/js/thirdparty/ws.js
@@ -700,6 +700,23 @@ class BunWebSocketMocked extends EventEmitter {
       this.#ws.close(code, reason);
     }
   }
+
+  terminate() {
+    let state = this.#state;
+    if (state === 3) return;
+    if (state === 0) {
+      const msg = "WebSocket was closed before the connection was established";
+      abortHandshake(this, this._req, msg);
+      return;
+    }
+
+    let ws = this.#ws;
+    if (ws) {
+      this.#state = WebSocket.CLOSING;
+      ws.terminate();
+    }
+  }
+
   get binaryType() {
     return this.#binaryType;
   }
diff --git a/test/js/first_party/ws/ws.test.ts b/test/js/first_party/ws/ws.test.ts
index 58468b5a6f..abc031bb7d 100644
--- a/test/js/first_party/ws/ws.test.ts
+++ b/test/js/first_party/ws/ws.test.ts
@@ -276,6 +276,25 @@ describe("WebSocketServer", () => {
     new WebSocket("ws://localhost:" + wss.address().port);
     await promise;
   });
+
+  it("sockets can be terminated", async () => {
+    const wss = new WebSocketServer({ port: 0 });
+    const { resolve, reject, promise } = Promise.withResolvers();
+
+    wss.on("connection", ws => {
+      ws.on("close", () => {
+        resolve();
+      });
+      try {
+        ws.terminate();
+      } catch (err) {
+        reject(err);
+      }
+    });
+
+    new WebSocket("ws://localhost:" + wss.address().port);
+    await promise;
+  });
 });
 
 describe("Server", () => {