server.cjs 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363
  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 = process.env.BRAINSTORM_PORT || (49152 + Math.floor(Math.random() * 16383));
  70. const HOST = process.env.BRAINSTORM_HOST || '127.0.0.1';
  71. const URL_HOST = process.env.BRAINSTORM_URL_HOST || (HOST === '127.0.0.1' ? 'localhost' : HOST);
  72. const SESSION_DIR = process.env.BRAINSTORM_DIR || '/tmp/brainstorm';
  73. const CONTENT_DIR = path.join(SESSION_DIR, 'content');
  74. const STATE_DIR = path.join(SESSION_DIR, 'state');
  75. let ownerPid = process.env.BRAINSTORM_OWNER_PID ? Number(process.env.BRAINSTORM_OWNER_PID) : null;
  76. const MIME_TYPES = {
  77. '.html': 'text/html', '.css': 'text/css', '.js': 'application/javascript',
  78. '.json': 'application/json', '.png': 'image/png', '.jpg': 'image/jpeg',
  79. '.jpeg': 'image/jpeg', '.gif': 'image/gif', '.svg': 'image/svg+xml'
  80. };
  81. // ========== Templates and Constants ==========
  82. const WAITING_PAGE = `<!DOCTYPE html>
  83. <html>
  84. <head><meta charset="utf-8"><title>Brainstorm Companion</title>
  85. <style>body { font-family: system-ui, sans-serif; padding: 2rem; max-width: 800px; margin: 0 auto; }
  86. h1 { color: #333; } p { color: #666; }</style>
  87. </head>
  88. <body><h1>Brainstorm Companion</h1>
  89. <p>Waiting for the agent to push a screen...</p></body></html>`;
  90. const frameTemplate = fs.readFileSync(path.join(__dirname, 'frame-template.html'), 'utf-8');
  91. const helperScript = fs.readFileSync(path.join(__dirname, 'helper.js'), 'utf-8');
  92. const helperInjection = '<script>\n' + helperScript + '\n</script>';
  93. // ========== Helper Functions ==========
  94. function isFullDocument(html) {
  95. const trimmed = html.trimStart().toLowerCase();
  96. return trimmed.startsWith('<!doctype') || trimmed.startsWith('<html');
  97. }
  98. function wrapInFrame(content) {
  99. return frameTemplate.replace('<!-- CONTENT -->', content);
  100. }
  101. function getNewestScreen() {
  102. const files = fs.readdirSync(CONTENT_DIR)
  103. .filter(f => f.endsWith('.html'))
  104. .map(f => {
  105. const fp = path.join(CONTENT_DIR, f);
  106. return { path: fp, mtime: fs.statSync(fp).mtime.getTime() };
  107. })
  108. .sort((a, b) => b.mtime - a.mtime);
  109. return files.length > 0 ? files[0].path : null;
  110. }
  111. // ========== HTTP Request Handler ==========
  112. function handleRequest(req, res) {
  113. touchActivity();
  114. if (req.method === 'GET' && req.url === '/') {
  115. const screenFile = getNewestScreen();
  116. let html = screenFile
  117. ? (raw => isFullDocument(raw) ? raw : wrapInFrame(raw))(fs.readFileSync(screenFile, 'utf-8'))
  118. : WAITING_PAGE;
  119. if (html.includes('</body>')) {
  120. html = html.replace('</body>', helperInjection + '\n</body>');
  121. } else {
  122. html += helperInjection;
  123. }
  124. res.writeHead(200, { 'Content-Type': 'text/html; charset=utf-8' });
  125. res.end(html);
  126. } else if (req.method === 'GET' && req.url.startsWith('/files/')) {
  127. const fileName = req.url.slice(7);
  128. const filePath = path.join(CONTENT_DIR, path.basename(fileName));
  129. if (!fs.existsSync(filePath)) {
  130. res.writeHead(404);
  131. res.end('Not found');
  132. return;
  133. }
  134. const ext = path.extname(filePath).toLowerCase();
  135. const contentType = MIME_TYPES[ext] || 'application/octet-stream';
  136. res.writeHead(200, { 'Content-Type': contentType });
  137. res.end(fs.readFileSync(filePath));
  138. } else {
  139. res.writeHead(404);
  140. res.end('Not found');
  141. }
  142. }
  143. // ========== WebSocket Connection Handling ==========
  144. const clients = new Set();
  145. function handleUpgrade(req, socket) {
  146. const key = req.headers['sec-websocket-key'];
  147. if (!key) { socket.destroy(); return; }
  148. const accept = computeAcceptKey(key);
  149. socket.write(
  150. 'HTTP/1.1 101 Switching Protocols\r\n' +
  151. 'Upgrade: websocket\r\n' +
  152. 'Connection: Upgrade\r\n' +
  153. 'Sec-WebSocket-Accept: ' + accept + '\r\n\r\n'
  154. );
  155. let buffer = Buffer.alloc(0);
  156. clients.add(socket);
  157. socket.on('data', (chunk) => {
  158. buffer = Buffer.concat([buffer, chunk]);
  159. while (buffer.length > 0) {
  160. let result;
  161. try {
  162. result = decodeFrame(buffer);
  163. } catch (e) {
  164. socket.end(encodeFrame(OPCODES.CLOSE, Buffer.alloc(0)));
  165. clients.delete(socket);
  166. return;
  167. }
  168. if (!result) break;
  169. buffer = buffer.slice(result.bytesConsumed);
  170. switch (result.opcode) {
  171. case OPCODES.TEXT:
  172. handleMessage(result.payload.toString());
  173. break;
  174. case OPCODES.CLOSE:
  175. socket.end(encodeFrame(OPCODES.CLOSE, Buffer.alloc(0)));
  176. clients.delete(socket);
  177. return;
  178. case OPCODES.PING:
  179. socket.write(encodeFrame(OPCODES.PONG, result.payload));
  180. break;
  181. case OPCODES.PONG:
  182. break;
  183. default: {
  184. const closeBuf = Buffer.alloc(2);
  185. closeBuf.writeUInt16BE(1003);
  186. socket.end(encodeFrame(OPCODES.CLOSE, closeBuf));
  187. clients.delete(socket);
  188. return;
  189. }
  190. }
  191. }
  192. });
  193. socket.on('close', () => clients.delete(socket));
  194. socket.on('error', () => clients.delete(socket));
  195. }
  196. function handleMessage(text) {
  197. let event;
  198. try {
  199. event = JSON.parse(text);
  200. } catch (e) {
  201. console.error('Failed to parse WebSocket message:', e.message);
  202. return;
  203. }
  204. touchActivity();
  205. console.log(JSON.stringify({ source: 'user-event', ...event }));
  206. if (event.choice) {
  207. const eventsFile = path.join(STATE_DIR, 'events');
  208. fs.appendFileSync(eventsFile, JSON.stringify(event) + '\n');
  209. }
  210. }
  211. function broadcast(msg) {
  212. const frame = encodeFrame(OPCODES.TEXT, Buffer.from(JSON.stringify(msg)));
  213. for (const socket of clients) {
  214. try { socket.write(frame); } catch (e) { clients.delete(socket); }
  215. }
  216. }
  217. // ========== Activity Tracking ==========
  218. const IDLE_TIMEOUT_MS = 30 * 60 * 1000; // 30 minutes
  219. let lastActivity = Date.now();
  220. function touchActivity() {
  221. lastActivity = Date.now();
  222. }
  223. // ========== File Watching ==========
  224. const debounceTimers = new Map();
  225. // ========== Server Startup ==========
  226. function startServer() {
  227. if (!fs.existsSync(CONTENT_DIR)) fs.mkdirSync(CONTENT_DIR, { recursive: true });
  228. if (!fs.existsSync(STATE_DIR)) fs.mkdirSync(STATE_DIR, { recursive: true });
  229. // Track known files to distinguish new screens from updates.
  230. // macOS fs.watch reports 'rename' for both new files and overwrites,
  231. // so we can't rely on eventType alone.
  232. const knownFiles = new Set(
  233. fs.readdirSync(CONTENT_DIR).filter(f => f.endsWith('.html'))
  234. );
  235. const server = http.createServer(handleRequest);
  236. server.on('upgrade', handleUpgrade);
  237. const watcher = fs.watch(CONTENT_DIR, (eventType, filename) => {
  238. if (!filename || !filename.endsWith('.html')) return;
  239. if (debounceTimers.has(filename)) clearTimeout(debounceTimers.get(filename));
  240. debounceTimers.set(filename, setTimeout(() => {
  241. debounceTimers.delete(filename);
  242. const filePath = path.join(CONTENT_DIR, filename);
  243. if (!fs.existsSync(filePath)) return; // file was deleted
  244. touchActivity();
  245. if (!knownFiles.has(filename)) {
  246. knownFiles.add(filename);
  247. const eventsFile = path.join(STATE_DIR, 'events');
  248. if (fs.existsSync(eventsFile)) fs.unlinkSync(eventsFile);
  249. console.log(JSON.stringify({ type: 'screen-added', file: filePath }));
  250. } else {
  251. console.log(JSON.stringify({ type: 'screen-updated', file: filePath }));
  252. }
  253. broadcast({ type: 'reload' });
  254. }, 100));
  255. });
  256. watcher.on('error', (err) => console.error('fs.watch error:', err.message));
  257. function shutdown(reason) {
  258. console.log(JSON.stringify({ type: 'server-stopped', reason }));
  259. const infoFile = path.join(STATE_DIR, 'server-info');
  260. if (fs.existsSync(infoFile)) fs.unlinkSync(infoFile);
  261. fs.writeFileSync(
  262. path.join(STATE_DIR, 'server-stopped'),
  263. JSON.stringify({ reason, timestamp: Date.now() }) + '\n'
  264. );
  265. watcher.close();
  266. clearInterval(lifecycleCheck);
  267. server.close(() => process.exit(0));
  268. }
  269. function ownerAlive() {
  270. if (!ownerPid) return true;
  271. try { process.kill(ownerPid, 0); return true; } catch (e) { return e.code === 'EPERM'; }
  272. }
  273. // Check every 60s: exit if owner process died or idle for 30 minutes
  274. const lifecycleCheck = setInterval(() => {
  275. if (!ownerAlive()) shutdown('owner process exited');
  276. else if (Date.now() - lastActivity > IDLE_TIMEOUT_MS) shutdown('idle timeout');
  277. }, 60 * 1000);
  278. lifecycleCheck.unref();
  279. // Validate owner PID at startup. If it's already dead, the PID resolution
  280. // was wrong (common on WSL, Tailscale SSH, and cross-user scenarios).
  281. // Disable monitoring and rely on the idle timeout instead.
  282. if (ownerPid) {
  283. try { process.kill(ownerPid, 0); }
  284. catch (e) {
  285. if (e.code !== 'EPERM') {
  286. console.log(JSON.stringify({ type: 'owner-pid-invalid', pid: ownerPid, reason: 'dead at startup' }));
  287. ownerPid = null;
  288. }
  289. }
  290. }
  291. server.listen(PORT, HOST, () => {
  292. const info = JSON.stringify({
  293. type: 'server-started', port: Number(PORT), host: HOST,
  294. url_host: URL_HOST, url: 'http://' + URL_HOST + ':' + PORT,
  295. screen_dir: CONTENT_DIR, state_dir: STATE_DIR
  296. });
  297. console.log(info);
  298. fs.writeFileSync(path.join(STATE_DIR, 'server-info'), info + '\n');
  299. });
  300. }
  301. if (require.main === module) {
  302. startServer();
  303. }
  304. module.exports = { computeAcceptKey, encodeFrame, decodeFrame, OPCODES, MAX_FRAME_PAYLOAD_BYTES };