feat: initial webssh tooling

app/server/README.md

@@ -7,6 +7,10 @@ many side-effects (in this case, importing a module may run code).
 ```
 server
 ├── index.ts: Loads everything and starts the web server.
+├── agent/
+│   ├── dispatcher.ts: Serializes commands for the agent control fd (stdin).
+│   ├── ssh.ts: Manages & multiplexes the active web SSH connections
+│   ├── env.ts: Checks the environment variables for custom overrides.
 ├── config/
 │   ├── integration/
 │   │   ├── abstract.ts: Defines the abstract class for integrations.

app/server/agent/dispatcher.ts

@@ -0,0 +1,36 @@
+import type { Writable } from 'node:stream';
+import { decode, encode } from 'cbor2';
+
+interface Command {
+	op: string;
+	payload: unknown;
+}
+
+interface SSHConnectCommand extends Command {
+	op: 'ssh_conn';
+	payload: {
+		sessionId: string;
+		username: string;
+		hostname: string;
+		port: number;
+	};
+}
+
+type AgentCommand = SSHConnectCommand;
+
+export async function dispatchCommand(
+	dispatcher: Writable,
+	command: AgentCommand,
+) {
+	return new Promise<void>((resolve, reject) => {
+		const encodedCommand = Buffer.concat([encode(command), Buffer.from('\n')]);
+		dispatcher.write(encodedCommand, (err) => {
+			console.log('Command dispatched:', command, err);
+			if (err) {
+				reject(err);
+			} else {
+				resolve();
+			}
+		});
+	});
+}

app/server/agent/ssh.ts

@@ -0,0 +1,185 @@
+import { ChildProcess } from 'node:child_process';
+import { randomUUID } from 'node:crypto';
+import type { Readable, Writable } from 'node:stream';
+import { Context } from 'hono';
+import { WSContext, WSEvents } from 'hono/ws';
+import log from '~/utils/log';
+import { dispatchCommand } from './dispatcher';
+
+interface SSHConnection {
+	username: string;
+	hostname: string;
+	port: number;
+}
+
+interface SSHSession {
+	connectionDetails: SSHConnection;
+	connected: boolean;
+	sessionId: string;
+	ws: WSContext;
+}
+
+export function createSSHMultiplexer(proc: ChildProcess): SSHMultiplexer {
+	return new SSHMultiplexer(proc);
+}
+
+export class SSHMultiplexer {
+	private connections: Map<string, SSHSession>;
+	private child: ChildProcess;
+
+	constructor(proc: ChildProcess) {
+		this.connections = new Map();
+		this.child = proc;
+
+		this.handleStdout();
+	}
+
+	// TODO: Determine if we want to allow multiple connections for the same
+	// target or attempt to reuse the existing connection (sounds stupid)
+	private async connect(conn: SSHConnection, ws: WSContext<string>) {
+		const sessionId = randomUUID();
+		const session: SSHSession = {
+			connectionDetails: conn,
+			connected: true,
+			sessionId,
+			ws,
+		};
+
+		log.info('agent', 'Dispatching SSH connection for %s', sessionId);
+		await dispatchCommand(this.child.stdin, {
+			op: 'ssh_conn',
+			payload: {
+				sessionId,
+				...conn,
+			},
+		});
+
+		this.connections.set(sessionId, session);
+		return sessionId;
+	}
+
+	websocketHandler(c: Context): WSEvents<string> {
+		return {
+			onOpen: async (_, ws) => {
+				const { username, hostname, port } = c.req.query();
+				if (!username || !hostname || !port) {
+					ws.close(1008, 'Missing connection parameters');
+					return;
+				}
+
+				const conn: SSHConnection = {
+					username,
+					hostname,
+					port: Number.parseInt(port, 10),
+				};
+
+				try {
+					const sessionId = await this.connect(conn, ws);
+					ws.raw = sessionId;
+					ws.send(JSON.stringify({ status: 'connected', sessionId }));
+				} catch (error) {
+					ws.close(1011, `Connection failed: ${error.message}`);
+				}
+			},
+
+			onMessage: async (event, ws) => {
+				const sessionId = ws.raw;
+				if (!sessionId || !this.connections.has(sessionId)) {
+					ws.close(1008, 'Invalid session ID');
+					return;
+				}
+
+				const session = this.connections.get(sessionId);
+				if (!session || !session.connected) {
+					ws.close(1008, 'Session not connected');
+					return;
+				}
+
+				const encodedFrame = this.encodeFrame(sessionId, event.data);
+				this.child.stdio[3]?.write(encodedFrame);
+			},
+
+			onClose: (_, ws) => {
+				const sessionId = ws.raw;
+				if (sessionId && this.connections.has(sessionId)) {
+					const session = this.connections.get(sessionId);
+					if (session) {
+						session.connected = false;
+						this.connections.delete(sessionId);
+					}
+				}
+			},
+
+			onError: (event, ws) => {
+				const sessionId = ws.raw;
+				if (sessionId && this.connections.has(sessionId)) {
+					const session = this.connections.get(sessionId);
+					if (session) {
+						session.connected = false;
+						this.connections.delete(sessionId);
+					}
+				}
+
+				log.error('agent', 'SSH WebSocket Error with %s', sessionId);
+				console.log(event);
+			},
+		};
+	}
+
+	private encodeFrame(id: string, data: string | Buffer): Buffer {
+		const sid = Buffer.from(id, 'utf8');
+		const payload = Buffer.isBuffer(data) ? data : Buffer.from(data, 'utf8');
+
+		// FIX: include +4 for the payload length
+		const frame = Buffer.alloc(1 + sid.length + 4 + payload.length);
+		frame.writeUint8(sid.length, 0); // 1 byte for sid length
+		sid.copy(frame, 1); // SID
+		frame.writeUint32BE(payload.length, 1 + sid.length); // 4 bytes for payload length
+		payload.copy(frame, 1 + sid.length + 4); // Payload
+
+		return frame;
+	}
+
+	private decodeFrame(frame: Buffer): { id: string; data: Buffer } | undefined {
+		if (frame.length < 5) return;
+
+		const sidLength = frame.readUint8(0);
+		if (frame.length < 1 + sidLength + 4) return;
+
+		const id = frame.slice(1, 1 + sidLength).toString('utf8');
+		const payloadLength = frame.readUint32BE(1 + sidLength);
+		if (frame.length < 1 + sidLength + 4 + payloadLength) return;
+
+		const data = frame.slice(
+			1 + sidLength + 4,
+			1 + sidLength + 4 + payloadLength,
+		);
+
+		return { id, data };
+	}
+
+	private handleStdout() {
+		const stdout = this.child.stdio[4];
+		if (!stdout) {
+			return;
+		}
+
+		stdout.on('data', (bytes) => {
+			console.log(Buffer.from(bytes).toString('utf8'));
+			const decoded = this.decodeFrame(bytes);
+			if (!decoded) {
+				return;
+			}
+
+			const { id, data } = decoded;
+			console.log(id, data);
+			const session = this.connections.get(id);
+			if (!session || !session.connected) {
+				log.warn('agent', 'Received data for disconnected session %s', id);
+				return;
+			}
+
+			session.ws.send(data);
+		});
+	}
+}

app/server/index.ts

@@ -28,6 +28,11 @@ const config = await loadConfig(
 	}),
 );
 
+const agentManager = await loadAgentSocket(
+	config.integration?.agent,
+	config.headscale.url,
+);
+
 // We also use this file to load anything needed by the react router code.
 // These are usually per-request things that we need access to, like the
 // helper that can issue and revoke cookies.
@@ -56,10 +61,7 @@ const appLoadContext = {
 		config.headscale.tls_cert_path,
 	),
 
-	agents: await loadAgentSocket(
-		config.integration?.agent,
-		config.headscale.url,
-	),
+	agents: agentManager,
 	integration: await loadIntegration(config.integration),
 	oidc: config.oidc ? await createOidcClient(config.oidc) : undefined,
 };
@@ -85,4 +87,18 @@ export default createHonoServer({
 	listeningListener(info) {
 		log.info('server', 'Running on %s:%s', info.address, info.port);
 	},
+
+	useWebSocket: true,
+	configure: (app, { upgradeWebSocket }) => {
+		if (agentManager === undefined) {
+			return;
+		}
+
+		app.get(
+			'/_ssh_plexer',
+			upgradeWebSocket((c) => {
+				return agentManager.multiplexer!.websocketHandler(c);
+			}),
+		);
+	},
 });

app/server/web/agent.ts

@@ -10,10 +10,12 @@ import {
 } from 'node:fs/promises';
 import { exit } from 'node:process';
 import { createInterface } from 'node:readline';
+import { Readable, Writable } from 'node:stream';
 import { setTimeout } from 'node:timers/promises';
 import { type } from 'arktype';
 import { HostInfo } from '~/types';
 import log from '~/utils/log';
+import { SSHMultiplexer, createSSHMultiplexer } from '../agent/ssh';
 import type { HeadplaneConfig } from '../config/schema';
 
 interface LogResponse {
@@ -113,6 +115,7 @@ class AgentManager {
 	>;
 
 	private spawnProcess: ChildProcess | null;
+	multiplexer: SSHMultiplexer | null;
 	private agentId: string | null;
 
 	constructor(
@@ -124,6 +127,7 @@ class AgentManager {
 		this.config = config;
 		this.headscaleUrl = headscaleUrl;
 		this.spawnProcess = null;
+		this.multiplexer = null;
 		this.agentId = null;
 		this.startAgent();
 
@@ -184,7 +188,7 @@ class AgentManager {
 		);
 		this.spawnProcess = spawn(this.config.executable_path, [], {
 			detached: false,
-			stdio: ['pipe', 'pipe', 'pipe'],
+			stdio: ['pipe', 'pipe', 'pipe', 'pipe', 'pipe'],
 			env: {
 				HOME: process.env.HOME,
 				HEADPLANE_EMBEDDED: 'true',
@@ -210,6 +214,14 @@ class AgentManager {
 			return;
 		}
 
+		const sshInput = this.spawnProcess.stdio[3];
+		const sshOutput = this.spawnProcess.stdio[4];
+
+		if (sshInput && sshOutput) {
+			log.info('agent', 'Using SSH multiplexer manager');
+			this.multiplexer = createSSHMultiplexer(this.spawnProcess);
+		}
+
 		const rlStdout = createInterface({
 			input: this.spawnProcess.stdout,
 			crlfDelay: Number.POSITIVE_INFINITY,