lvmh/plugin/mini-daemon.ts

/**
 * Dev-only minimal WebSocket "daemon" used by smoke.ts and e2e.ts.
 * Zero npm dependencies: hand-rolled WS handshake + frame codec.
 * Not covered by tsc beyond what tsconfig includes; Node type-strips it.
 */

import { createHash } from "node:crypto";
import * as http from "node:http";
import type { AddressInfo } from "node:net";
import type { Duplex } from "node:stream";

export const WS_GUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";

export interface ClientFrame {
	v: number;
	sessionId: string;
	seq: number;
	ts: number;
	type: string;
	[key: string]: unknown;
}

export interface MiniDaemon {
	url: string;
	frames: ClientFrame[];
	authHeaders: string[];
	connections(): number;
	pushAll(text: string): void;
	dropConnections(): void;
	/** Pause reads on all live sockets: simulate a daemon wedged without TCP close. */
	wedge(): void;
	unwedge(): void;
	close(): void;
}

export function check(name: string, ok: boolean, detail = ""): void {
	if (ok) console.log(`ok   ${name}`);
	else console.log(`FAIL ${name} ${detail}`);
}

export function sleep(ms: number): Promise<void> {
	return new Promise((resolve) => setTimeout(resolve, ms));
}

export async function waitFor(
	cond: () => boolean,
	timeoutMs: number,
	stepMs = 25,
): Promise<boolean> {
	const deadline = Date.now() + timeoutMs;
	while (Date.now() < deadline) {
		if (cond()) return true;
		await sleep(stepMs);
	}
	return cond();
}

function encodeTextFrame(text: string): Buffer {
	const payload = Buffer.from(text, "utf8");
	const len = payload.length;
	let header: Buffer;
	if (len < 126) header = Buffer.from([0x81, len]);
	else if (len < 65536) {
		header = Buffer.alloc(4);
		header[0] = 0x81;
		header[1] = 126;
		header.writeUInt16BE(len, 2);
	} else {
		header = Buffer.alloc(10);
		header[0] = 0x81;
		header[1] = 127;
		header.writeBigUInt64BE(BigInt(len), 2);
	}
	return Buffer.concat([header, payload]);
}

function decodeFrames(chunk: Buffer): {
	frames: Array<{ opcode: number; data: Buffer }>;
	consumed: number;
} {
	const frames: Array<{ opcode: number; data: Buffer }> = [];
	let offset = 0;
	while (offset + 2 <= chunk.length) {
		const opcode = chunk[offset] & 0x0f;
		const masked = (chunk[offset + 1] & 0x80) !== 0;
		let len = chunk[offset + 1] & 0x7f;
		let cursor = offset + 2;
		if (len === 126) {
			if (cursor + 2 > chunk.length) break;
			len = chunk.readUInt16BE(cursor);
			cursor += 2;
		} else if (len === 127) {
			if (cursor + 8 > chunk.length) break;
			len = Number(chunk.readBigUInt64BE(cursor));
			cursor += 8;
		}
		let mask: Buffer | null = null;
		if (masked) {
			if (cursor + 4 > chunk.length) break;
			mask = chunk.subarray(cursor, cursor + 4);
			cursor += 4;
		}
		if (cursor + len > chunk.length) break;
		let data = chunk.subarray(cursor, cursor + len);
		if (mask !== null) {
			const unmasked = Buffer.allocUnsafe(len);
			for (let i = 0; i < len; i++) unmasked[i] = data[i] ^ mask[i % 4];
			data = unmasked;
		}
		frames.push({ opcode, data });
		offset = cursor + len;
	}
	return { frames, consumed: offset };
}

export function startMiniDaemon(
	getLastSeq: () => number,
	opts: { suppressWelcome?: boolean } = {},
): Promise<MiniDaemon> {
	const server = http.createServer();
	const sockets = new Set<Duplex>();
	const daemon: MiniDaemon = {
		url: "",
		frames: [],
		authHeaders: [],
		connections: () => sockets.size,
		pushAll(text: string) {
			const frame = encodeTextFrame(text);
			for (const s of sockets) s.write(frame);
		},
		dropConnections() {
			for (const s of sockets) s.destroy();
			sockets.clear();
		},
		wedge() {
			for (const s of sockets) s.pause();
		},
		unwedge() {
			for (const s of sockets) s.resume();
		},
		close() {
			daemon.dropConnections();
			server.close();
		},
	};
	server.on("upgrade", (req: http.IncomingMessage, socket: Duplex) => {
		daemon.authHeaders.push(String(req.headers.authorization ?? ""));
		const key = String(req.headers["sec-websocket-key"] ?? "");
		const accept = createHash("sha1")
			.update(key + WS_GUID)
			.digest("base64");
		socket.write(
			"HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n" +
				`Sec-WebSocket-Accept: ${accept}\r\n\r\n`,
		);
		sockets.add(socket);
		// Frames can split across TCP chunks: buffer per socket and decode only
		// complete frames, retaining the remainder for the next chunk.
		let recvBuf = Buffer.alloc(0);
		socket.on("data", (chunk: Buffer) => {
			recvBuf = Buffer.concat([recvBuf, chunk]);
			const { frames: decoded, consumed } = decodeFrames(recvBuf);
			recvBuf = recvBuf.subarray(consumed);
			for (const f of decoded) {
				if (f.opcode === 0x8) {
					// proper close handshake: echo close frame, then end the socket
					socket.write(Buffer.from([0x88, 0x02, 0x03, 0xe8]));
					socket.end();
					continue;
				}
				if (f.opcode !== 0x1) continue;
				try {
					const frame = JSON.parse(f.data.toString("utf8")) as ClientFrame;
					daemon.frames.push(frame);
					if (frame.type === "hello" && !opts.suppressWelcome) {
						socket.write(
							encodeTextFrame(
								JSON.stringify({
									v: 1,
									type: "welcome",
									sessionId: frame.sessionId,
									seq: 0,
									ts: Date.now(),
									lastSeq: getLastSeq(),
								}),
							),
						);
					}
				} catch {
					// malformed client frame: ignore
				}
			}
		});
		socket.on("close", () => sockets.delete(socket));
		socket.on("error", () => sockets.delete(socket));
	});
	return new Promise((resolve) => {
		server.listen(0, "127.0.0.1", () => {
			daemon.url = `ws://127.0.0.1:${(server.address() as AddressInfo).port}/agent/ws`;
			resolve(daemon);
		});
	});
}