pi-gui/src/bridge/unix-server.js

import { createServer } from "node:net";
import { StringDecoder } from "node:string_decoder";
import { chmod, lstat, unlink } from "node:fs/promises";
import {
	MAX_FRAME_BYTES,
	ProtocolError,
	encodeFrame,
	errorResponse,
	parseRequestFrame,
	successResponse,
} from "../protocol/index.js";

const OWNER_ONLY_MODE = 0o600;

function isMissing(error) {
	return error && error.code === "ENOENT";
}

async function assertUnusedSocketPath(socketPath) {
	try {
		const stats = await lstat(socketPath);
		if (stats.isSocket())
			throw new Error(`socket path is already in use: ${socketPath}`);
		throw new Error(`refusing to replace non-socket path: ${socketPath}`);
	} catch (error) {
		if (!isMissing(error)) throw error;
	}
}

function requestIdFromLine(line) {
	try {
		const value = JSON.parse(line);
		return typeof value?.id === "string" ? value.id : undefined;
	} catch {
		return undefined;
	}
}

function write(socket, frame) {
	if (!socket.destroyed) socket.write(encodeFrame(frame));
}

export async function startUnixSocketServer({
	socketPath,
	handleRequest,
	subscribe,
	maxFrameBytes = MAX_FRAME_BYTES,
}) {
	if (typeof socketPath !== "string" || !socketPath.startsWith("/")) {
		throw new Error("socketPath must be an absolute Unix-socket path");
	}
	if (typeof handleRequest !== "function")
		throw new TypeError("handleRequest must be a function");
	if (subscribe !== undefined && typeof subscribe !== "function")
		throw new TypeError("subscribe must be a function");
	await assertUnusedSocketPath(socketPath);

	const sockets = new Set();
	const server = createServer((socket) => {
		sockets.add(socket);
		socket.on("error", () => {});
		const decoder = new StringDecoder("utf8");
		const cleanup = new Set();
		let pending = "";
		let closedForFrameLimit = false;

		socket.on("close", () => {
			sockets.delete(socket);
			for (const unsubscribe of cleanup) unsubscribe();
			cleanup.clear();
		});

		const handleSubscription = async (request) => {
			const queued = [];
			let ready = false;
			const subscription = await subscribe(request, (event) => {
				if (ready) write(socket, { version: "v1", type: "event", event });
				else queued.push(event);
			});
			cleanup.add(subscription.unsubscribe);
			write(socket, successResponse(request.id, subscription.result));
			ready = true;
			for (const event of queued)
				write(socket, { version: "v1", type: "event", event });
		};

		const handleLine = async (line) => {
			const normalized = line.endsWith("\r") ? line.slice(0, -1) : line;
			try {
				const request = parseRequestFrame(normalized, maxFrameBytes);
				if (
					(request.op === "subscribe" || request.op === "subscribe_workspace") &&
					subscribe
				) {
					await handleSubscription(request);
					return;
				}
				const result = await handleRequest(request);
				write(socket, successResponse(request.id, result));
			} catch (error) {
				write(socket, errorResponse(requestIdFromLine(normalized), error));
			}
		};

		socket.on("data", (chunk) => {
			if (closedForFrameLimit) return;
			pending += decoder.write(chunk);
			let newline;
			while ((newline = pending.indexOf("\n")) !== -1) {
				const line = pending.slice(0, newline);
				pending = pending.slice(newline + 1);
				void handleLine(line);
			}
			if (Buffer.byteLength(pending, "utf8") > maxFrameBytes) {
				closedForFrameLimit = true;
				write(
					socket,
					errorResponse(
						undefined,
						new ProtocolError(
							"frame_too_large",
							`frame exceeds ${maxFrameBytes} bytes`,
						),
					),
				);
				socket.end();
			}
		});

		socket.on("end", () => {
			if (closedForFrameLimit) return;
			const finalLine = pending + decoder.end();
			if (finalLine.length > 0) void handleLine(finalLine);
		});
	});

	await new Promise((resolve, reject) => {
		server.once("error", reject);
		server.listen(socketPath, () => {
			server.off("error", reject);
			resolve();
		});
	});

	try {
		await chmod(socketPath, OWNER_ONLY_MODE);
	} catch (error) {
		await new Promise((resolve) => server.close(resolve));
		await unlink(socketPath).catch(() => {});
		throw error;
	}

	return {
		socketPath,
		async close() {
			for (const socket of sockets) socket.destroy();
			await new Promise((resolve, reject) =>
				server.close((error) => (error ? reject(error) : resolve())),
			);
			await unlink(socketPath).catch((error) => {
				if (!isMissing(error)) throw error;
			});
		},
	};
}