mirror of
https://github.com/tale/headplane.git
synced 2026-08-22 02:36:38 +00:00
feat: use a new ws implementation thats encapsulated
This commit is contained in:
Aarnav Tale
@@ -0,0 +1,126 @@
|
||||
import { createHash } from 'node:crypto';
|
||||
import { readFile, writeFile } from 'node:fs/promises';
|
||||
import { type } from 'arktype';
|
||||
import log from '~server/utils/log';
|
||||
import mutex from '~server/utils/mutex';
|
||||
|
||||
const diskSchema = type({
|
||||
key: 'string',
|
||||
value: 'unknown',
|
||||
expires: 'number?',
|
||||
}).array();
|
||||
|
||||
// A persistent HashMap with a TTL for each key
|
||||
export class TimedCache<V> {
|
||||
private _cache = new Map<string, V>();
|
||||
private _timings = new Map<string, number>();
|
||||
|
||||
// Default TTL is 1 minute
|
||||
private defaultTTL: number;
|
||||
private filePath: string;
|
||||
private writeLock = mutex();
|
||||
|
||||
// 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<string, V> = {};
|
||||
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('CACH', 'Failed to load cache at %s', this.filePath);
|
||||
return;
|
||||
}
|
||||
|
||||
const cacheData = diskSchema(diskData);
|
||||
if (cacheData instanceof type.errors) {
|
||||
log.error('CACH', 'Failed to load cache at %s', this.filePath);
|
||||
log.debug('CACHE', '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('CACH', 'Loaded cache from %s', this.filePath);
|
||||
}
|
||||
|
||||
private async flush() {
|
||||
this.writeLock.acquire();
|
||||
const data = Array.from(this._cache.entries()).map(([key, value]) => {
|
||||
return { key, value, expires: this._timings.get(key) };
|
||||
});
|
||||
|
||||
if (data.length === 0) {
|
||||
this.writeLock.release();
|
||||
return;
|
||||
}
|
||||
|
||||
// Calculate the hash of the data
|
||||
const dumpData = JSON.stringify(data);
|
||||
const sha = createHash('sha256').update(dumpData).digest('hex');
|
||||
if (sha === this.lastFlushId) {
|
||||
this.writeLock.release();
|
||||
return;
|
||||
}
|
||||
|
||||
await writeFile(this.filePath, dumpData, 'utf-8');
|
||||
this.lastFlushId = sha;
|
||||
this.writeLock.release();
|
||||
log.debug('CACH', 'Flushed cache to %s', this.filePath);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,61 @@
|
||||
import { open } from 'node:fs/promises';
|
||||
import type { HostInfo } from '~/types';
|
||||
import log from '~server/utils/log';
|
||||
import { TimedCache } from './cache';
|
||||
import { hp_getAgents } from './socket';
|
||||
|
||||
let cache: TimedCache<HostInfo> | undefined;
|
||||
export async function hp_loadAgentCache(defaultTTL: number, filepath: string) {
|
||||
log.debug('CACH', `Loading agent cache from ${filepath}`);
|
||||
|
||||
try {
|
||||
const handle = await open(filepath, 'w');
|
||||
log.info('CACH', `Using agent cache file at ${filepath}`);
|
||||
await handle.close();
|
||||
} catch (e) {
|
||||
log.info('CACH', `Agent cache file not found at ${filepath}`);
|
||||
return;
|
||||
}
|
||||
|
||||
cache = new TimedCache(defaultTTL, filepath);
|
||||
}
|
||||
|
||||
export function hp_getAgentCache() {
|
||||
return cache;
|
||||
}
|
||||
|
||||
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);
|
||||
|
||||
// Deduplicate the list of nodes
|
||||
const NodeIDs = [...new Set(nodeList)];
|
||||
NodeIDs.map((node) => {
|
||||
log.debug('CACH', 'Requesting agent data for', node);
|
||||
});
|
||||
|
||||
// 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) => {
|
||||
agent.send(JSON.stringify({ NodeIDs }));
|
||||
await new Promise<void>((resolve) => {
|
||||
// Just as a safety measure, we set a maximum timeout of 3 seconds
|
||||
setTimeout(() => resolve(), 3000);
|
||||
|
||||
agent.once('message', (data) => {
|
||||
const parsed = JSON.parse(data.toString());
|
||||
for (const [node, info] of Object.entries<HostInfo>(parsed)) {
|
||||
cache?.set(node, info);
|
||||
log.debug('CACH', 'Cached %s', node);
|
||||
}
|
||||
|
||||
resolve();
|
||||
});
|
||||
});
|
||||
}),
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,59 @@
|
||||
import WebSocket, { WebSocketServer } from 'ws';
|
||||
import log from '~server/utils/log';
|
||||
|
||||
const server = new WebSocketServer({ noServer: true });
|
||||
export function initWebsocket(authKey: string) {
|
||||
if (authKey.length === 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
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) {
|
||||
log.warn(
|
||||
'SRVX',
|
||||
'Rejecting an agent WebSocket connection without a tailnet ID',
|
||||
);
|
||||
ws.close(1008, 'ERR_INVALID_TAILNET_ID');
|
||||
return;
|
||||
}
|
||||
|
||||
if (req.headers.authorization !== `Bearer ${authKey}`) {
|
||||
log.warn('SRVX', 'Rejecting an unauthorized WebSocket connection');
|
||||
if (req.socket.remoteAddress) {
|
||||
log.warn('SRVX', 'Agent source IP: %s', req.socket.remoteAddress);
|
||||
}
|
||||
|
||||
ws.close(1008, 'ERR_UNAUTHORIZED');
|
||||
return;
|
||||
}
|
||||
|
||||
const pinger = setInterval(() => {
|
||||
if (ws.readyState !== WebSocket.OPEN) {
|
||||
clearInterval(pinger);
|
||||
return;
|
||||
}
|
||||
|
||||
ws.ping();
|
||||
}, 30000);
|
||||
|
||||
ws.on('close', () => {
|
||||
clearInterval(pinger);
|
||||
});
|
||||
|
||||
ws.on('error', (error) => {
|
||||
clearInterval(pinger);
|
||||
log.error('SRVX', 'Agent WebSocket error: %s', error);
|
||||
log.debug('SRVX', 'Error details: %o', error);
|
||||
log.error('SRVX', 'Closing agent WebSocket connection');
|
||||
ws.close(1011, 'ERR_INTERNAL_ERROR');
|
||||
});
|
||||
});
|
||||
|
||||
return server;
|
||||
}
|
||||
|
||||
export function hp_getAgents() {
|
||||
return server.clients;
|
||||
}
|
||||
Reference in New Issue
Block a user