dahjah/headplane
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 1 Wiki Activity

feat: keep track of agent ids

Browse Source
This commit is contained in:
Aarnav Tale 2025-03-06 17:37:41 -05:00
parent 5dd4c41291
commit d9314887fb
No known key found for this signature in database
3 changed files with 14 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
server/context/app.ts
View File

@ -1,19 +1,22 @@
import type { HostInfo } from '~/types';
import { TimedCache } from '~server/ws/cache';
import { hp_agentRequest, hp_getAgentCache } from '~server/ws/data';
import { hp_getAgents } from '~server/ws/socket';
import { hp_getConfig } from './loader';
import type { HeadplaneConfig } from './parser';
export interface AppContext {
context: HeadplaneConfig;
agentData?: TimedCache<HostInfo>;
hp_agentRequest: typeof hp_agentRequest;
agents: string[];
agentData?: TimedCache<HostInfo>;
}
export default function appContext() {
export default function appContext(): AppContext {
return {
context: hp_getConfig(),
agentData: hp_getAgentCache(),
hp_agentRequest,
agents: [...hp_getAgents().keys()],
agentData: hp_getAgentCache(),
};
}

6
server/ws/data.ts
View File

@ -28,8 +28,7 @@ export async function hp_agentRequest(nodeList: string[]) {
// Request to all connected agents (we can have multiple)
// Luckily we can parse all the data at once through message parsing
// and then overlapping cache entries will be overwritten by time
const agents = [...hp_getAgents()];
console.log(agents);
const agents = hp_getAgents();
// Deduplicate the list of nodes
const NodeIDs = [...new Set(nodeList)];
@ -40,7 +39,7 @@ export async function hp_agentRequest(nodeList: string[]) {
// Await so that data loads on first request without racing
// Since we do agent.once() we NEED to wait for it to finish
await Promise.allSettled(
agents.map(async (agent) => {
[...agents].map(async ([id, agent]) => {
agent.send(JSON.stringify({ NodeIDs }));
await new Promise<void>((resolve) => {
// Just as a safety measure, we set a maximum timeout of 3 seconds
@ -48,6 +47,7 @@ export async function hp_agentRequest(nodeList: string[]) {
agent.once('message', (data) => {
const parsed = JSON.parse(data.toString());
log.debug('CACH', 'Received agent data from %s', id);
for (const [node, info] of Object.entries<HostInfo>(parsed)) {
cache?.set(node, info);
log.debug('CACH', 'Cached %s', node);

7
server/ws/socket.ts
View File

@ -10,7 +10,7 @@ export function initWebsocket(authKey: string) {
log.info('SRVX', 'Starting a WebSocket server for agent connections');
server.on('connection', (ws, req) => {
const tailnetID = req.headers['x-headplane-tailnet-id'];
if (!tailnetID) {
if (!tailnetID || typeof tailnetID !== 'string') {
log.warn(
'SRVX',
'Rejecting an agent WebSocket connection without a tailnet ID',
@ -29,6 +29,7 @@ export function initWebsocket(authKey: string) {
return;
}
agents.set(tailnetID, ws);
const pinger = setInterval(() => {
if (ws.readyState !== WebSocket.OPEN) {
clearInterval(pinger);
@ -40,6 +41,7 @@ export function initWebsocket(authKey: string) {
ws.on('close', () => {
clearInterval(pinger);
agents.delete(tailnetID);
});
ws.on('error', (error) => {
@ -54,6 +56,7 @@ export function initWebsocket(authKey: string) {
return server;
}
const agents = new Map<string, WebSocket>();
export function hp_getAgents() {
return server.clients;
return agents;
}