gsadmin/headplane
1
0
Fork 0
mirror of https://github.com/tale/headplane.git synced 2026-08-20 18:02:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: switch away from websocket to stdout messaging for agent

Browse Source
This commit is contained in:
Aarnav Tale
2025-04-11 12:00:19 -04:00
parent 608fcacdb1
commit c5e6a6c407
19 changed files with 522 additions and 420 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/server/config/schema.ts
+28 -12
View File
@@ -1,22 +1,26 @@
import { type } from 'arktype';
const stringToBool = type('string | boolean').pipe((v) => Boolean(v));
const stringToBool = type('string | boolean').pipe((v) => {
if (typeof v === 'string') {
if (v === '1' || v === 'true' || v === 'yes') {
return true;
}
if (v === '0' || v === 'false' || v === 'no') {
return false;
}
throw new Error(`Invalid string value for boolean: ${v}`);
}
return Boolean(v);
});
const serverConfig = type({
host: 'string.ip',
port: type('string | number.integer').pipe((v) => Number(v)),
cookie_secret: '32 <= string <= 32',
cookie_secure: stringToBool,
agent: type({
authkey: 'string = ""',
ttl: 'number.integer = 180000', // Default to 3 minutes
cache_path: 'string = "/var/lib/headplane/agent_cache.json"',
})
.onDeepUndeclaredKey('reject')
.default(() => ({
authkey: '',
ttl: 180000,
cache_path: '/var/lib/headplane/agent_cache.json',
})),
});
const oidcConfig = type({
@@ -46,6 +50,16 @@ const containerLabel = type({
value: 'string',
}).optional();
const agentConfig = type({
enabled: stringToBool.default(false),
host_name: 'string = "headplane-agent"',
pre_authkey: 'string = ""',
cache_ttl: 'number.integer = 180000',
cache_path: 'string = "/var/lib/headplane/agent_cache.json"',
executable_path: 'string = "/usr/libexec/headplane/agent"',
work_dir: 'string = "/var/lib/headplane/agent"',
});
const dockerConfig = type({
enabled: stringToBool,
container_name: 'string',
@@ -67,12 +81,14 @@ const integrationConfig = type({
'docker?': dockerConfig,
'kubernetes?': kubernetesConfig,
'proc?': procConfig,
'agent?': agentConfig,
}).onDeepUndeclaredKey('reject');
const partialIntegrationConfig = type({
'docker?': dockerConfig.partial(),
'kubernetes?': kubernetesConfig.partial(),
'proc?': procConfig.partial(),
'agent?': agentConfig.partial(),
}).partial();
export const headplaneConfig = type({
app/server/index.ts
+3 -21
View File
@@ -1,7 +1,6 @@
import { env, versions } from 'node:process';
import type { UpgradeWebSocket } from 'hono/ws';
import { createHonoServer } from 'react-router-hono-server/node';
import type { WebSocket } from 'ws';
import log from '~/utils/log';
import { configureConfig, configureLogger, envVariables } from './config/env';
import { loadIntegration } from './config/integration';
@@ -57,11 +56,9 @@ const appLoadContext = {
),
agents: await loadAgentSocket(
config.server.agent.authkey,
config.server.agent.cache_path,
config.server.agent.ttl,
config.integration?.agent,
config.headscale.url,
),
integration: await loadIntegration(config.integration),
oidc: config.oidc ? await createOidcClient(config.oidc) : undefined,
};
@@ -71,7 +68,6 @@ declare module 'react-router' {
}
export default createHonoServer({
useWebSocket: true,
overrideGlobalObjects: true,
port: config.server.port,
hostname: config.server.host,
@@ -85,20 +81,6 @@ export default createHonoServer({
return appLoadContext;
},
configure(app, { upgradeWebSocket }) {
const agentManager = appLoadContext.agents;
if (agentManager) {
app.get(
`${__PREFIX__}/_dial`,
// We need this since we cannot pass the WSEvents context
// Also important to not pass the callback directly
// since we need to retain `this` context
(upgradeWebSocket as UpgradeWebSocket<WebSocket>)((c) =>
agentManager.configureSocket(c),
),
);
}
},
listeningListener(info) {
log.info('server', 'Running on %s:%s', info.address, info.port);
},
app/server/web/agent.ts
+236 -117
View File
@@ -1,159 +1,278 @@
import { ChildProcess, spawn } from 'node:child_process';
import { createHash } from 'node:crypto';
import { open, readFile, writeFile } from 'node:fs/promises';
import { constants, access, open, readFile, writeFile } from 'node:fs/promises';
import { exit } from 'node:process';
import { createInterface } from 'node:readline';
import { setTimeout } from 'node:timers/promises';
import { getConnInfo } from '@hono/node-server/conninfo';
import { type } from 'arktype';
import type { Context } from 'hono';
import type { WSContext, WSEvents } from 'hono/ws';
import { WebSocket } from 'ws';
import { HostInfo } from '~/types';
import log from '~/utils/log';
import type { HeadplaneConfig } from '../config/schema';
interface LogResponse {
Level: 'info' | 'debug' | 'error' | 'fatal';
Message: string;
}
interface RegisterMessage {
Type: 'register';
ID: string;
}
interface StatusMessage {
Type: 'status';
Data: Record<string, HostInfo>;
}
interface MessageResponse {
Level: 'msg';
Message: RegisterMessage | StatusMessage;
}
type AgentResponse = LogResponse | MessageResponse;
export async function loadAgentSocket(
authkey: string,
path: string,
ttl: number,
config: NonNullable<HeadplaneConfig['integration']>['agent'] | undefined,
headscaleUrl: string,
) {
if (authkey.length === 0) {
if (!config?.enabled) {
return;
}
if (config.pre_authkey.trim().length === 0) {
log.error('agent', 'Agent `pre_authkey` is not set');
log.warn('agent', 'The agent will not run until resolved');
return;
}
try {
const handle = await open(path, 'w');
log.info('agent', 'Using agent cache file at %s', path);
await access(config.work_dir, constants.R_OK | constants.W_OK);
log.debug('config', 'Using agent work dir at %s', config.work_dir);
} catch (error) {
log.info('config', 'Agent work dir not accessible at %s', config.work_dir);
log.debug('config', 'Error details: %s', error);
return;
}
try {
const handle = await open(config.cache_path, 'a+');
log.info('agent', 'Using agent cache file at %s', config.cache_path);
await handle.close();
} catch (error) {
log.info('agent', 'Agent cache file not accessible at %s', path);
log.info(
'agent',
'Agent cache file not accessible at %s',
config.cache_path,
);
log.debug('agent', 'Error details: %s', error);
return;
}
const cache = new TimedCache<HostInfo>(ttl, path);
return new AgentManager(cache, authkey);
const cache = new TimedCache<HostInfo>(config.cache_ttl, config.cache_path);
return new AgentManager(cache, config, headscaleUrl);
}
class AgentManager {
private static readonly MAX_RESTARTS = 5;
private restartCounter = 0;
private cache: TimedCache<HostInfo>;
private agents: Map<string, WSContext>;
private timers: Map<string, NodeJS.Timeout>;
private authkey: string;
private headscaleUrl: string;
private config: NonNullable<
NonNullable<HeadplaneConfig['integration']>['agent']
>;
constructor(cache: TimedCache<HostInfo>, authkey: string) {
private spawnProcess: ChildProcess | null;
private agentId: string | null;
constructor(
cache: TimedCache<HostInfo>,
config: NonNullable<NonNullable<HeadplaneConfig['integration']>['agent']>,
headscaleUrl: string,
) {
this.cache = cache;
this.authkey = authkey;
this.agents = new Map();
this.timers = new Map();
this.config = config;
this.headscaleUrl = headscaleUrl;
this.spawnProcess = null;
this.agentId = null;
this.startAgent();
process.on('SIGINT', () => {
this.spawnProcess?.kill();
exit(0);
});
process.on('SIGTERM', () => {
this.spawnProcess?.kill();
exit(0);
});
}
tailnetIDs() {
return Array.from(this.agents.keys());
/**
* Used by the UI to indicate why the agent is not running.
* Exhaustion requires a manual restart of the agent.
* (Which can be invoked via the UI)
* @returns true if the agent is exhausted
*/
exhausted() {
return this.restartCounter >= AgentManager.MAX_RESTARTS;
}
lookup(nodeIds: string[]) {
/*
* Called by the UI to manually force a restart of the agent.
*/
deExhaust() {
this.restartCounter = 0;
this.startAgent();
}
/*
* Stored agent ID for the current process. This is caught by the agent
* when parsing the stdout on agent startup.
*/
agentID() {
return this.agentId;
}
private startAgent() {
if (this.spawnProcess) {
log.debug('agent', 'Agent already running');
return;
}
if (this.exhausted()) {
log.error('agent', 'Agent is exhausted, cannot start');
return;
}
// Cannot be detached since we want to follow our process lifecycle
// We also need to be able to send data to the process by using stdin
log.info(
'agent',
'Starting agent process (attempt %d)',
this.restartCounter,
);
this.spawnProcess = spawn(this.config.executable_path, [], {
detached: false,
stdio: ['pipe', 'pipe', 'pipe'],
env: {
HOME: process.env.HOME,
HEADPLANE_EMBEDDED: 'true',
HEADPLANE_AGENT_WORK_DIR: this.config.work_dir,
HEADPLANE_AGENT_DEBUG: log.debugEnabled ? 'true' : 'false',
HEADPLANE_AGENT_HOSTNAME: this.config.host_name,
HEADPLANE_AGENT_TS_SERVER: this.headscaleUrl,
HEADPLANE_AGENT_TS_AUTHKEY: this.config.pre_authkey,
},
});
if (!this.spawnProcess?.pid) {
log.error('agent', 'Failed to start agent process');
this.restartCounter++;
global.setTimeout(() => this.startAgent(), 1000);
return;
}
if (this.spawnProcess.stdin === null || this.spawnProcess.stdout === null) {
log.error('agent', 'Failed to connect to agent process');
this.restartCounter++;
global.setTimeout(() => this.startAgent(), 1000);
return;
}
const rlStdout = createInterface({
input: this.spawnProcess.stdout,
crlfDelay: Number.POSITIVE_INFINITY,
});
rlStdout.on('line', (line) => {
try {
const parsed = JSON.parse(line) as AgentResponse;
if (parsed.Level === 'msg') {
switch (parsed.Message.Type) {
case 'register':
this.agentId = parsed.Message.ID;
break;
case 'status':
for (const [key, value] of Object.entries(parsed.Message.Data)) {
// Mark the agent as the one that is running
// We store it in the cache so that it shows
// itself later
if (key === this.agentId) {
value.HeadplaneAgent = true;
}
this.cache.set(key, value);
}
break;
}
return;
}
switch (parsed.Level) {
case 'info':
case 'debug':
case 'error':
log[parsed.Level]('agent', parsed.Message);
break;
case 'fatal':
log.error('agent', parsed.Message);
break;
default:
log.debug('agent', 'Unknown agent response: %s', line);
break;
}
} catch (error) {
log.debug('agent', 'Failed to parse agent response: %s', error);
log.debug('agent', 'Raw data: %s', line);
}
});
this.spawnProcess.on('error', (error) => {
log.error('agent', 'Failed to start agent process: %s', error);
this.restartCounter++;
this.spawnProcess = null;
global.setTimeout(() => this.startAgent(), 1000);
});
this.spawnProcess.on('exit', (code) => {
log.error('agent', 'Agent process exited with code %d', code ?? -1);
this.restartCounter++;
this.spawnProcess = null;
global.setTimeout(() => this.startAgent(), 1000);
});
}
async lookup(nodeIds: string[]) {
const entries = this.cache.toJSON();
const missing = nodeIds.filter((nodeId) => !entries[nodeId]);
if (missing.length > 0) {
this.requestData(missing);
await this.requestData(missing);
}
return entries;
}
// Request data from all connected agents
// This does not return anything, but caches the data which then needs to be
// queried by the caller separately.
private requestData(nodeList: string[]) {
const NodeIDs = [...new Set(nodeList)];
NodeIDs.map((node) => {
log.debug('agent', 'Requesting agent data for %s', node);
});
for (const agent of this.agents.values()) {
agent.send(JSON.stringify({ NodeIDs }));
// Request data from the internal agent by sending a message to the process
// via stdin. This is a blocking call, so it will wait for the agent to
// respond before returning.
private async requestData(nodeList: string[]) {
if (this.exhausted()) {
return;
}
}
// Since we are using Node, Hono is built on 'ws' WebSocket types.
configureSocket(c: Context): WSEvents<WebSocket> {
return {
onOpen: (_, ws) => {
const id = c.req.header('x-headplane-tailnet-id');
if (!id) {
log.warn(
'agent',
'Rejecting an agent WebSocket connection without a tailnet ID',
);
ws.close(1008, 'ERR_INVALID_TAILNET_ID');
return;
}
// Wait for the process to be spawned, busy waiting is gross
while (this.spawnProcess === null) {
await setTimeout(100);
}
const auth = c.req.header('authorization');
if (auth !== `Bearer ${this.authkey}`) {
log.warn('agent', 'Rejecting an unauthorized WebSocket connection');
const info = getConnInfo(c);
if (info.remote.address) {
log.warn('agent', 'Agent source IP: %s', info.remote.address);
}
ws.close(1008, 'ERR_UNAUTHORIZED');
return;
}
const pinger = setInterval(() => {
if (ws.readyState !== 1) {
clearInterval(pinger);
return;
}
ws.raw?.ping();
}, 30000);
this.agents.set(id, ws);
this.timers.set(id, pinger);
},
onClose: () => {
const id = c.req.header('x-headplane-tailnet-id');
if (!id) {
return;
}
clearInterval(this.timers.get(id));
this.agents.delete(id);
},
onError: (event, ws) => {
const id = c.req.header('x-headplane-tailnet-id');
if (!id) {
return;
}
clearInterval(this.timers.get(id));
if (event instanceof ErrorEvent) {
log.error('agent', 'WebSocket error: %s', event.message);
}
log.debug('agent', 'Closing agent WebSocket connection');
ws.close(1011, 'ERR_INTERNAL_ERROR');
},
// This is where we receive the data from the agent
// Requests are made in the AgentManager.requestData function
onMessage: (event, ws) => {
const id = c.req.header('x-headplane-tailnet-id');
if (!id) {
return;
}
const data = JSON.parse(event.data.toString());
log.debug('agent', 'Received agent data from %s', id);
for (const [node, info] of Object.entries<HostInfo>(data)) {
this.cache.set(node, info);
log.debug('agent', 'Cached HostInfo for %s', node);
}
},
};
// Send the request to the agent, without waiting for a response.
// The live data invalidator will re-request the data if it is not
// available in the cache anyways.
const data = JSON.stringify({ NodeIDs: nodeList });
this.spawnProcess.stdin?.write(`${data}\n`);
}
}