diff --git a/server/src/__tests__/live-events-ws.test.ts b/server/src/__tests__/live-events-ws.test.ts new file mode 100644 index 00000000..b52c1106 --- /dev/null +++ b/server/src/__tests__/live-events-ws.test.ts @@ -0,0 +1,123 @@ +import { EventEmitter } from "node:events"; +import type { IncomingMessage } from "node:http"; +import type { Duplex } from "node:stream"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { setupLiveEventsWebSocketServer } from "../realtime/live-events-ws.js"; +import { logger } from "../middleware/logger.js"; + +vi.mock("../middleware/logger.js", () => ({ + logger: { + warn: vi.fn(), + error: vi.fn(), + }, +})); + +class FakeUpgradeSocket extends EventEmitter { + destroyed = false; + writable = true; + writableEnded = false; + writableDestroyed = false; + endedChunks: string[] = []; + destroyCalls = 0; + + end(chunk?: string) { + if (chunk) this.endedChunks.push(chunk); + this.writableEnded = true; + this.writable = false; + setImmediate(() => { + if (this.destroyed) return; + this.emit("finish"); + if (!this.destroyed) { + this.emit("close"); + } + }); + return this; + } + + destroy() { + this.destroyCalls += 1; + this.destroyed = true; + this.writable = false; + this.writableDestroyed = true; + this.emit("close"); + return this; + } + + emitSocketError(err: Error) { + this.writable = false; + this.writableDestroyed = true; + this.emit("error", err); + } +} + +function createUpgradeRequest(overrides: Partial = {}) { + return { + url: "/api/companies/company-1/events/ws", + headers: {}, + ...overrides, + } as IncomingMessage; +} + +async function flushPromises() { + await new Promise((resolve) => setImmediate(resolve)); +} + +describe("setupLiveEventsWebSocketServer", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("does not write a rejection response after the raw upgrade socket is already closed", async () => { + const server = new EventEmitter(); + setupLiveEventsWebSocketServer(server as never, {} as never, { deploymentMode: "authenticated" }); + const socket = new FakeUpgradeSocket(); + + server.emit("upgrade", createUpgradeRequest(), socket as unknown as Duplex, Buffer.alloc(0)); + socket.destroy(); + await flushPromises(); + + expect(socket.endedChunks).toEqual([]); + expect(socket.destroyCalls).toBe(1); + }); + + it("handles raw upgrade socket errors during async authorization", async () => { + const server = new EventEmitter(); + let resolveSession: (value: null) => void = () => undefined; + setupLiveEventsWebSocketServer(server as never, {} as never, { + deploymentMode: "authenticated", + resolveSessionFromHeaders: () => + new Promise((resolve) => { + resolveSession = resolve; + }), + }); + const socket = new FakeUpgradeSocket(); + + server.emit("upgrade", createUpgradeRequest(), socket as unknown as Duplex, Buffer.alloc(0)); + expect(() => socket.emitSocketError(new Error("write EPIPE"))).not.toThrow(); + resolveSession(null); + await flushPromises(); + + expect(logger.warn).toHaveBeenCalledWith( + expect.objectContaining({ err: expect.any(Error), path: "/api/companies/company-1/events/ws" }), + "live websocket upgrade socket error", + ); + expect(socket.endedChunks).toEqual([]); + expect(socket.destroyed).toBe(true); + }); + + it("destroys and cleans up listeners after flushing a rejection response", async () => { + const server = new EventEmitter(); + setupLiveEventsWebSocketServer(server as never, {} as never, { deploymentMode: "authenticated" }); + const socket = new FakeUpgradeSocket(); + + server.emit("upgrade", createUpgradeRequest(), socket as unknown as Duplex, Buffer.alloc(0)); + await flushPromises(); + await flushPromises(); + + expect(socket.endedChunks[0]).toContain("403 Forbidden"); + expect(socket.destroyed).toBe(true); + expect(socket.listenerCount("error")).toBe(0); + expect(socket.listenerCount("close")).toBe(0); + expect(socket.listenerCount("finish")).toBe(0); + }); +}); diff --git a/server/src/realtime/live-events-ws.ts b/server/src/realtime/live-events-ws.ts index d18e2930..07f4b05f 100644 --- a/server/src/realtime/live-events-ws.ts +++ b/server/src/realtime/live-events-ws.ts @@ -54,10 +54,31 @@ function hashToken(token: string) { return createHash("sha256").update(token).digest("hex"); } +function isWritableUpgradeSocket(socket: Duplex) { + const maybeWritableState = socket as Duplex & { writable?: boolean; writableEnded?: boolean; writableDestroyed?: boolean }; + return !socket.destroyed && maybeWritableState.writable !== false && !maybeWritableState.writableEnded && !maybeWritableState.writableDestroyed; +} + +function closeUpgradeSocket(socket: Duplex) { + if (!socket.destroyed) { + socket.destroy(); + } +} + function rejectUpgrade(socket: Duplex, statusLine: string, message: string) { const safe = message.replace(/[\r\n]+/g, " ").trim(); - socket.write(`HTTP/1.1 ${statusLine}\r\nConnection: close\r\nContent-Type: text/plain\r\n\r\n${safe}`); - socket.destroy(); + if (!isWritableUpgradeSocket(socket)) { + closeUpgradeSocket(socket); + return; + } + + try { + socket.once("finish", () => closeUpgradeSocket(socket)); + socket.end(`HTTP/1.1 ${statusLine}\r\nConnection: close\r\nContent-Type: text/plain\r\n\r\n${safe}`); + } catch (err) { + logger.warn({ err }, "failed to reject live websocket upgrade"); + closeUpgradeSocket(socket); + } } function parseCompanyId(pathname: string) { @@ -234,6 +255,17 @@ export function setupLiveEventsWebSocketServer( }); server.on("upgrade", (req, socket, head) => { + const onRawSocketError = (err: Error) => { + logger.warn({ err, path: req.url }, "live websocket upgrade socket error"); + }; + const cleanupRawSocketListeners = () => { + socket.off("error", onRawSocketError); + socket.off("close", cleanupRawSocketListeners); + }; + + socket.on("error", onRawSocketError); + socket.once("close", cleanupRawSocketListeners); + if (!req.url) { rejectUpgrade(socket, "400 Bad Request", "missing url"); return; @@ -242,7 +274,7 @@ export function setupLiveEventsWebSocketServer( const url = new URL(req.url, "http://localhost"); const companyId = parseCompanyId(url.pathname); if (!companyId) { - socket.destroy(); + closeUpgradeSocket(socket); return; } @@ -256,9 +288,15 @@ export function setupLiveEventsWebSocketServer( return; } + if (!isWritableUpgradeSocket(socket)) { + cleanupRawSocketListeners(); + return; + } + const reqWithContext = req as IncomingMessageWithContext; reqWithContext.paperclipUpgradeContext = context; + cleanupRawSocketListeners(); wss.handleUpgrade(req, socket, head, (ws: WsSocket) => { wss.emit("connection", ws, reqWithContext); });