From 872c698132aa2e0ee4f105ad1c872c1cd3b9f78a Mon Sep 17 00:00:00 2001 From: nntrivi2001 Date: Wed, 12 Aug 2026 10:13:14 +0700 Subject: [PATCH] feat(run): add /ws-pty/:runId PTY transport bridging WS to tmux attach --- server/__tests__/pty-attach.test.js | 105 ++++++++++++++++++++++++++++ server/index.js | 3 +- server/lib/pty-attach.js | 104 +++++++++++++++++++++++++++ server/websocket.js | 93 +++++++++++++++++++++--- 4 files changed, 293 insertions(+), 12 deletions(-) create mode 100644 server/__tests__/pty-attach.test.js create mode 100644 server/lib/pty-attach.js diff --git a/server/__tests__/pty-attach.test.js b/server/__tests__/pty-attach.test.js new file mode 100644 index 0000000..8482438 --- /dev/null +++ b/server/__tests__/pty-attach.test.js @@ -0,0 +1,105 @@ +/** + * @file pty-attach.test.js + * @description Unit tests for the PTY-attach helper's runId validation and + * data/control framing, using a fake node-pty implementation (no real tmux + * or PTY spawned). + * @author Nguyễn Ngọc Trí Vĩ + */ +const { describe, it, beforeEach } = require("node:test"); +const assert = require("node:assert/strict"); +const { EventEmitter } = require("node:events"); +const ptyAttach = require("../lib/pty-attach"); + +function makeFakePty() { + const emitter = new EventEmitter(); + const writes = []; + const pty = { + onData: (fn) => emitter.on("data", fn), + onExit: (fn) => emitter.on("exit", fn), + write: (d) => writes.push(d), + resize: (cols, rows) => writes.push({ resize: [cols, rows] }), + kill: () => emitter.emit("exit", { exitCode: 0 }), + __emitter: emitter, + __writes: writes, + }; + return pty; +} + +function makeFakeWs() { + const emitter = new EventEmitter(); + const sent = []; + return { + on: (evt, fn) => emitter.on(evt, fn), + send: (data, opts) => sent.push({ data, binary: !!(opts && opts.binary) }), + __emitter: emitter, + __sent: sent, + }; +} + +describe("pty-attach", () => { + let fakePty; + beforeEach(() => { + fakePty = makeFakePty(); + ptyAttach.__setSpawnImpl(() => fakePty); + }); + + it("rejects a runId that doesn't match ccam-lane-", () => { + assert.throws(() => ptyAttach.validateRunId("../../etc/passwd"), /EBADRUNID/); + assert.throws(() => ptyAttach.validateRunId("some-other-session"), /EBADRUNID/); + }); + + it("accepts a well-formed runId", () => { + assert.doesNotThrow(() => ptyAttach.validateRunId("ccam-lane-42")); + }); + + it("wires PTY data to binary WS frames and WS binary frames to PTY writes", () => { + const ws = makeFakeWs(); + ptyAttach.attach(ws, "ccam-lane-1", { cols: 80, rows: 24 }); + + fakePty.__emitter.emit("data", "hello from claude"); + assert.equal(ws.__sent.length, 1); + assert.equal(ws.__sent[0].data, "hello from claude"); + assert.equal(ws.__sent[0].binary, true); + + ws.__emitter.emit("message", Buffer.from("typed text"), { binary: true }); + assert.deepEqual(fakePty.__writes[0], "typed text"); + }); + + it("routes a JSON text frame with type resize to pty.resize", () => { + const ws = makeFakeWs(); + ptyAttach.attach(ws, "ccam-lane-1", { cols: 80, rows: 24 }); + ws.__emitter.emit( + "message", + Buffer.from(JSON.stringify({ type: "resize", cols: 100, rows: 40 })), + { + binary: false, + } + ); + assert.deepEqual(fakePty.__writes[0], { resize: [100, 40] }); + }); + + it("sends an exit control message and closes on PTY exit", () => { + const ws = makeFakeWs(); + let closed = false; + ws.close = () => { + closed = true; + }; + ptyAttach.attach(ws, "ccam-lane-1", { cols: 80, rows: 24 }); + fakePty.__emitter.emit("exit", { exitCode: 0 }); + const last = ws.__sent[ws.__sent.length - 1]; + assert.equal(last.binary, false); + assert.deepEqual(JSON.parse(last.data), { type: "exit", code: 0 }); + assert.equal(closed, true); + }); + + it("kills the PTY attach process when the WS connection closes", () => { + const ws = makeFakeWs(); + ptyAttach.attach(ws, "ccam-lane-1", { cols: 80, rows: 24 }); + let killed = false; + fakePty.kill = () => { + killed = true; + }; + ws.__emitter.emit("close"); + assert.equal(killed, true); + }); +}); diff --git a/server/index.js b/server/index.js index 04687bf..2410efe 100644 --- a/server/index.js +++ b/server/index.js @@ -33,7 +33,7 @@ const cors = require("cors"); const path = require("path"); const http = require("http"); const swaggerUi = require("swagger-ui-express"); -const { initWebSocket } = require("./websocket"); +const { initWebSocket, initPtyWebSocket } = require("./websocket"); const { createOpenApiSpec } = require("./openapi"); const { redocBundlePath, renderRedocHtml } = require("./lib/redoc"); const { writeServerInfo, removeServerInfo, peersSharingDataDir } = require("./lib/server-info"); @@ -150,6 +150,7 @@ function createApp() { function startServer(app, port) { const server = http.createServer(app); initWebSocket(server); + initPtyWebSocket(server); const isProduction = process.env.NODE_ENV === "production"; if (isProduction) { diff --git a/server/lib/pty-attach.js b/server/lib/pty-attach.js new file mode 100644 index 0000000..ed1948f --- /dev/null +++ b/server/lib/pty-attach.js @@ -0,0 +1,104 @@ +/** + * @file pty-attach.js + * @description Bridges one WebSocket connection to a `node-pty`-backed + * `tmux attach-session` process. Binary WS frames carry raw PTY bytes in + * both directions; text WS frames carry small JSON control messages + * (`resize`, and an outbound `exit` sent once when the pane process/tmux + * session ends). Multiple browser tabs each get their own PTY attach + * process — tmux itself is what keeps them all in sync, this module does no + * cross-connection coordination. + * @author Nguyễn Ngọc Trí Vĩ + */ + +const RUN_ID_RE = /^ccam-lane-\d+$/; + +/** + * Reject anything that isn't exactly `ccam-lane-` before it can ever + * reach a tmux/PTY command — the trust boundary for this WS path, since a + * validated runId is the only thing standing between an authenticated WS + * client and naming an arbitrary session on the host. + */ +function validateRunId(runId) { + if (typeof runId !== "string" || !RUN_ID_RE.test(runId)) { + const err = new Error(`EBADRUNID: invalid runId: ${runId}`); + err.code = "EBADRUNID"; + throw err; + } +} + +// Test seam — real implementation set in Step 4. +let spawnImpl = null; +function __setSpawnImpl(fn) { + spawnImpl = fn; +} + +/** + * Attach `ws` to the tmux session `runId`. Spawns one PTY-backed + * `tmux attach-session -t ` per call. + */ +function attach(ws, runId, { cols, rows }) { + validateRunId(runId); + const pty = spawnImpl("tmux", ["attach-session", "-t", runId], { + name: "xterm-256color", + cols: cols || 80, + rows: rows || 24, + }); + + pty.onData((data) => { + try { + ws.send(data, { binary: true }); + } catch { + /* client gone between data event and send — safe to ignore */ + } + }); + + pty.onExit(({ exitCode }) => { + try { + ws.send(JSON.stringify({ type: "exit", code: exitCode }), { binary: false }); + } catch { + /* ignore */ + } + try { + ws.close(); + } catch { + /* ignore */ + } + }); + + ws.on("message", (data, isBinary) => { + // node-pty attach processes for `tmux attach` are already tolerant of + // resize mid-stream; the `ws` lib passes isBinary either as the second + // callback arg (newer) or via `data.binary` on some transports — this + // helper's own tests exercise the `{binary}` option shape used above. + const binary = typeof isBinary === "boolean" ? isBinary : !!(isBinary && isBinary.binary); + if (binary) { + pty.write(data.toString("utf8")); + return; + } + let msg; + try { + msg = JSON.parse(data.toString("utf8")); + } catch { + return; + } + if (msg && msg.type === "resize" && Number.isFinite(msg.cols) && Number.isFinite(msg.rows)) { + pty.resize(msg.cols, msg.rows); + } + }); + + ws.on("close", () => { + try { + pty.kill(); + } catch { + /* already gone */ + } + }); + + return pty; +} + +// Real spawn implementation — lazy-required so unit tests never load the +// native node-pty addon unless they explicitly opt in. +__setSpawnImpl((...args) => require("node-pty").spawn(...args)); + +module.exports = { attach, validateRunId, __setSpawnImpl }; diff --git a/server/websocket.js b/server/websocket.js index c4fbe4a..6ca4423 100644 --- a/server/websocket.js +++ b/server/websocket.js @@ -59,6 +59,61 @@ function initWebSocket(server) { return wss; } +let ptyWss = null; +const PTY_PATH_RE = /^\/ws-pty\/(ccam-lane-\d+)$/; + +/** + * Second WebSocket server, dedicated to the terminal-run PTY transport + * (`/ws-pty/:runId`). Kept separate from the `/ws` JSON-broadcast path so + * raw binary PTY frames never have to coexist with the typed + * `{type, data, timestamp}` envelope the rest of the app relies on. Reuses + * the exact same auth guard as `/ws`. + * + * `runId` is a path SEGMENT, not a fixed string, so this can't use the `ws` + * library's `{server, path}` shorthand (that option only matches an exact + * string). Instead this server is created with `noServer: true` and the + * upgrade is handled manually — the same `server.on("upgrade", ...)` pattern + * `ws` itself uses internally, just filtered to `/ws-pty/*` first so `/ws`'s + * own upgrade handling (already registered by `initWebSocket`) is untouched. + */ +function initPtyWebSocket(server) { + const { attach } = require("./lib/pty-attach"); + ptyWss = new WebSocketServer({ noServer: true, maxPayload: 1024 * 1024 }); + + ptyWss.on("connection", (ws, runId) => { + try { + attach(ws, runId, { cols: 80, rows: 24 }); + } catch (err) { + try { + ws.close(1008, err.message); + } catch { + /* ignore */ + } + } + }); + + server.on("upgrade", (req, socket, head) => { + const url = new URL(req.url, "http://localhost"); + const match = url.pathname.match(PTY_PATH_RE); + if (!match) return; // not ours — the `/ws` WebSocketServer (registered + // by initWebSocket, also attached to this same http.Server) handles its + // own path independently and ignores upgrades it doesn't match too. + if (!isHostAllowed(req.headers.host)) { + socket.destroy(); + return; + } + if (!isWebSocketAuthorized(req)) { + socket.destroy(); + return; + } + ptyWss.handleUpgrade(req, socket, head, (ws) => { + ptyWss.emit("connection", ws, match[1]); + }); + }); + + return ptyWss; +} + function broadcast(type, data) { if (!wss) return; const message = JSON.stringify({ type, data, timestamp: new Date().toISOString() }); @@ -90,20 +145,36 @@ function getConnectionCount() { * clients first lets the HTTP server drain and close promptly. */ function closeWebSocket() { - if (!wss) return; - wss.clients.forEach((client) => { + if (wss) { + wss.clients.forEach((client) => { + try { + client.terminate(); + } catch { + /* already gone */ + } + }); try { - client.terminate(); + wss.close(); } catch { - /* already gone */ + /* ignore */ } - }); - try { - wss.close(); - } catch { - /* ignore */ + wss = null; + } + if (ptyWss) { + ptyWss.clients.forEach((client) => { + try { + client.terminate(); + } catch { + /* already gone */ + } + }); + try { + ptyWss.close(); + } catch { + /* ignore */ + } + ptyWss = null; } - wss = null; } -module.exports = { initWebSocket, broadcast, getConnectionCount, closeWebSocket }; +module.exports = { initWebSocket, initPtyWebSocket, broadcast, getConnectionCount, closeWebSocket };