From d9314887fb87f1421736309dcd31289fc28a6998 Mon Sep 17 00:00:00 2001
From: Aarnav Tale <aarnavtale@icloud.com>
Date: Thu, 6 Mar 2025 17:37:41 -0500
Subject: [PATCH] feat: keep track of agent ids

---
 server/context/app.ts | 9 ++++++---
 server/ws/data.ts     | 6 +++---
 server/ws/socket.ts   | 7 +++++--
 3 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/server/context/app.ts b/server/context/app.ts
index 23449ec..729886f 100644
--- a/server/context/app.ts
+++ b/server/context/app.ts
@@ -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(),
 	};
 }
diff --git a/server/ws/data.ts b/server/ws/data.ts
index e2245d3..38664fa 100644
--- a/server/ws/data.ts
+++ b/server/ws/data.ts
@@ -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);
diff --git a/server/ws/socket.ts b/server/ws/socket.ts
index 6fc10d0..11716cc 100644
--- a/server/ws/socket.ts
+++ b/server/ws/socket.ts
@@ -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;
 }