Add test for calling websocket server publish/send methods repeatedly on closed sockets (#13131)

src/bun.js/api/server.zig

@@ -4356,6 +4356,11 @@ pub const ServerWebSocket = struct {
             return .zero;
         }
 
+        if (this.isClosed() and !publish_to_self) {
+            // We can't access the socket context on a closed socket.
+            return JSValue.jsNumber(0);
+        }
+
         if (message_value.asArrayBuffer(globalThis)) |array_buffer| {
             const buffer = array_buffer.slice();
 
@@ -4438,6 +4443,11 @@ pub const ServerWebSocket = struct {
             return .zero;
         }
 
+        if (this.isClosed() and !publish_to_self) {
+            // Can't publish on a closed socket.
+            return JSValue.jsNumber(0);
+        }
+
         var string_slice = message_value.toSlice(globalThis, bun.default_allocator);
         defer string_slice.deinit();
 
@@ -4503,6 +4513,12 @@ pub const ServerWebSocket = struct {
             globalThis.throw("publishBinary requires a non-empty message", .{});
             return .zero;
         }
+
+        if (this.isClosed() and !publish_to_self) {
+            // Can't publish on a closed socket.
+            return JSValue.jsNumber(0);
+        }
+
         const array_buffer = message_value.asArrayBuffer(globalThis) orelse {
             globalThis.throw("publishBinary expects an ArrayBufferView", .{});
             return .zero;
@@ -4549,6 +4565,11 @@ pub const ServerWebSocket = struct {
             return JSC.JSValue.jsNumber(0);
         }
 
+        if (this.isClosed() and !publish_to_self) {
+            // We can't access the socket context on a closed socket.
+            return JSValue.jsNumber(0);
+        }
+
         const result = if (!publish_to_self)
             this.websocket().publish(topic_slice.slice(), buffer, .binary, compress)
         else
@@ -4591,6 +4612,12 @@ pub const ServerWebSocket = struct {
         if (buffer.len == 0) {
             return JSC.JSValue.jsNumber(0);
         }
+
+        if (this.isClosed() and !publish_to_self) {
+            // We can't access the socket context on a closed socket.
+            return JSValue.jsNumber(0);
+        }
+
         const result = if (!publish_to_self)
             this.websocket().publish(topic_slice.slice(), buffer, .text, compress)
         else

test/harness.ts

@@ -1173,3 +1173,25 @@ export function isMacOSVersionAtLeast(minVersion: number): boolean {
   }
   return parseFloat(macOSVersion) >= minVersion;
 }
+
+let hasGuardMalloc = -1;
+export function forceGuardMalloc(env) {
+  if (process.platform !== "darwin") {
+    return;
+  }
+
+  if (hasGuardMalloc === -1) {
+    hasGuardMalloc = Number(fs.existsSync("/usr/lib/libgmalloc.dylib"));
+  }
+
+  if (hasGuardMalloc === 1) {
+    env.DYLD_INSERT_LIBRARIES = "/usr/lib/libgmalloc.dylib";
+    env.MALLOC_PROTECT_BEFORE = "1";
+    env.MallocScribble = "1";
+    env.MallocGuardEdges = "1";
+    env.MALLOC_FILL_SPACE = "1";
+    env.MALLOC_STRICT_SIZE = "1";
+  } else {
+    console.warn("Guard malloc is not available on this platform for some reason.");
+  }
+}

test/js/bun/websocket/websocket-server-fixture.js

@@ -0,0 +1,67 @@
+// For this test to consistently reproduce the original issue, you need guard malloc enabled.
+// DYLD_INSERT_LIBRARIES=(realpath  /usr/lib/libgmalloc.dylib)
+// MALLOC_PROTECT_BEFORE=1
+// MallocScribble=1
+// MallocGuardEdges=1
+// MALLOC_FILL_SPACE=1
+// MALLOC_STRICT_SIZE=1
+
+let pending = [];
+using server = Bun.serve({
+  websocket: {
+    open(ws) {
+      globalThis.sockets ??= [];
+      globalThis.sockets.push(ws);
+      ws.data = Promise.withResolvers();
+      pending.push(ws.data.promise);
+      ws.subscribe("bye");
+      setTimeout(() => {
+        ws.close();
+      });
+    },
+    close(ws, code, reason) {
+      setTimeout(
+        ws => {
+          Bun.gc();
+          for (let i = 0; i < 10; i++) {
+            ws.publishText("bye", "ok");
+            ws.publishBinary("bye", Buffer.from("ok"));
+            ws.publish("bye", "ok");
+            ws.subscribe("bye", "ok");
+            ws.isSubscribed("bye", "ok");
+            ws.send("bye");
+            ws.sendText("bye");
+            ws.sendBinary(Buffer.from("bye"));
+          }
+          ws.data.resolve();
+        },
+        10,
+        ws,
+      );
+      Bun.gc();
+    },
+  },
+  fetch(req, server) {
+    return server.upgrade(req);
+  },
+});
+
+for (let i = 0; i < 5; i++) {
+  const ws = new WebSocket(`ws://${server.hostname}:${server.port}`);
+
+  let { promise, resolve } = Promise.withResolvers();
+  ws.addEventListener("open", () => {});
+  ws.addEventListener("message", e => {});
+
+  ws.addEventListener("close", () => {
+    console.count("Closed");
+    resolve();
+  });
+  pending.push(promise);
+}
+
+await Bun.sleep(1);
+Bun.gc(true);
+await Promise.all(pending);
+Bun.gc(true);
+console.log("Exiting");

test/js/bun/websocket/websocket-server.test.ts

@@ -1,7 +1,7 @@
 import { describe, it, expect, afterEach } from "bun:test";
 import type { Server, Subprocess, WebSocketHandler } from "bun";
 import { serve, spawn } from "bun";
-import { bunEnv, bunExe, nodeExe } from "harness";
+import { bunEnv, bunExe, forceGuardMalloc, nodeExe } from "harness";
 import { isIP } from "node:net";
 import path from "node:path";
 
@@ -62,6 +62,21 @@ const binaryTypes = [
 let servers: Server[] = [];
 let clients: Subprocess[] = [];
 
+it("should work fine if you repeatedly call methods on closed websockets", async () => {
+  let env = { ...bunEnv };
+  forceGuardMalloc(env);
+
+  const { exited } = Bun.spawn({
+    cmd: [bunExe(), path.join(import.meta.dir, "websocket-server-fixture.js")],
+    env,
+    stderr: "inherit",
+    stdout: "inherit",
+    stdin: "inherit",
+  });
+
+  expect(await exited).toBe(0);
+});
+
 afterEach(() => {
   for (const server of servers) {
     server.stop(true);