gsadmin/headplane
1
0
Fork 0
mirror of https://github.com/tale/headplane.git synced 2026-08-10 14:06:51 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: cleanup removal of old ssh plexer and logic

Browse Source 
we also have added the necessary logic to auto prune ephemeral nodes because
headscale doesn't seem to automatically remove them. this change made use of a database
which is now stored in the persistent headplane directory.
This commit is contained in:
Aarnav Tale
2025-06-20 00:14:00 -04:00
parent bf1d75a27a
commit b18147fa82
24 changed files with 963 additions and 547 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/server/agent/dispatcher.ts
-89
View File
@@ -1,89 +0,0 @@
import type { Writable } from 'node:stream';
import { encode } from 'cborg';
import { WSContext } from 'hono/ws';
import { ChannelType } from './encoder';
export interface Command {
op: string;
payload: unknown;
}
export interface SSHConnectCommand extends Command {
op: 'ssh_conn';
payload: {
sessionId: string;
username: string;
hostname: string;
port: number;
};
}
export interface SSHCloseCommand extends Command {
op: 'ssh_term';
payload: {
sessionId: string;
};
}
export interface SSHResizeCommand extends Command {
op: 'ssh_resize';
payload: {
sessionId: string;
width: number;
height: number;
};
}
export interface SSHDataCommand extends Command {
op: 'ssh_data';
payload: {
sessionId: string;
data: Uint8Array;
};
}
type AgentCommand = SSHConnectCommand | SSHCloseCommand | SSHResizeCommand;
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) => {
if (err) {
reject(err);
} else {
resolve();
}
});
});
}
export interface SSHConnectData extends Command {
op: 'ssh_conn_successful';
payload: {
sessionId: string;
};
}
export interface SSHConnectFailedData extends Command {
op: 'ssh_conn_failed';
payload: {
reason: string;
};
}
export interface SSHFrameData extends Command {
op: 'ssh_frame';
payload: {
channel: ChannelType;
frame: Buffer;
};
}
type WebData = SSHConnectData | SSHConnectFailedData | SSHFrameData;
export function dispatchWeb<T>(dispatcher: WSContext<T>, data: WebData) {
return dispatcher.send(encode(data));
}
app/server/agent/encoder.ts
-91
View File
@@ -1,91 +0,0 @@
// Refer to agent/internal/sshutil/encoder.go for more details
// This is the Node.js implementation of the SSH encoder
import log from '~/utils/log';
const MAGIC = 'HPLS';
const VERSION = 1;
// 0 -> Stdin
// 1 -> Stdout
// 2 -> Stderr
export type ChannelType = 0 | 1 | 2;
interface SSHFrame {
sessionId: string;
channel: ChannelType;
payload: Buffer;
}
export async function encodeSSHFrame(frame: SSHFrame) {
const sid = Buffer.from(frame.sessionId, 'utf8');
if (sid.length > 255) {
log.error('agent', 'SSH session ID too long: %s', frame.sessionId);
return;
}
// Size can only hold 4 bytes
if (frame.payload.length > 0xffffffff) {
log.error('agent', 'SSH payload too large: %d bytes', frame.payload.length);
return;
}
const frameSize =
4 + // Magic
1 + // Version
1 + // Channel Type
(1 + sid.length) + // Session ID length + SID
(4 + frame.payload.length); // Payload length + Payload
const buf = Buffer.alloc(frameSize);
buf.write(MAGIC, 0, 'utf8');
buf.writeUInt8(VERSION, 4);
buf.writeUInt8(frame.channel, 5);
buf.writeUInt8(sid.length, 6);
const offset = 7 + sid.length;
sid.copy(buf, 7);
buf.writeUInt32BE(frame.payload.length, offset);
frame.payload.copy(buf, offset + 4);
return buf;
}
export function decodeSSHFrame(data: Buffer) {
if (data.length < 5) {
log.error('agent', 'SSH frame too short: %d bytes', data.length);
return;
}
const magic = data.toString('utf8', 0, 4);
const version = data.readUInt8(4);
if (magic !== MAGIC || version !== VERSION) {
log.error('agent', 'Invalid SSH frame magic or version');
return;
}
const channel = data.readUInt8(5) as ChannelType;
const sidLength = data.readUInt8(6);
if (data.length < 7 + sidLength + 4) {
log.error('agent', 'SSH frame too short for session ID and payload');
return;
}
const sessionId = data.toString('utf8', 7, 7 + sidLength);
const payloadLength = data.readUInt32BE(7 + sidLength);
if (data.length < 7 + sidLength + 4 + payloadLength) {
log.error('agent', 'SSH frame too short for payload');
return;
}
const payload = data.subarray(
7 + sidLength + 4,
7 + sidLength + 4 + payloadLength,
);
return {
sessionId,
channel,
payload,
};
}
app/server/agent/ssh.ts
-255
View File
@@ -1,255 +0,0 @@
import { ChildProcess } from 'node:child_process';
import { randomUUID } from 'node:crypto';
import type { Readable, Writable } from 'node:stream';
import { decode } from 'cborg';
import { Context } from 'hono';
import { WSContext, WSEvents } from 'hono/ws';
import log from '~/utils/log';
import {
Command,
SSHDataCommand,
SSHResizeCommand,
dispatchCommand,
dispatchWeb,
} from './dispatcher';
import { decodeSSHFrame, encodeSSHFrame } from './encoder';
interface SSHConnection {
username: string;
hostname: string;
port: number;
}
interface SSHSession {
connectionDetails: SSHConnection;
connected: boolean;
sessionId: string;
ws: WSContext;
}
interface FrameDecodeSuccess {
id: string;
data: Buffer;
}
interface FrameDecodeFailure {
id: undefined;
data: undefined;
}
export function createSSHMultiplexer(proc: ChildProcess): SSHMultiplexer {
const control = proc.stdin;
const sshInput = proc.stdio[3];
const sshOutput = proc.stdio[4];
if (!control || !sshInput || !sshOutput) {
throw new Error('Invalid SSH multiplexer process: missing stdio streams');
}
return new SSHMultiplexer(
control,
sshInput as Writable,
sshOutput as Readable,
);
}
export class SSHMultiplexer {
private connections: Map<string, SSHSession>;
private control: Writable;
private sshInput: Writable;
private sshOutput: Readable;
constructor(control: Writable, sshInput: Writable, sshOutput: Readable) {
this.connections = new Map();
this.control = control;
this.sshInput = sshInput;
this.sshOutput = sshOutput;
this.configureStdout();
}
// 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.debug('agent', 'Dispatching SSH connection for %s', sessionId);
await dispatchCommand(this.control, {
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;
dispatchWeb(ws, {
op: 'ssh_conn_successful',
payload: { sessionId },
});
} catch (error) {
const errorMessage =
error instanceof Error ? error.message : 'Unknown error';
dispatchWeb(ws, {
op: 'ssh_conn_failed',
payload: {
reason: errorMessage,
},
});
ws.close(1011, 'Connection failed');
}
},
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 wsData = Buffer.isBuffer(event.data)
? event.data
: typeof event.data === 'string'
? Buffer.from(event.data, 'utf8')
: event.data instanceof Blob
? Buffer.from(await event.data.arrayBuffer())
: Buffer.from(event.data);
const obj = decode(wsData) as Command;
if (obj.op === 'ssh_data') {
const data = obj as SSHDataCommand;
if (data.payload.sessionId !== sessionId) {
log.warn(
'agent',
'Received data for mismatched SSH session %s',
data.payload.sessionId,
);
return;
}
const encodedFrame = await encodeSSHFrame({
sessionId,
channel: 0, // stdin
payload: Buffer.from(data.payload.data),
});
this.sshInput.write(encodedFrame);
}
if (obj.op === 'ssh_resize') {
const resize = obj as SSHResizeCommand;
if (resize.payload.sessionId !== sessionId) {
log.warn(
'agent',
'Received resize for mismatched SSH session %s',
resize.payload.sessionId,
);
return;
}
await dispatchCommand(this.control, resize);
}
},
onClose: async (_, ws) => {
const sessionId = ws.raw;
if (sessionId && this.connections.has(sessionId)) {
const session = this.connections.get(sessionId);
if (session) {
await dispatchCommand(this.control, {
op: 'ssh_term',
payload: {
sessionId,
},
});
session.connected = false;
this.connections.delete(sessionId);
}
}
},
onError: async (event, ws) => {
const sessionId = ws.raw;
if (sessionId && this.connections.has(sessionId)) {
const session = this.connections.get(sessionId);
if (session) {
await dispatchCommand(this.control, {
op: 'ssh_term',
payload: {
sessionId,
},
});
session.connected = false;
this.connections.delete(sessionId);
}
}
log.error('agent', 'SSH WebSocket Error with %s', sessionId);
log.debug('agent', 'Error details: %o', event);
},
};
}
private configureStdout() {
this.sshOutput.on('data', (bytes) => {
const frame = decodeSSHFrame(bytes);
if (!frame) {
return;
}
const session = this.connections.get(frame.sessionId);
if (!session || !session.connected) {
log.warn(
'agent',
'Received data for invalid SSH session %s',
frame.sessionId,
);
return;
}
dispatchWeb(session.ws, {
op: 'ssh_frame',
payload: {
channel: frame.channel,
frame: frame.payload,
},
});
});
}
}
app/server/config/schema.ts
+1
View File
@@ -19,6 +19,7 @@ const stringToBool = type('string | boolean').pipe((v) => {
const serverConfig = type({
host: 'string.ip',
port: type('string | number.integer').pipe((v) => Number(v)),
data_path: 'string = "/var/lib/headplane/"',
cookie_secret: '32 <= string <= 32',
cookie_secure: stringToBool,
});
app/server/db/client.ts
+26
View File
@@ -0,0 +1,26 @@
import { mkdir } from 'node:fs/promises';
import { dirname } from 'node:path';
import { drizzle } from 'drizzle-orm/better-sqlite3';
import { migrate } from 'drizzle-orm/better-sqlite3/migrator';
import log from '~/utils/log';
export async function createDbClient(path: string) {
try {
await mkdir(dirname(path), { recursive: true });
} catch (error) {
log.error(
'server',
'Failed to create directory for database at %s: %s',
path,
error instanceof Error ? error.message : String(error),
);
throw new Error(`Could not create directory for database at ${path}`);
}
const db = drizzle(path);
migrate(db, {
migrationsFolder: './drizzle',
});
return db;
}
app/server/db/pruner.ts
+58
View File
@@ -0,0 +1,58 @@
import { eq, isNotNull } from 'drizzle-orm';
import { LoaderFunctionArgs } from 'react-router';
import { Machine } from '~/types';
import log from '~/utils/log';
import { LoadContext } from '..';
import { ephemeralNodes } from './schema';
export async function pruneEphemeralNodes({
context,
request,
}: LoaderFunctionArgs<LoadContext>) {
const session = await context.sessions.auth(request);
const ephemerals = await context.db
.select()
.from(ephemeralNodes)
.where(isNotNull(ephemeralNodes.node_key));
if (ephemerals.length === 0) {
log.debug('api', 'No ephemeral nodes to prune');
return;
}
const { nodes } = await context.client.get<{ nodes: Machine[] }>(
'v1/node',
session.get('api_key')!,
);
const toPrune = nodes.filter((node) => {
if (node.online) {
return false;
}
return ephemerals.some((ephemeral) => node.nodeKey === ephemeral.node_key);
});
if (toPrune.length === 0) {
log.debug('api', 'No SSH nodes to prune');
return;
}
// Delete from the Headscale nodes list and then from the database
const promises = toPrune.map((node) => {
return async () => {
log.info('api', `Pruning node ${node.name}`);
await context.client.delete(
`v1/node/${node.id}`,
session.get('api_key')!,
);
await context.db
.delete(ephemeralNodes)
.where(eq(ephemeralNodes.node_key, node.nodeKey));
log.info('api', `Node ${node.name} pruned successfully`);
};
});
await Promise.all(promises.map((p) => p()));
}
app/server/db/schema.ts
+9
View File
@@ -0,0 +1,9 @@
import { sqliteTable, text } from 'drizzle-orm/sqlite-core';
export const ephemeralNodes = sqliteTable('ephemeral_nodes', {
auth_key: text('auth_key').primaryKey(),
node_key: text('node_key'),
});
export type EphemeralNode = typeof ephemeralNodes.$inferSelect;
export type EphemeralNodeInsert = typeof ephemeralNodes.$inferInsert;
app/server/index.ts
+5 -33
View File
@@ -1,10 +1,11 @@
import { join } from 'node:path';
import { env, versions } from 'node:process';
import type { WSEvents } from 'hono/ws';
import { createHonoServer } from 'react-router-hono-server/node';
import log from '~/utils/log';
import { configureConfig, configureLogger, envVariables } from './config/env';
import { loadIntegration } from './config/integration';
import { loadConfig } from './config/loader';
import { createDbClient } from './db/client';
import { createApiClient } from './headscale/api-client';
import { loadHeadscaleConfig } from './headscale/config-loader';
import { loadAgentSocket } from './web/agent';
@@ -33,6 +34,8 @@ const agentManager = await loadAgentSocket(
config.headscale.url,
);
const db = await createDbClient(join(config.server.data_path, 'hp_persist.db'));
// 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.
@@ -64,6 +67,7 @@ const appLoadContext = {
agents: agentManager,
integration: await loadIntegration(config.integration),
oidc: config.oidc ? await createOidcClient(config.oidc) : undefined,
db,
};
declare module 'react-router' {
@@ -87,36 +91,4 @@ 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) => {
// MARK: This is a limitation of the hono NPM module we use
const wsHandler = agentManager.multiplexer?.websocketHandler(
c,
) as WSEvents<unknown>;
return {
onOpen: wsHandler
? wsHandler.onOpen
: (_, ws) => ws.close(1000, 'Multiplexer not available'),
onClose: wsHandler
? wsHandler.onClose
: (_, ws) => ws.close(1000, 'Multiplexer not available'),
onMessage: wsHandler
? wsHandler.onMessage
: (_, ws) => ws.close(1000, 'Multiplexer not available'),
onError: wsHandler
? wsHandler.onError
: (_, ws) => ws.close(1000, 'Multiplexer error'),
};
}),
);
},
});
app/server/web/agent.ts
-12
View File
@@ -10,12 +10,10 @@ 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 {
@@ -115,7 +113,6 @@ class AgentManager {
>;
private spawnProcess: ChildProcess | null;
multiplexer: SSHMultiplexer | null;
private agentId: string | null;
constructor(
@@ -127,7 +124,6 @@ class AgentManager {
this.config = config;
this.headscaleUrl = headscaleUrl;
this.spawnProcess = null;
this.multiplexer = null;
this.agentId = null;
this.startAgent();
@@ -214,14 +210,6 @@ 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,