From 3876ecfde8acabc69a6fa83520166ed8fe9cb05f Mon Sep 17 00:00:00 2001 From: Jarred Sumner Date: Tue, 6 Aug 2024 16:22:21 -0700 Subject: [PATCH] Add test for calling websocket server publish/send methods repeatedly on closed sockets (#13131) --- src/bun.js/api/server.zig | 27 ++++++++ test/harness.ts | 22 ++++++ .../bun/websocket/websocket-server-fixture.js | 67 +++++++++++++++++++ .../js/bun/websocket/websocket-server.test.ts | 17 ++++- 4 files changed, 132 insertions(+), 1 deletion(-) create mode 100644 test/js/bun/websocket/websocket-server-fixture.js diff --git a/src/bun.js/api/server.zig b/src/bun.js/api/server.zig index 58c890215a..6c8c7d3b3e 100644 --- a/src/bun.js/api/server.zig +++ b/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 diff --git a/test/harness.ts b/test/harness.ts index 4dd30f8483..d689c612f7 100644 --- a/test/harness.ts +++ b/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."); + } +} diff --git a/test/js/bun/websocket/websocket-server-fixture.js b/test/js/bun/websocket/websocket-server-fixture.js new file mode 100644 index 0000000000..82fcc8844d --- /dev/null +++ b/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"); diff --git a/test/js/bun/websocket/websocket-server.test.ts b/test/js/bun/websocket/websocket-server.test.ts index f53b2aaf77..e3f06290ac 100644 --- a/test/js/bun/websocket/websocket-server.test.ts +++ b/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);