From 8cb91cd45b437b791e7543482e88893e2b231f06 Mon Sep 17 00:00:00 2001 From: Aarnav Tale Date: Tue, 19 Aug 2025 00:10:06 -0400 Subject: [PATCH] feat: overhaul hp_agent lifecycle handling * Added backoff and liveness probes for better management * Switched IPC to a simple text based system * Lookups don't directly touch the agent now * Use the database as a source of truth --- CHANGELOG.md | 3 +- app/server/db/schema.ts | 14 +- app/server/hp-agent.ts | 400 ++++++++++++++++++++++++++++ app/server/index.ts | 10 +- app/server/web/agent.ts | 447 -------------------------------- cmd/hp_agent/hp_agent.go | 3 + drizzle/0001_naive_lilith.sql | 5 + drizzle/meta/0001_snapshot.json | 73 ++++++ drizzle/meta/_journal.json | 7 + internal/hpagent/handler.go | 84 ++---- internal/tsnet/peers.go | 82 ++++++ mise.toml | 3 + 12 files changed, 619 insertions(+), 512 deletions(-) create mode 100644 app/server/hp-agent.ts delete mode 100644 app/server/web/agent.ts create mode 100644 drizzle/0001_naive_lilith.sql create mode 100644 drizzle/meta/0001_snapshot.json diff --git a/CHANGELOG.md b/CHANGELOG.md index c8f4e8b..8e01999 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,7 +4,7 @@ - This is built on top of a Go binary that runs in WebAssembly, using Xterm.js for the terminal interface. - Begin using a new SQLite database file in `/var/lib/headplane/hp_persist.db`. - The database is created automatically if it does not exist. - - It currently stores SSH connection details and will migrate older data. + - It currently stores SSH connection details and HostInfo for the agent. - The docker container now runs in a distroless image (closes [#255](https://github.com/tale/headplane/issues/255)). - A debug version of the container that runs as root and has a shell is available as `ghcr.io/tale/headplane:-shell`. - Removing a Split DNS record will no longer make the split domain unresolvable by clients (closes [#231](https://github.com/tale/headplane/issues/231)). @@ -18,6 +18,7 @@ - Environment variables are interpolatable into these paths - See the full reference in the [docs](https://github.com/tale/headplane/blob/main/docs/Configuration.md#sensitive-values) - The nix overlay build is fixed for the SSH module (via [#282](https://github.com/tale/headplane/pull/282)) +- Switch our build processes to use TypeScript Go and Rolldown Vite for better build and type-check performance. ### 0.6.0 (May 25, 2025) - Headplane 0.6.0 now requires **Headscale 0.26.0** or newer. diff --git a/app/server/db/schema.ts b/app/server/db/schema.ts index b0034fe..9aa6f69 100644 --- a/app/server/db/schema.ts +++ b/app/server/db/schema.ts @@ -1,4 +1,5 @@ -import { sqliteTable, text } from 'drizzle-orm/sqlite-core'; +import { integer, sqliteTable, text } from 'drizzle-orm/sqlite-core'; +import { HostInfo } from '~/types'; export const ephemeralNodes = sqliteTable('ephemeral_nodes', { auth_key: text('auth_key').primaryKey(), @@ -7,3 +8,14 @@ export const ephemeralNodes = sqliteTable('ephemeral_nodes', { export type EphemeralNode = typeof ephemeralNodes.$inferSelect; export type EphemeralNodeInsert = typeof ephemeralNodes.$inferInsert; + +export const hostInfo = sqliteTable('host_info', { + host_id: text('host_id').primaryKey(), + payload: text('payload', { mode: 'json' }).$type(), + updated_at: integer('updated_at', { mode: 'timestamp' }).$default( + () => new Date(), + ), +}); + +export type HostInfoRecord = typeof hostInfo.$inferSelect; +export type HostInfoInsert = typeof hostInfo.$inferInsert; diff --git a/app/server/hp-agent.ts b/app/server/hp-agent.ts new file mode 100644 index 0000000..c8050d8 --- /dev/null +++ b/app/server/hp-agent.ts @@ -0,0 +1,400 @@ +import { ChildProcessWithoutNullStreams, spawn } from 'node:child_process'; +import EventEmitter from 'node:events'; +import { constants, access, mkdir, open } from 'node:fs/promises'; +import { getegid, geteuid } from 'node:process'; +import { Interface, createInterface } from 'node:readline'; +import { inArray } from 'drizzle-orm'; +import { LibSQLDatabase } from 'drizzle-orm/libsql/driver-core'; +import { HostInfo } from '~/types'; +import log from '~/utils/log'; +import { HeadplaneConfig } from './config/schema'; +import { hostInfo } from './db/schema'; + +export async function createHeadplaneAgent( + config: NonNullable['agent'] | undefined, + headscaleUrl: string, + db: LibSQLDatabase, +) { + if (!config?.enabled) { + return; + } + + if (!config.pre_authkey) { + log.error('agent', 'Agent `pre_authkey` is not set'); + log.warn('agent', 'The agent will not run until resolved'); + return; + } + + try { + 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) { + // Try to create the directory just in case + try { + await mkdir(config.work_dir, { recursive: true }); + log.debug('config', 'Created agent work dir at %s', config.work_dir); + log.info( + 'config', + 'Created missing agent work dir at %s', + config.work_dir, + ); + + return; + } catch (innerError) { + log.error( + 'config', + 'Failed to create agent work dir at %s', + config.work_dir, + ); + log.info( + 'config', + 'Agent work dir not accessible at %s', + config.work_dir, + ); + log.debug('config', 'Error details: %s', error); + log.debug('config', 'Create error details: %s', innerError); + 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', + config.cache_path, + ); + log.debug('agent', 'Error details: %s', error); + return; + } + + const agent = new HeadplaneAgent({ + ...config, + headscaleUrl, + }); + + agent.on('spawn', () => { + log.info('agent', 'Headplane agent started'); + }); + + agent.on('ready', () => { + log.info('agent', 'Headplane agent is ready and serving queries'); + }); + + agent.on('error', (err) => { + log.warn('agent', 'Headplane agent experienced an error: %s', err.message); + log.debug('agent', 'Error details: %o', err); + }); + + agent.on('exit', ({ code, signal }) => { + log.warn( + 'agent', + 'Headplane agent exited with code %s and signal %s', + code, + signal, + ); + }); + + agent.on('restart', ({ delay, attempt }) => { + log.warn( + 'agent', + 'Headplane agent will restart in %f seconds (attempt %d)', + delay / 1000, + attempt, + ); + }); + + agent.on('stderr', (data) => { + log.error('agent', 'Headplane agent stderr:', data); + }); + + agent.on('info', async ({ id, info }) => { + log.debug('agent', 'Received HostInfo for %s', id); + try { + const parsedInfo = JSON.parse(info) as HostInfo; + await db + .insert(hostInfo) + .values({ + host_id: id, + payload: parsedInfo, + updated_at: new Date(), + }) + .onConflictDoUpdate({ + target: hostInfo.host_id, + set: { + payload: parsedInfo, + updated_at: new Date(), + }, + }); + } catch (error) { + log.error( + 'agent', + 'Failed to parse HostInfo for %s: %s', + id, + error instanceof Error ? error.message : String(error), + ); + return; + } + }); + + agent.start(); + + process.on('SIGTERM', () => agent.shutdown()); + process.on('SIGINT', () => agent.shutdown()); + + return { + agentID: () => agent.agentID(), + lookup: async (nodes: string[]) => { + const results = await db + .select() + .from(hostInfo) + .where(inArray(hostInfo.host_id, nodes)); + + return Object.fromEntries( + results.filter((r) => r.payload).map((r) => [r.host_id, r.payload]), + ) as Record; + }, + }; +} + +type AgentOptions = NonNullable< + NonNullable['agent'] +> & { + headscaleUrl: string; +}; + +interface AgentEvents { + ready: []; + spawn: []; + error: [Error]; + exit: [{ code?: number; signal?: NodeJS.Signals }]; + restart: [{ delay: number; attempt: number }]; + stderr: [string]; + info: [{ id: string; info: string }]; +} + +/** + * A custom class that turns the lifecycle of the agent into an event emitter. + * It has many different responsibilities ensuring that: + * - The agent is spawned with the correct configuration + * - The agent is ready and still running (ping/heartbeat) + * - The agent is restarted on a backoff strategy + */ +class HeadplaneAgent extends EventEmitter { + private child?: ChildProcessWithoutNullStreams; + private readline?: Interface; + + private options: AgentOptions; + + private hbInterval?: NodeJS.Timeout; + private hbDeadline?: NodeJS.Timeout; + private restartTimer?: NodeJS.Timeout; + private isWaitingForAck = false; + private isShuttingDown = false; + private backoffAttempt = 0; + private agentId?: string; + + private BASE_BACKOFF_MS = 1.5 * 1000; // 1.5 seconds + private MAX_BACKOFF_MS = 30 * 1000; // 30 seconds + private PROBE_COOLDOWN_MS = 5 * 60_000; // 5 minutes + private PROBE_ATTEMPT_INTERVAL = 10; // Every 10th attempt + + private HEARTBEAT_INTERVAL_MS = 5 * 1000; // 5 seconds + private HEARTBEAT_TIMEOUT_MS = 3 * 1000; // 3 seconds + + constructor(options: AgentOptions) { + super(); + this.options = options; + } + + agentID() { + return this.agentId; + } + + start() { + this.isShuttingDown = false; + this.spawnInternalChild(); + } + + shutdown() { + this.isShuttingDown = true; + this.agentId = undefined; + + clearTimeout(this.restartTimer); + clearInterval(this.hbInterval); + clearTimeout(this.hbDeadline); + this.isWaitingForAck = false; + + this.send('SHUTDOWN'); + this.child?.kill('SIGTERM'); + this.readline?.close(); + } + + private spawnInternalChild() { + this.child = spawn(this.options.executable_path, { + stdio: ['pipe', 'pipe', 'pipe'], + uid: geteuid?.() ?? undefined, + gid: getegid?.() ?? undefined, + env: { + HOME: process.env.HOME, + HEADPLANE_EMBEDDED: 'true', + HEADPLANE_AGENT_WORK_DIR: this.options.work_dir, + HEADPLANE_AGENT_DEBUG: log.debugEnabled ? 'true' : 'false', + HEADPLANE_AGENT_HOSTNAME: this.options.host_name, + HEADPLANE_AGENT_TS_SERVER: this.options.headscaleUrl, + HEADPLANE_AGENT_TS_AUTHKEY: this.options.pre_authkey, + }, + }); + + this.emit('spawn'); + this.child.on('error', (err) => this.emit('error', err)); + this.child.stderr.on('data', (data) => + this.emit('stderr', data.toString()), + ); + + this.child.on('exit', (code, signal) => { + this.agentId = undefined; + this.emit('exit', { + code: code ?? undefined, + signal: signal ?? undefined, + }); + + this.readline?.close(); + clearInterval(this.hbInterval); + clearTimeout(this.hbDeadline); + this.isWaitingForAck = false; + + if (this.isShuttingDown) { + log.info('agent', 'Child process exited gracefully'); + return; + } + + this.backoffAttempt++; + const delay = this.calculateBackoff(); + this.emit('restart', { delay, attempt: this.backoffAttempt }); + this.restartTimer = setTimeout(() => this.spawnInternalChild(), delay); + }); + + this.readline = createInterface({ input: this.child.stdout }); + this.readline.on('line', (line) => this.readlineHandler(line)); + this.send('START'); + + // Start the heartbeat loop with our custom interval + this.hbInterval = setInterval(() => { + if (!this.child || this.child.killed) return; + + // If we get here, we missed the last PONG response and can die + if (this.isWaitingForAck) { + this.agentId = undefined; + this.emit('error', new Error('Agent heartbeat missed')); + this.child.kill('SIGTERM'); + return; + } + + this.isWaitingForAck = true; + this.send('PING'); + + clearTimeout(this.hbDeadline); + this.hbDeadline = setTimeout(() => { + if (this.isWaitingForAck) { + this.agentId = undefined; + this.emit('error', new Error('Agent heartbeat timeout')); + this.child?.kill('SIGTERM'); + } + }, this.HEARTBEAT_TIMEOUT_MS); + }, this.HEARTBEAT_INTERVAL_MS); + } + + private send(s: string) { + if (!this.child || this.child.killed) return; + const ok = this.child.stdin.write(`${s}\n`); + if (!ok) this.child.stdin.once('drain', () => {}); + } + + /** + * Calculates a backoff time based on the current attempt. + * Supports a randomized jitter to avoid thundering herd problems. + * + * @param min The minimum backoff time in milliseconds. + * @param max The maximum backoff time in milliseconds. + * @returns The calculated backoff time in milliseconds. + */ + private calculateBackoff() { + const attempt = this.backoffAttempt; + if (attempt > 0 && attempt % this.PROBE_ATTEMPT_INTERVAL === 0) { + const jitter = Math.floor(Math.random() * (this.MAX_BACKOFF_MS + 1)); + const sign = Math.random() < 0.5 ? -1 : 1; + + return Math.max(0, this.PROBE_COOLDOWN_MS + jitter * sign); + } + + const cap = Math.min( + this.MAX_BACKOFF_MS, + this.BASE_BACKOFF_MS * 2 ** attempt, + ); + + return Math.floor(Math.random() * (cap + 1)); + } + + /** + * Processes and dispatches the appropriate response based on the message. + * @param line The message to process (piped straight from readline) + */ + private readlineHandler(line: string) { + // When we are ready we force a refresh so that the UI has the most + // up-to-date information and will gracefully handle new info being sent + if (line.startsWith('READY')) { + this.backoffAttempt = 0; + this.send('REFRESH'); + this.emit('ready'); + + const agentId = line.slice(5).trim(); + if (this.agentId && this.agentId !== agentId) { + log.warn( + 'agent', + 'Agent ID changed from %s to %s', + this.agentId, + agentId, + ); + } + + this.agentId = agentId; + return; + } + + if (line.startsWith('PONG')) { + this.isWaitingForAck = false; + clearTimeout(this.hbDeadline); + + const agentId = line.slice(5).trim(); + if (this.agentId && this.agentId !== agentId) { + log.warn( + 'agent', + 'Agent ID changed from %s to %s', + this.agentId, + agentId, + ); + } + + this.agentId = agentId; + return; + } + + if (line.startsWith('HOSTINFO')) { + const data = line.slice(9).trim(); + const [id, ...infoParts] = data.split(' '); + const info = infoParts.join(' '); + this.emit('info', { id, info }); + return; + } + + if (line.startsWith('ERROR')) { + const error = line.slice(6).trim(); + this.emit('error', new Error(error)); + return; + } + } +} diff --git a/app/server/index.ts b/app/server/index.ts index e169e95..acb6fdd 100644 --- a/app/server/index.ts +++ b/app/server/index.ts @@ -8,7 +8,7 @@ import { loadConfig } from './config/loader'; import { createDbClient } from './db/client.server'; import { createApiClient } from './headscale/api-client'; import { loadHeadscaleConfig } from './headscale/config-loader'; -import { loadAgentSocket } from './web/agent'; +import { createHeadplaneAgent } from './hp-agent'; import { createOidcClient } from './web/oidc'; import { createSessionStorage } from './web/sessions'; @@ -29,13 +29,13 @@ const config = await loadConfig( }), ); -const agentManager = await loadAgentSocket( +const db = await createDbClient(join(config.server.data_path, 'hp_persist.db')); +const agents = await createHeadplaneAgent( config.integration?.agent, config.headscale.url, + db, ); -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,7 +64,7 @@ const appLoadContext = { config.headscale.tls_cert_path, ), - agents: agentManager, + agents, integration: await loadIntegration(config.integration), oidc: config.oidc ? await createOidcClient(config.oidc) : undefined, db, diff --git a/app/server/web/agent.ts b/app/server/web/agent.ts deleted file mode 100644 index d090373..0000000 --- a/app/server/web/agent.ts +++ /dev/null @@ -1,447 +0,0 @@ -import { ChildProcess, spawn } from 'node:child_process'; -import { createHash } from 'node:crypto'; -import { - constants, - access, - mkdir, - open, - readFile, - writeFile, -} from 'node:fs/promises'; -import { exit, geteuid, getegid } from 'node:process'; -import { createInterface } from 'node:readline'; -import { setTimeout } from 'node:timers/promises'; -import { type } from 'arktype'; -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; -} - -interface MessageResponse { - Level: 'msg'; - Message: RegisterMessage | StatusMessage; -} - -type AgentResponse = LogResponse | MessageResponse; - -export async function loadAgentSocket( - config: NonNullable['agent'] | undefined, - headscaleUrl: string, -) { - 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 { - 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) { - // Try to create the directory just in case - try { - await mkdir(config.work_dir, { recursive: true }); - log.debug('config', 'Created agent work dir at %s', config.work_dir); - log.info( - 'config', - 'Created missing agent work dir at %s', - config.work_dir, - ); - - return; - } catch (innerError) { - log.error( - 'config', - 'Failed to create agent work dir at %s', - config.work_dir, - ); - log.info( - 'config', - 'Agent work dir not accessible at %s', - config.work_dir, - ); - log.debug('config', 'Error details: %s', error); - log.debug('config', 'Create error details: %s', innerError); - 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', - config.cache_path, - ); - log.debug('agent', 'Error details: %s', error); - return; - } - - const cache = new TimedCache(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; - private headscaleUrl: string; - private config: NonNullable< - NonNullable['agent'] - >; - - private spawnProcess: ChildProcess | null; - private agentId: string | null; - private uid: number | null; - private gid: number | null; - - constructor( - cache: TimedCache, - config: NonNullable['agent']>, - headscaleUrl: string, - ) { - this.cache = cache; - this.config = config; - this.headscaleUrl = headscaleUrl; - this.spawnProcess = null; - this.agentId = null; - this.startAgent(); - this.uid = geteuid ? geteuid() : null; - this.gid = getegid ? getegid() : null; - - process.on('SIGINT', () => { - this.spawnProcess?.kill(); - exit(0); - }); - - process.on('SIGTERM', () => { - this.spawnProcess?.kill(); - exit(0); - }); - } - - /** - * 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; - } - - /* - * 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, [], { - ...(this.uid ? { uid: this.uid } : {}), - ...(this.gid ? { gid: this.gid } : {}), - detached: false, - stdio: ['pipe', 'pipe', '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) { - await this.requestData(missing); - } - - return Object.entries(entries).reduce>( - (acc, [key, value]) => { - if (nodeIds.includes(key)) { - acc[key] = value; - } - - return acc; - }, - {}, - ); - } - - // 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; - } - - // Wait for the process to be spawned, busy waiting is gross - while (this.spawnProcess === null) { - await setTimeout(100); - } - - // 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`); - } -} - -const diskSchema = type({ - key: 'string', - value: 'unknown', - expires: 'number?', -}).array(); - -// A persistent HashMap with a TTL for each key -class TimedCache { - private _cache = new Map(); - private _timings = new Map(); - - // Default TTL is 1 minute - private defaultTTL: number; - private filePath: string; - private writeLock = false; - - // Last flush ID is essentially a hash of the flush contents - // Prevents unnecessary flushing if nothing has changed - private lastFlushId = ''; - - constructor(defaultTTL: number, filePath: string) { - this.defaultTTL = defaultTTL; - this.filePath = filePath; - - // Load the cache from disk and then queue flushes every 10 seconds - this.load().then(() => { - setInterval(() => this.flush(), 10000); - }); - } - - set(key: string, value: V, ttl: number = this.defaultTTL) { - this._cache.set(key, value); - this._timings.set(key, Date.now() + ttl); - } - - get(key: string) { - const value = this._cache.get(key); - if (!value) { - return; - } - - const expires = this._timings.get(key); - if (!expires || expires < Date.now()) { - this._cache.delete(key); - this._timings.delete(key); - return; - } - - return value; - } - - // Map into a Record without any TTLs - toJSON() { - const result: Record = {}; - for (const [key, value] of this._cache.entries()) { - result[key] = value; - } - - return result; - } - - // WARNING: This function expects that this.filePath is NOT ENOENT - private async load() { - const data = await readFile(this.filePath, 'utf-8'); - const cache = () => { - try { - return JSON.parse(data); - } catch (e) { - return undefined; - } - }; - - const diskData = cache(); - if (diskData === undefined) { - log.error('agent', 'Failed to load cache at %s', this.filePath); - return; - } - - const cacheData = diskSchema(diskData); - if (cacheData instanceof type.errors) { - log.debug('agent', 'Failed to load cache at %s', this.filePath); - log.debug('agent', 'Error details: %s', cacheData.toString()); - - // Skip loading the cache (it should be overwritten soon) - return; - } - - for (const { key, value, expires } of diskData) { - this._cache.set(key, value); - this._timings.set(key, expires); - } - - log.info('agent', 'Loaded cache from %s', this.filePath); - } - - private async flush() { - const data = Array.from(this._cache.entries()).map(([key, value]) => { - return { key, value, expires: this._timings.get(key) }; - }); - - if (data.length === 0) { - return; - } - - // Calculate the hash of the data - const dumpData = JSON.stringify(data); - const sha = createHash('sha256').update(dumpData).digest('hex'); - if (sha === this.lastFlushId) { - return; - } - - // We need to lock the writeLock so that we don't try to write - // to the file while we're already writing to it - while (this.writeLock) { - await setTimeout(100); - } - - this.writeLock = true; - await writeFile(this.filePath, dumpData, 'utf-8'); - log.debug('agent', 'Flushed cache to %s', this.filePath); - this.lastFlushId = sha; - this.writeLock = false; - } -} diff --git a/cmd/hp_agent/hp_agent.go b/cmd/hp_agent/hp_agent.go index b20c34d..c1400be 100644 --- a/cmd/hp_agent/hp_agent.go +++ b/cmd/hp_agent/hp_agent.go @@ -1,6 +1,8 @@ package main import ( + "fmt" + _ "github.com/joho/godotenv/autoload" "github.com/tale/headplane/internal/config" "github.com/tale/headplane/internal/hpagent" @@ -31,5 +33,6 @@ func main() { ID: agent.ID, }) + fmt.Println("READY") hpagent.FollowMaster(agent) } diff --git a/drizzle/0001_naive_lilith.sql b/drizzle/0001_naive_lilith.sql new file mode 100644 index 0000000..aaa61db --- /dev/null +++ b/drizzle/0001_naive_lilith.sql @@ -0,0 +1,5 @@ +CREATE TABLE `host_info` ( + `host_id` text PRIMARY KEY NOT NULL, + `payload` text, + `updated_at` integer +); diff --git a/drizzle/meta/0001_snapshot.json b/drizzle/meta/0001_snapshot.json new file mode 100644 index 0000000..4980bd6 --- /dev/null +++ b/drizzle/meta/0001_snapshot.json @@ -0,0 +1,73 @@ +{ + "version": "6", + "dialect": "sqlite", + "id": "16f780a3-a6e7-4810-94bb-fad5c6446ab4", + "prevId": "ab03ffcd-9aa5-4b4b-9f38-322acc6899a3", + "tables": { + "ephemeral_nodes": { + "name": "ephemeral_nodes", + "columns": { + "auth_key": { + "name": "auth_key", + "type": "text", + "primaryKey": true, + "notNull": true, + "autoincrement": false + }, + "node_key": { + "name": "node_key", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "checkConstraints": {} + }, + "host_info": { + "name": "host_info", + "columns": { + "host_id": { + "name": "host_id", + "type": "text", + "primaryKey": true, + "notNull": true, + "autoincrement": false + }, + "payload": { + "name": "payload", + "type": "text", + "primaryKey": false, + "notNull": false, + "autoincrement": false + }, + "updated_at": { + "name": "updated_at", + "type": "integer", + "primaryKey": false, + "notNull": false, + "autoincrement": false + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "checkConstraints": {} + } + }, + "views": {}, + "enums": {}, + "_meta": { + "schemas": {}, + "tables": {}, + "columns": {} + }, + "internal": { + "indexes": {} + } +} diff --git a/drizzle/meta/_journal.json b/drizzle/meta/_journal.json index dc3d0f8..e7c8c17 100644 --- a/drizzle/meta/_journal.json +++ b/drizzle/meta/_journal.json @@ -8,6 +8,13 @@ "when": 1750355487927, "tag": "0000_spicy_bloodscream", "breakpoints": true + }, + { + "idx": 1, + "version": "6", + "when": 1755554742267, + "tag": "0001_naive_lilith", + "breakpoints": true } ] } diff --git a/internal/hpagent/handler.go b/internal/hpagent/handler.go index 826b74f..c6d5327 100644 --- a/internal/hpagent/handler.go +++ b/internal/hpagent/handler.go @@ -2,26 +2,15 @@ package hpagent import ( "bufio" + "context" + "fmt" - "encoding/json" "os" - "sync" "github.com/tale/headplane/internal/tsnet" "github.com/tale/headplane/internal/util" - "tailscale.com/tailcfg" ) -// Represents messages from the Headplane master -type RecvMessage struct { - NodeIDs []string -} - -type SendMessage struct { - Type string - Data any -} - // Starts listening for messages from stdin func FollowMaster(agent *tsnet.TSAgent) { log := util.GetLogger() @@ -30,55 +19,34 @@ func FollowMaster(agent *tsnet.TSAgent) { for scanner.Scan() { line := scanner.Bytes() - var msg RecvMessage - err := json.Unmarshal(line, &msg) - if err != nil { - log.Error("Unable to decode message from master: %s", err) + directive := string(line) + + log.Debug("Received directive from master: %s", directive) + switch directive { + case "SHUTDOWN": + log.Debug("Received SHUTDOWN directive from master, shutting down agent") + agent.Shutdown() + return + + case "START": + log.Debug("Received START directive from master, starting agent") + // TODO: Start the agent here instead of in main + fmt.Println("READY " + agent.ID) continue - } - log.Debug("Recieved message from master: %v", line) - - if len(msg.NodeIDs) == 0 { - log.Debug("Message recieved had no node IDs") - log.Debug("Full message: %s", line) + case "PING": + log.Debug("Received PING directive from master, responding with PONG") + fmt.Println("PONG " + agent.ID) continue + + case "REFRESH": + log.Debug("Received REFRESH directive from master, refreshing status for all nodes") + err := agent.DispatchHostInfo(context.Background()) + if err != nil { + log.Error("Error refreshing host info: %s", err) + fmt.Println("ERR " + err.Error()) + } } - - // Accumulate the results since we invoke via gofunc - results := make(map[string]*tailcfg.HostinfoView) - mu := sync.Mutex{} - wg := sync.WaitGroup{} - - for _, nodeID := range msg.NodeIDs { - wg.Add(1) - go func(nodeID string) { - defer wg.Done() - result, err := agent.GetStatusForPeer(nodeID) - if err != nil { - log.Error("Unable to get status for node %s: %s", nodeID, err) - return - } - - if result == nil { - log.Debug("No status for node %s", nodeID) - return - } - - mu.Lock() - results[nodeID] = result - mu.Unlock() - }(nodeID) - } - - wg.Wait() - - // Send the results back to the Headplane master - log.Debug("Sending status back to master: %v", results) - log.Msg(&SendMessage{ - Type: "status", - Data: results, - }) } if err := scanner.Err(); err != nil { diff --git a/internal/tsnet/peers.go b/internal/tsnet/peers.go index 1ce2c2b..74a1ec1 100644 --- a/internal/tsnet/peers.go +++ b/internal/tsnet/peers.go @@ -3,10 +3,14 @@ package tsnet import ( "context" "encoding/hex" + "encoding/json" "fmt" "strings" + "sync" + "time" "github.com/tale/headplane/internal/util" + "tailscale.com/ipn/ipnstate" "tailscale.com/tailcfg" "tailscale.com/types/key" @@ -64,3 +68,81 @@ func (s *TSAgent) GetStatusForPeer(id string) (*tailcfg.HostinfoView, error) { log.Debug("Got whois for peer %s: %v", id, whois) return &whois.Node.Hostinfo, nil } + +// Dispatches ALL the HostInfo entries in our Tailnet to the master +func (s *TSAgent) DispatchHostInfo(ctx context.Context) error { + log := util.GetLogger() + + stat, err := s.Lc.Status(ctx) + if err != nil { + log.Debug("Failed to get status: %s", err) + return fmt.Errorf("failed to get status: %w", err) + } + + // Do lookups for all peers with a hint of parallelism for speed! + const maxParallel = 8 + sema := make(chan struct{}, maxParallel) + var wg sync.WaitGroup + var mu sync.Mutex + + nodeMap := make(map[key.NodePublic]*ipnstate.PeerStatus) + nodeMap[stat.Self.PublicKey] = stat.Self + for nodeKey, peer := range stat.Peer { + if peer == nil { + log.Debug("Skipping nil peer for node key: %s", nodeKey) + continue + } + + nodeMap[nodeKey] = peer + } + + for nodeKey, peer := range nodeMap { + idBytes, err := nodeKey.MarshalText() + if err != nil { + log.Debug("Failed to marshal node key: %s", err) + continue + } + + nodeID := string(idBytes) + wg.Add(1) + sema <- struct{}{} + + go func() { + defer wg.Done() + defer func() { <-sema }() + + wctx, cancel := context.WithTimeout(ctx, 3*time.Second) + defer cancel() + + ip := peer.TailscaleIPs[0].String() + if len(ip) == 0 { + log.Debug("Peer %s has no Tailscale IPs", nodeID) + return + } + + whois, err := s.Lc.WhoIs(wctx, ip) + if err != nil { + log.Debug("WhoIs failed for %s (%s): %s", nodeID, ip, err) + return + } + + if whois == nil || whois.Node == nil { + log.Debug("WhoIs returned nil node for %s (%s)", nodeID, ip) + return + } + + data, err := json.Marshal(whois.Node.Hostinfo) + if err != nil { + log.Debug("Failed to marshal hostinfo for %s (%s): %s", nodeID, ip, err) + return + } + + mu.Lock() + fmt.Println("HOSTINFO " + nodeID + " " + string(data)) + mu.Unlock() + }() + } + + wg.Wait() + return nil +} diff --git a/mise.toml b/mise.toml index a158439..345ecdf 100644 --- a/mise.toml +++ b/mise.toml @@ -23,6 +23,9 @@ run = "go build -o app/hp_ssh.wasm ./cmd/hp_ssh" alias = ["agent"] description = "Builds the Go agent for HostInfo" run = "go build -o build/hp_agent ./cmd/hp_agent" +# env = { CGO_ENABLED = "1" } +# run = "go build -o build/hp_agent.so -buildmode=c-shared ./cmd/hp_agent" + [tasks.build-fake-shell] alias = ["fake-shell"]