diff --git a/server/__tests__/websocket-shutdown.test.js b/server/__tests__/websocket-shutdown.test.js new file mode 100644 index 0000000..5b62d20 --- /dev/null +++ b/server/__tests__/websocket-shutdown.test.js @@ -0,0 +1,56 @@ +/** + * @file websocket-shutdown.test.js + * @description Regression test for the shutdown race in server/websocket.js: + * `closeWebSocket()` nulls both WebSocketServer references but cannot + * unregister the `server.on("upgrade")` listeners it installed, so an upgrade + * arriving mid-shutdown used to throw an unhandled TypeError + * ("Cannot read properties of null (reading 'handleUpgrade')") and kill the + * process. Uses a real http.Server and raw upgrade requests — no ws client. + * @author Nguyễn Ngọc Trí Vĩ + */ +const { describe, it } = require("node:test"); +const assert = require("node:assert/strict"); +const http = require("node:http"); +const net = require("node:net"); +const ws = require("../websocket"); + +/** Send a raw WS upgrade for `path` and resolve once the socket settles. */ +function rawUpgrade(port, path) { + return new Promise((resolve) => { + const socket = net.connect(port, "127.0.0.1", () => { + socket.write( + `GET ${path} HTTP/1.1\r\nHost: 127.0.0.1:${port}\r\n` + + `Upgrade: websocket\r\nConnection: Upgrade\r\n` + + `Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\nSec-WebSocket-Version: 13\r\n\r\n` + ); + }); + let body = ""; + socket.on("data", (d) => { + body += d.toString("utf8"); + }); + socket.on("close", () => resolve(body)); + socket.on("error", () => resolve(body)); + }); +} + +describe("websocket shutdown race", () => { + it("drops upgrades that land after closeWebSocket instead of crashing", async () => { + const server = http.createServer((req, res) => res.end("ok")); + ws.initWebSocket(server); + ws.initPtyWebSocket(server); + await new Promise((r) => server.listen(0, "127.0.0.1", r)); + const { port } = server.address(); + + // The shutdown the real server performs on SIGTERM. + ws.closeWebSocket(); + + // Both paths, because each has its own upgrade listener and its own null + // reference. An unguarded handleUpgrade here takes the process down, so + // simply reaching the assertions below is the test passing. + await rawUpgrade(port, "/ws"); + await rawUpgrade(port, "/ws-pty/ccam-lane-1"); + + assert.equal(ws.getConnectionCount(), 0); + await new Promise((r) => server.close(r)); + }); +}); diff --git a/server/websocket.js b/server/websocket.js index b4cffd9..fdf99d8 100644 --- a/server/websocket.js +++ b/server/websocket.js @@ -29,6 +29,14 @@ function initWebSocket(server) { server.on("upgrade", (req, socket, head) => { const url = new URL(req.url, "http://localhost"); if (url.pathname !== "/ws") return; // not ours — `/ws-pty/*` handles its own path. + // closeWebSocket() nulls `wss` but cannot unregister this listener, so an + // upgrade arriving mid-shutdown would throw an unhandled TypeError and take + // the process down instead of letting it exit gracefully. Clients retry, so + // dropping the socket is the correct answer here. + if (!wss) { + socket.destroy(); + return; + } wss.handleUpgrade(req, socket, head, (ws) => { wss.emit("connection", ws, req); }); @@ -118,6 +126,12 @@ function initPtyWebSocket(server) { socket.destroy(); return; } + // Same shutdown race as `/ws` above — and far easier to hit here, since + // TerminalView re-attaches every 1.5s while a run console is open. + if (!ptyWss) { + socket.destroy(); + return; + } ptyWss.handleUpgrade(req, socket, head, (ws) => { ptyWss.emit("connection", ws, match[1]); });