mirror of
https://github.com/tale/headplane.git
synced 2026-08-21 02:06:37 +00:00
feat: expand frame type to support stdout/stderr chan
This commit is contained in:
Aarnav Tale
@@ -1,5 +1,7 @@
|
||||
import type { Writable } from 'node:stream';
|
||||
import { decode, encode } from 'cbor2';
|
||||
import { encode } from 'cborg';
|
||||
import { WSContext } from 'hono/ws';
|
||||
import { ChannelType } from './encoder';
|
||||
|
||||
interface Command {
|
||||
op: string;
|
||||
@@ -25,7 +27,6 @@ export async function dispatchCommand(
|
||||
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 {
|
||||
@@ -34,3 +35,31 @@ export async function dispatchCommand(
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
interface SSHConnectData extends Command {
|
||||
op: 'ssh_conn_successful';
|
||||
payload: {
|
||||
sessionId: string;
|
||||
};
|
||||
}
|
||||
|
||||
interface SSHConnectFailedData extends Command {
|
||||
op: 'ssh_conn_failed';
|
||||
payload: {
|
||||
reason: string;
|
||||
};
|
||||
}
|
||||
|
||||
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));
|
||||
}
|
||||
|
||||
@@ -0,0 +1,99 @@
|
||||
// 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: Blob | ArrayBufferLike | string;
|
||||
}
|
||||
|
||||
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;
|
||||
}
|
||||
|
||||
const payload = Buffer.isBuffer(frame.payload)
|
||||
? frame.payload
|
||||
: typeof frame.payload === 'string'
|
||||
? Buffer.from(frame.payload, 'utf8')
|
||||
: frame.payload instanceof Blob
|
||||
? Buffer.from(await frame.payload.arrayBuffer())
|
||||
: Buffer.from(frame.payload);
|
||||
|
||||
// Size can only hold 4 bytes
|
||||
if (payload.length > 0xffffffff) {
|
||||
log.error('agent', 'SSH payload too large: %d bytes', payload.length);
|
||||
return;
|
||||
}
|
||||
|
||||
const frameSize =
|
||||
4 + // Magic
|
||||
1 + // Version
|
||||
1 + // Channel Type
|
||||
(1 + sid.length) + // Session ID length + SID
|
||||
(4 + 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(payload.length, offset);
|
||||
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,
|
||||
};
|
||||
}
|
||||
+75
-61
@@ -4,7 +4,8 @@ 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';
|
||||
import { dispatchCommand, dispatchWeb } from './dispatcher';
|
||||
import { decodeSSHFrame, encodeSSHFrame } from './encoder';
|
||||
|
||||
interface SSHConnection {
|
||||
username: string;
|
||||
@@ -19,19 +20,44 @@ interface SSHSession {
|
||||
ws: WSContext;
|
||||
}
|
||||
|
||||
interface FrameDecodeSuccess {
|
||||
id: string;
|
||||
data: Buffer;
|
||||
}
|
||||
|
||||
interface FrameDecodeFailure {
|
||||
id: undefined;
|
||||
data: undefined;
|
||||
}
|
||||
|
||||
export function createSSHMultiplexer(proc: ChildProcess): SSHMultiplexer {
|
||||
return new SSHMultiplexer(proc);
|
||||
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 child: ChildProcess;
|
||||
private control: Writable;
|
||||
private sshInput: Writable;
|
||||
private sshOutput: Readable;
|
||||
|
||||
constructor(proc: ChildProcess) {
|
||||
constructor(control: Writable, sshInput: Writable, sshOutput: Readable) {
|
||||
this.connections = new Map();
|
||||
this.child = proc;
|
||||
|
||||
this.handleStdout();
|
||||
this.control = control;
|
||||
this.sshInput = sshInput;
|
||||
this.sshOutput = sshOutput;
|
||||
this.configureStdout();
|
||||
}
|
||||
|
||||
// TODO: Determine if we want to allow multiple connections for the same
|
||||
@@ -45,8 +71,8 @@ export class SSHMultiplexer {
|
||||
ws,
|
||||
};
|
||||
|
||||
log.info('agent', 'Dispatching SSH connection for %s', sessionId);
|
||||
await dispatchCommand(this.child.stdin, {
|
||||
log.debug('agent', 'Dispatching SSH connection for %s', sessionId);
|
||||
await dispatchCommand(this.control, {
|
||||
op: 'ssh_conn',
|
||||
payload: {
|
||||
sessionId,
|
||||
@@ -76,9 +102,22 @@ export class SSHMultiplexer {
|
||||
try {
|
||||
const sessionId = await this.connect(conn, ws);
|
||||
ws.raw = sessionId;
|
||||
ws.send(JSON.stringify({ status: 'connected', sessionId }));
|
||||
dispatchWeb(ws, {
|
||||
op: 'ssh_conn_successful',
|
||||
payload: { sessionId },
|
||||
});
|
||||
} catch (error) {
|
||||
ws.close(1011, `Connection failed: ${error.message}`);
|
||||
const errorMessage =
|
||||
error instanceof Error ? error.message : 'Unknown error';
|
||||
|
||||
dispatchWeb(ws, {
|
||||
op: 'ssh_conn_failed',
|
||||
payload: {
|
||||
reason: errorMessage,
|
||||
},
|
||||
});
|
||||
|
||||
ws.close(1011, 'Connection failed');
|
||||
}
|
||||
},
|
||||
|
||||
@@ -95,8 +134,13 @@ export class SSHMultiplexer {
|
||||
return;
|
||||
}
|
||||
|
||||
const encodedFrame = this.encodeFrame(sessionId, event.data);
|
||||
this.child.stdio[3]?.write(encodedFrame);
|
||||
const encodedFrame = await encodeSSHFrame({
|
||||
sessionId,
|
||||
channel: 0, // stdin
|
||||
payload: event.data,
|
||||
});
|
||||
|
||||
this.sshInput.write(encodedFrame);
|
||||
},
|
||||
|
||||
onClose: (_, ws) => {
|
||||
@@ -121,65 +165,35 @@ export class SSHMultiplexer {
|
||||
}
|
||||
|
||||
log.error('agent', 'SSH WebSocket Error with %s', sessionId);
|
||||
console.log(event);
|
||||
log.debug('agent', 'Error details: %o', 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) {
|
||||
private configureStdout() {
|
||||
this.sshOutput.on('data', (bytes) => {
|
||||
const frame = decodeSSHFrame(bytes);
|
||||
if (!frame) {
|
||||
return;
|
||||
}
|
||||
|
||||
const { id, data } = decoded;
|
||||
console.log(id, data);
|
||||
const session = this.connections.get(id);
|
||||
const session = this.connections.get(frame.sessionId);
|
||||
if (!session || !session.connected) {
|
||||
log.warn('agent', 'Received data for disconnected session %s', id);
|
||||
log.warn(
|
||||
'agent',
|
||||
'Received data for invalid SSH session %s',
|
||||
frame.sessionId,
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
session.ws.send(data);
|
||||
dispatchWeb(session.ws, {
|
||||
op: 'ssh_frame',
|
||||
payload: {
|
||||
channel: frame.channel,
|
||||
data: frame.payload,
|
||||
},
|
||||
});
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
+20
-2
@@ -1,6 +1,6 @@
|
||||
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';
|
||||
@@ -97,7 +97,25 @@ export default createHonoServer({
|
||||
app.get(
|
||||
'/_ssh_plexer',
|
||||
upgradeWebSocket((c) => {
|
||||
return agentManager.multiplexer!.websocketHandler(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'),
|
||||
};
|
||||
}),
|
||||
);
|
||||
},
|
||||
|
||||
Reference in New Issue
Block a user