1
0

server.cjs 18 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519
  1. const crypto = require('crypto');
  2. const http = require('http');
  3. const fs = require('fs');
  4. const path = require('path');
  5. // ========== WebSocket Protocol (RFC 6455) ==========
  6. const OPCODES = { TEXT: 0x01, CLOSE: 0x08, PING: 0x09, PONG: 0x0A };
  7. const WS_MAGIC = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11';
  8. const MAX_FRAME_PAYLOAD_BYTES = 10 * 1024 * 1024;
  9. function computeAcceptKey(clientKey) {
  10. return crypto.createHash('sha1').update(clientKey + WS_MAGIC).digest('base64');
  11. }
  12. function encodeFrame(opcode, payload) {
  13. const fin = 0x80;
  14. const len = payload.length;
  15. let header;
  16. if (len < 126) {
  17. header = Buffer.alloc(2);
  18. header[0] = fin | opcode;
  19. header[1] = len;
  20. } else if (len < 65536) {
  21. header = Buffer.alloc(4);
  22. header[0] = fin | opcode;
  23. header[1] = 126;
  24. header.writeUInt16BE(len, 2);
  25. } else {
  26. header = Buffer.alloc(10);
  27. header[0] = fin | opcode;
  28. header[1] = 127;
  29. header.writeBigUInt64BE(BigInt(len), 2);
  30. }
  31. return Buffer.concat([header, payload]);
  32. }
  33. function decodeFrame(buffer) {
  34. if (buffer.length < 2) return null;
  35. const secondByte = buffer[1];
  36. const opcode = buffer[0] & 0x0F;
  37. const masked = (secondByte & 0x80) !== 0;
  38. let payloadLen = secondByte & 0x7F;
  39. let offset = 2;
  40. if (!masked) throw new Error('Client frames must be masked');
  41. if (payloadLen === 126) {
  42. if (buffer.length < 4) return null;
  43. payloadLen = buffer.readUInt16BE(2);
  44. offset = 4;
  45. } else if (payloadLen === 127) {
  46. if (buffer.length < 10) return null;
  47. const extendedLen = buffer.readBigUInt64BE(2);
  48. if (extendedLen > BigInt(MAX_FRAME_PAYLOAD_BYTES)) {
  49. throw new Error('WebSocket frame payload exceeds maximum allowed size');
  50. }
  51. payloadLen = Number(extendedLen);
  52. offset = 10;
  53. }
  54. if (payloadLen > MAX_FRAME_PAYLOAD_BYTES) {
  55. throw new Error('WebSocket frame payload exceeds maximum allowed size');
  56. }
  57. const maskOffset = offset;
  58. const dataOffset = offset + 4;
  59. const totalLen = dataOffset + payloadLen;
  60. if (buffer.length < totalLen) return null;
  61. const mask = buffer.slice(maskOffset, dataOffset);
  62. const data = Buffer.alloc(payloadLen);
  63. for (let i = 0; i < payloadLen; i++) {
  64. data[i] = buffer[dataOffset + i] ^ mask[i % 4];
  65. }
  66. return { opcode, payload: data, bytesConsumed: totalLen };
  67. }
  68. // ========== Configuration ==========
  69. const PORT_FILE = process.env.BRAINSTORM_PORT_FILE || null;
  70. const randomPort = () => 49152 + Math.floor(Math.random() * 16383);
  71. // Prefer an explicit port, else the port this session last bound (so a restart
  72. // reuses it and an already-open browser tab reconnects), else a random high port.
  73. function preferredPort() {
  74. if (process.env.BRAINSTORM_PORT) return Number(process.env.BRAINSTORM_PORT);
  75. if (PORT_FILE) {
  76. try {
  77. const p = Number(fs.readFileSync(PORT_FILE, 'utf-8').trim());
  78. if (Number.isInteger(p) && p > 1023 && p < 65536) return p;
  79. } catch (e) { /* no prior port recorded */ }
  80. }
  81. return randomPort();
  82. }
  83. let PORT = preferredPort();
  84. const HOST = process.env.BRAINSTORM_HOST || '127.0.0.1';
  85. const URL_HOST = process.env.BRAINSTORM_URL_HOST || (HOST === '127.0.0.1' ? 'localhost' : HOST);
  86. const SESSION_DIR = process.env.BRAINSTORM_DIR || '/tmp/brainstorm';
  87. const CONTENT_DIR = path.join(SESSION_DIR, 'content');
  88. const STATE_DIR = path.join(SESSION_DIR, 'state');
  89. let ownerPid = process.env.BRAINSTORM_OWNER_PID ? Number(process.env.BRAINSTORM_OWNER_PID) : null;
  90. // Per-session secret key. The companion is reachable by any local browser tab
  91. // and, when bound to a non-loopback host, by any host that can route to it.
  92. // The key authenticates the real client uniformly across loopback, tunnel, and
  93. // remote binds — and defeats DNS rebinding — where a Host/Origin allowlist
  94. // cannot. It rides the served URL as ?key= and is mirrored into a cookie on
  95. // first load so same-origin subresources and the WebSocket carry it for free.
  96. const TOKEN = process.env.BRAINSTORM_TOKEN || crypto.randomBytes(32).toString('hex');
  97. const COOKIE_NAME = 'brainstorm-key-' + PORT;
  98. const MIME_TYPES = {
  99. '.html': 'text/html', '.css': 'text/css', '.js': 'application/javascript',
  100. '.json': 'application/json', '.png': 'image/png', '.jpg': 'image/jpeg',
  101. '.jpeg': 'image/jpeg', '.gif': 'image/gif', '.svg': 'image/svg+xml'
  102. };
  103. // ========== Templates and Constants ==========
  104. const WAITING_PAGE = `<!DOCTYPE html>
  105. <html>
  106. <head><meta charset="utf-8"><title>Brainstorm Companion</title>
  107. <style>body { font-family: system-ui, sans-serif; padding: 2rem; max-width: 800px; margin: 0 auto; }
  108. h1 { color: #333; } p { color: #666; }</style>
  109. </head>
  110. <body><h1>Brainstorm Companion</h1>
  111. <p>Waiting for the agent to push a screen...</p></body></html>`;
  112. const FORBIDDEN_PAGE = `<!DOCTYPE html>
  113. <html>
  114. <head><meta charset="utf-8"><title>Session key required</title>
  115. <style>body { font-family: system-ui, sans-serif; padding: 2rem; max-width: 800px; margin: 0 auto; }
  116. h1 { color: #333; } p { color: #666; } code { background: #f0f0f0; padding: 0.1em 0.3em; border-radius: 4px; }</style>
  117. </head>
  118. <body><h1>Session key required</h1>
  119. <p>This page needs the full URL your coding agent gave you, including the
  120. <code>?key=&hellip;</code> part. Copy the complete URL and open it again.</p></body></html>`;
  121. const frameTemplate = fs.readFileSync(path.join(__dirname, 'frame-template.html'), 'utf-8');
  122. const helperScript = fs.readFileSync(path.join(__dirname, 'helper.js'), 'utf-8');
  123. const helperInjection = '<script>\n' + helperScript + '\n</script>';
  124. // ========== Helper Functions ==========
  125. function isFullDocument(html) {
  126. const trimmed = html.trimStart().toLowerCase();
  127. return trimmed.startsWith('<!doctype') || trimmed.startsWith('<html');
  128. }
  129. function wrapInFrame(content) {
  130. return frameTemplate.replace('<!-- CONTENT -->', content);
  131. }
  132. function getNewestScreen() {
  133. const files = fs.readdirSync(CONTENT_DIR)
  134. .filter(f => !f.startsWith('.') && f.endsWith('.html'))
  135. .map(f => {
  136. const fp = path.join(CONTENT_DIR, f);
  137. return { path: fp, mtime: fs.statSync(fp).mtime.getTime() };
  138. })
  139. .sort((a, b) => b.mtime - a.mtime);
  140. return files.length > 0 ? files[0].path : null;
  141. }
  142. // ========== Authentication ==========
  143. function timingSafeEqualStr(a, b) {
  144. const ab = Buffer.from(String(a));
  145. const bb = Buffer.from(String(b));
  146. if (ab.length !== bb.length) return false;
  147. return crypto.timingSafeEqual(ab, bb);
  148. }
  149. function parseCookies(header) {
  150. const out = {};
  151. if (!header) return out;
  152. for (const part of header.split(';')) {
  153. const eq = part.indexOf('=');
  154. if (eq < 0) continue;
  155. out[part.slice(0, eq).trim()] = part.slice(eq + 1).trim();
  156. }
  157. return out;
  158. }
  159. // A request is authorized if it carries the session key as ?key= or as the
  160. // session cookie. Both are compared in constant time.
  161. function isAuthorized(req) {
  162. const q = req.url.indexOf('?');
  163. if (q >= 0) {
  164. const key = new URLSearchParams(req.url.slice(q + 1)).get('key');
  165. if (key && timingSafeEqualStr(key, TOKEN)) return true;
  166. }
  167. const cookie = parseCookies(req.headers['cookie'])[COOKIE_NAME];
  168. if (cookie && timingSafeEqualStr(cookie, TOKEN)) return true;
  169. return false;
  170. }
  171. function pathnameOf(url) {
  172. const q = url.indexOf('?');
  173. return q >= 0 ? url.slice(0, q) : url;
  174. }
  175. // ========== HTTP Request Handler ==========
  176. function handleRequest(req, res) {
  177. touchActivity();
  178. if (!isAuthorized(req)) {
  179. res.writeHead(403, { 'Content-Type': 'text/html; charset=utf-8' });
  180. res.end(FORBIDDEN_PAGE);
  181. return;
  182. }
  183. // Mirror the key into a cookie so same-origin subresources (/files/*) and the
  184. // WebSocket handshake carry it automatically, whatever URL style the agent
  185. // writes. SameSite=Strict: a cross-site page can neither read the key nor ride
  186. // the cookie; HttpOnly: page scripts can't exfiltrate it.
  187. res.setHeader('Set-Cookie',
  188. COOKIE_NAME + '=' + TOKEN + '; HttpOnly; SameSite=Strict; Path=/');
  189. const pathname = pathnameOf(req.url);
  190. if (req.method === 'GET' && pathname === '/') {
  191. const screenFile = getNewestScreen();
  192. let html = screenFile
  193. ? (raw => isFullDocument(raw) ? raw : wrapInFrame(raw))(fs.readFileSync(screenFile, 'utf-8'))
  194. : WAITING_PAGE;
  195. if (html.includes('</body>')) {
  196. html = html.replace('</body>', helperInjection + '\n</body>');
  197. } else {
  198. html += helperInjection;
  199. }
  200. res.writeHead(200, { 'Content-Type': 'text/html; charset=utf-8' });
  201. res.end(html);
  202. } else if (req.method === 'GET' && pathname.startsWith('/files/')) {
  203. const fileName = path.basename(pathname.slice(7));
  204. const filePath = path.join(CONTENT_DIR, fileName);
  205. if (fileName.startsWith('.') || !fs.existsSync(filePath)) {
  206. res.writeHead(404);
  207. res.end('Not found');
  208. return;
  209. }
  210. const ext = path.extname(filePath).toLowerCase();
  211. const contentType = MIME_TYPES[ext] || 'application/octet-stream';
  212. res.writeHead(200, { 'Content-Type': contentType });
  213. res.end(fs.readFileSync(filePath));
  214. } else {
  215. res.writeHead(404);
  216. res.end('Not found');
  217. }
  218. }
  219. // ========== WebSocket Connection Handling ==========
  220. const clients = new Set();
  221. function handleUpgrade(req, socket) {
  222. if (!isAuthorized(req)) { socket.destroy(); return; }
  223. const key = req.headers['sec-websocket-key'];
  224. if (!key) { socket.destroy(); return; }
  225. const accept = computeAcceptKey(key);
  226. socket.write(
  227. 'HTTP/1.1 101 Switching Protocols\r\n' +
  228. 'Upgrade: websocket\r\n' +
  229. 'Connection: Upgrade\r\n' +
  230. 'Sec-WebSocket-Accept: ' + accept + '\r\n\r\n'
  231. );
  232. let buffer = Buffer.alloc(0);
  233. clients.add(socket);
  234. socket.on('data', (chunk) => {
  235. buffer = Buffer.concat([buffer, chunk]);
  236. while (buffer.length > 0) {
  237. let result;
  238. try {
  239. result = decodeFrame(buffer);
  240. } catch (e) {
  241. socket.end(encodeFrame(OPCODES.CLOSE, Buffer.alloc(0)));
  242. clients.delete(socket);
  243. return;
  244. }
  245. if (!result) break;
  246. buffer = buffer.slice(result.bytesConsumed);
  247. switch (result.opcode) {
  248. case OPCODES.TEXT:
  249. handleMessage(result.payload.toString());
  250. break;
  251. case OPCODES.CLOSE:
  252. socket.end(encodeFrame(OPCODES.CLOSE, Buffer.alloc(0)));
  253. clients.delete(socket);
  254. return;
  255. case OPCODES.PING:
  256. socket.write(encodeFrame(OPCODES.PONG, result.payload));
  257. break;
  258. case OPCODES.PONG:
  259. break;
  260. default: {
  261. const closeBuf = Buffer.alloc(2);
  262. closeBuf.writeUInt16BE(1003);
  263. socket.end(encodeFrame(OPCODES.CLOSE, closeBuf));
  264. clients.delete(socket);
  265. return;
  266. }
  267. }
  268. }
  269. });
  270. socket.on('close', () => clients.delete(socket));
  271. socket.on('error', () => clients.delete(socket));
  272. }
  273. function handleMessage(text) {
  274. let event;
  275. try {
  276. event = JSON.parse(text);
  277. } catch (e) {
  278. console.error('Failed to parse WebSocket message:', e.message);
  279. return;
  280. }
  281. touchActivity();
  282. console.log(JSON.stringify({ source: 'user-event', ...event }));
  283. if (event && event.choice) {
  284. const eventsFile = path.join(STATE_DIR, 'events');
  285. fs.appendFileSync(eventsFile, JSON.stringify(event) + '\n');
  286. }
  287. }
  288. function broadcast(msg) {
  289. const frame = encodeFrame(OPCODES.TEXT, Buffer.from(JSON.stringify(msg)));
  290. for (const socket of clients) {
  291. try { socket.write(frame); } catch (e) { clients.delete(socket); }
  292. }
  293. }
  294. // Best-effort: open the user's browser the first time a screen is actually ready
  295. // to show. Skips when disabled, on a non-loopback (remote) bind, or when a
  296. // browser is already connected. Override the launcher with BRAINSTORM_OPEN_CMD.
  297. let browserOpened = false;
  298. function maybeOpenBrowser() {
  299. if (browserOpened) return;
  300. browserOpened = true;
  301. if (!process.env.BRAINSTORM_OPEN) return; // opt-in: only after the user approves the companion
  302. if (HOST !== '127.0.0.1' && HOST !== 'localhost') return;
  303. if (clients.size > 0) return; // the user already opened it
  304. const url = 'http://' + URL_HOST + ':' + PORT;
  305. const cp = require('child_process');
  306. // Operator-provided launcher: run as given (this env var is trusted operator input).
  307. if (process.env.BRAINSTORM_OPEN_CMD) {
  308. try { cp.exec(process.env.BRAINSTORM_OPEN_CMD + ' ' + JSON.stringify(url), () => {}); } catch (e) { /* best effort */ }
  309. return;
  310. }
  311. // Platform launchers: pass the URL as an argv element via execFile (no shell),
  312. // so a url-host containing shell metacharacters can't inject a command.
  313. const isWSL = process.platform === 'linux' && /microsoft/i.test(require('os').release());
  314. let bin, args;
  315. if (process.platform === 'darwin') { bin = 'open'; args = [url]; }
  316. else if (process.platform === 'win32' || isWSL) { bin = 'cmd.exe'; args = ['/c', 'start', '', url]; }
  317. else if (process.env.DISPLAY || process.env.WAYLAND_DISPLAY) { bin = 'xdg-open'; args = [url]; }
  318. else return; // headless: nothing to open
  319. try { cp.execFile(bin, args, () => {}); } catch (e) { /* best effort */ }
  320. }
  321. // ========== Activity Tracking ==========
  322. // Idle timeout: shut down after this long with no activity. Default 4 hours;
  323. // override with BRAINSTORM_IDLE_TIMEOUT_MS (start-server.sh: --idle-timeout-minutes).
  324. const IDLE_TIMEOUT_MS = (() => {
  325. const ms = Number(process.env.BRAINSTORM_IDLE_TIMEOUT_MS);
  326. return Number.isFinite(ms) && ms > 0 ? ms : 4 * 60 * 60 * 1000;
  327. })();
  328. // How often the watchdog checks for owner-death / idleness. Configurable mainly
  329. // so tests can run fast; production default is 60s.
  330. const LIFECYCLE_CHECK_MS = (() => {
  331. const ms = Number(process.env.BRAINSTORM_LIFECYCLE_CHECK_MS);
  332. return Number.isFinite(ms) && ms > 0 ? ms : 60 * 1000;
  333. })();
  334. let lastActivity = Date.now();
  335. function touchActivity() {
  336. lastActivity = Date.now();
  337. }
  338. // ========== File Watching ==========
  339. const debounceTimers = new Map();
  340. // ========== Server Startup ==========
  341. function startServer() {
  342. if (!fs.existsSync(CONTENT_DIR)) fs.mkdirSync(CONTENT_DIR, { recursive: true });
  343. if (!fs.existsSync(STATE_DIR)) fs.mkdirSync(STATE_DIR, { recursive: true });
  344. // Track known files to distinguish new screens from updates.
  345. // macOS fs.watch reports 'rename' for both new files and overwrites,
  346. // so we can't rely on eventType alone.
  347. const knownFiles = new Set(
  348. fs.readdirSync(CONTENT_DIR).filter(f => !f.startsWith('.') && f.endsWith('.html'))
  349. );
  350. const server = http.createServer(handleRequest);
  351. server.on('upgrade', handleUpgrade);
  352. const watcher = fs.watch(CONTENT_DIR, (eventType, filename) => {
  353. if (!filename || filename.startsWith('.') || !filename.endsWith('.html')) return;
  354. if (debounceTimers.has(filename)) clearTimeout(debounceTimers.get(filename));
  355. debounceTimers.set(filename, setTimeout(() => {
  356. debounceTimers.delete(filename);
  357. const filePath = path.join(CONTENT_DIR, filename);
  358. if (!fs.existsSync(filePath)) return; // file was deleted
  359. touchActivity();
  360. if (!knownFiles.has(filename)) {
  361. knownFiles.add(filename);
  362. const eventsFile = path.join(STATE_DIR, 'events');
  363. if (fs.existsSync(eventsFile)) fs.unlinkSync(eventsFile);
  364. console.log(JSON.stringify({ type: 'screen-added', file: filePath }));
  365. maybeOpenBrowser();
  366. } else {
  367. console.log(JSON.stringify({ type: 'screen-updated', file: filePath }));
  368. }
  369. broadcast({ type: 'reload' });
  370. }, 100));
  371. });
  372. watcher.on('error', (err) => console.error('fs.watch error:', err.message));
  373. function shutdown(reason) {
  374. console.log(JSON.stringify({ type: 'server-stopped', reason }));
  375. const infoFile = path.join(STATE_DIR, 'server-info');
  376. if (fs.existsSync(infoFile)) fs.unlinkSync(infoFile);
  377. fs.writeFileSync(
  378. path.join(STATE_DIR, 'server-stopped'),
  379. JSON.stringify({ reason, timestamp: Date.now() }) + '\n'
  380. );
  381. watcher.close();
  382. clearInterval(lifecycleCheck);
  383. // Close any upgraded WebSocket sockets so server.close() can complete and
  384. // the process actually exits instead of lingering on an open connection.
  385. for (const socket of clients) {
  386. try { socket.destroy(); } catch (e) { /* already gone */ }
  387. }
  388. server.close(() => process.exit(0));
  389. }
  390. function ownerAlive() {
  391. if (!ownerPid) return true;
  392. try { process.kill(ownerPid, 0); return true; } catch (e) { return e.code === 'EPERM'; }
  393. }
  394. // Periodically exit if the owner process died or we've been idle too long.
  395. const lifecycleCheck = setInterval(() => {
  396. if (!ownerAlive()) shutdown('owner process exited');
  397. else if (Date.now() - lastActivity > IDLE_TIMEOUT_MS) shutdown('idle timeout');
  398. }, LIFECYCLE_CHECK_MS);
  399. lifecycleCheck.unref();
  400. // Validate owner PID at startup. If it's already dead, the PID resolution
  401. // was wrong (common on WSL, Tailscale SSH, and cross-user scenarios).
  402. // Disable monitoring and rely on the idle timeout instead.
  403. if (ownerPid) {
  404. try { process.kill(ownerPid, 0); }
  405. catch (e) {
  406. if (e.code !== 'EPERM') {
  407. console.log(JSON.stringify({ type: 'owner-pid-invalid', pid: ownerPid, reason: 'dead at startup' }));
  408. ownerPid = null;
  409. }
  410. }
  411. }
  412. // If the preferred port is already taken (e.g. a previous server is still
  413. // alive), fall back to a random port once instead of failing.
  414. let triedFallback = false;
  415. function onListen() {
  416. // Record the bound port so the next restart of this session reuses it — but
  417. // ONLY when we got our preferred port. On a fallback we bound a *different*
  418. // port because someone else holds the preferred one; persisting it would
  419. // overwrite the shared .last-port and strand that other session's open tab.
  420. if (PORT_FILE && !triedFallback) {
  421. try { fs.writeFileSync(PORT_FILE, String(PORT)); } catch (e) { /* best effort */ }
  422. }
  423. const info = JSON.stringify({
  424. type: 'server-started', port: Number(PORT), host: HOST,
  425. url_host: URL_HOST, url: 'http://' + URL_HOST + ':' + PORT + '/?key=' + TOKEN,
  426. screen_dir: CONTENT_DIR, state_dir: STATE_DIR, idle_timeout_ms: IDLE_TIMEOUT_MS
  427. });
  428. console.log(info);
  429. fs.writeFileSync(path.join(STATE_DIR, 'server-info'), info + '\n');
  430. }
  431. server.on('error', (err) => {
  432. if (err.code === 'EADDRINUSE' && !triedFallback) {
  433. triedFallback = true;
  434. PORT = randomPort();
  435. server.listen(PORT, HOST, onListen);
  436. } else {
  437. console.error('Server failed to bind:', err.message);
  438. process.exit(1);
  439. }
  440. });
  441. server.listen(PORT, HOST, onListen);
  442. }
  443. if (require.main === module) {
  444. startServer();
  445. }
  446. module.exports = { computeAcceptKey, encodeFrame, decodeFrame, OPCODES, MAX_FRAME_PAYLOAD_BYTES };