feat: cleanup removal of old ssh plexer and logic

we also have added the necessary logic to auto prune ephemeral nodes because
headscale doesn't seem to automatically remove them. this change made use of a database
which is now stored in the persistent headplane directory.

app/routes/ssh/console.tsx

@@ -2,17 +2,22 @@ import { faker } from '@faker-js/faker';
 import { Loader2 } from 'lucide-react';
 import { useEffect, useState } from 'react';
 import {
+	ActionFunctionArgs,
 	LoaderFunctionArgs,
 	ShouldRevalidateFunction,
 	data,
 	useLoaderData,
+	useSubmit,
 } from 'react-router';
+import { useEventSource } from 'remix-utils/sse/react';
 import { LoadContext } from '~/server';
-import { PreAuthKey, User } from '~/types';
+import { Machine, PreAuthKey, User } from '~/types';
 import { useLiveData } from '~/utils/live-data';
 import XTerm from './xterm.client';
 
+import { eq } from 'drizzle-orm';
 import wasm from '~/hp_ssh.wasm?url';
+import { EphemeralNodeInsert, ephemeralNodes } from '~/server/db/schema';
 import '~/wasm_exec';
 
 export const shouldRevalidate: ShouldRevalidateFunction = () => {
@@ -54,12 +59,14 @@ export async function loader({
 		}
 		return subject === user.subject;
 	});
+
 	if (!lookup) {
 		throw data(
 			`User with subject ${user.subject} not found within Headscale`,
 			404,
 		);
 	}
+
 	const { preAuthKey } = await context.client.post<{ preAuthKey: PreAuthKey }>(
 		'v1/preauthkey',
 		session.get('api_key')!,
@@ -67,30 +74,71 @@ export async function loader({
 			user: lookup.id,
 			reusable: false,
 			ephemeral: true,
-			expiration: new Date(Date.now() + 60 * 60 * 1000).toISOString(), // 1 hour
+			expiration: new Date(Date.now() + 60 * 1000).toISOString(), // 1 minute
 		},
 	);
+
 	// TODO: Enable config to enforce generate_authkeys capability
 	// For now, any user is capable of WebSSH connections
 	// const check = await context.sessions.check(
 	// 	request,
 	// 	Capabilities.generate_authkeys,
 	// );
+
 	const qp = new URL(request.url).searchParams;
 	const username = qp.get('username') || undefined;
 	const hostname = qp.get('hostname') || undefined;
 	if (!username || !hostname) {
 		throw data('Missing required parameters: username, hostname', 400);
 	}
-	// TODO: Verify the Host headers to ensure CORS friendly
-	// TODO: Check if the URL actually resolves correctly
-	// TODO: Keep track of hostname since ephemeral nodes are broken atm
+
+	// We're making a request to <url>/key?v=116 to check the CORS headers
+	const u = context.config.headscale.public_url ?? context.config.headscale.url;
+	// const res = await fetch(`${u}/key?v=116`, {
+	// 	method: 'GET',
+	// });
+
+	// const corsOrigin = res.headers.get('Access-Control-Allow-Origin');
+	// const corsMethods = res.headers.get('Access-Control-Allow-Methods');
+	// const corsHeaders = res.headers.get('Access-Control-Allow-Headers');
+	// console.log(corsOrigin, corsMethods, corsHeaders);
+
+	// if (!corsOrigin || !corsMethods || !corsHeaders) {
+	// 	throw data(
+	// 		'Headscale server does not have the required CORS headers for WebSSH',
+	// 		500,
+	// 	);
+	// }
+
+	const { nodes } = await context.client.get<{ nodes: Machine[] }>(
+		'v1/node',
+		session.get('api_key')!,
+	);
+
+	// node.name is the hostname, given_name is the set name
+	const lookupNode = nodes.find((n) => n.name === hostname);
+	if (!lookupNode) {
+		throw data(`Node with hostname ${hostname} not found`, 404);
+	}
+
+	// Last thing is keeping track of the ephemeral node in the database
+	// because Headscale doesn't automatically delete ephemeral nodes???
+	const [ephemeralNode] = await context.db
+		.insert(ephemeralNodes)
+		.values({
+			auth_key: preAuthKey.key,
+		} satisfies EphemeralNodeInsert)
+		.returning();
+	console.log('Created ephemeral node:', ephemeralNode);
+
 	return {
-		PreAuthKey: preAuthKey.key,
-		ControlURL:
-			context.config.headscale.public_url ?? context.config.headscale.url,
-		Hostname: generateHostname(username),
-		ssh: {
+		ipnDetails: {
+			PreAuthKey: preAuthKey.key,
+			Hostname: generateHostname(username),
+			ControlURL: u,
+		},
+
+		sshDetails: {
 			username,
 			hostname,
 		},
@@ -115,11 +163,46 @@ function generateHostname(username: string) {
 	return `ssh-${adjective}-${noun}-${username}`;
 }
 
+export async function action({
+	request,
+	context,
+}: ActionFunctionArgs<LoadContext>) {
+	const session = await context.sessions.auth(request);
+	if (!context.agents?.agentID()) {
+		throw data(
+			'WebSSH is only available with the Headplane agent integration',
+			400,
+		);
+	}
+
+	const form = await request.formData();
+	const nodeKey = form.get('node_key');
+	const authKey = form.get('auth_key');
+
+	if (nodeKey === null || typeof nodeKey !== 'string') {
+		throw data('Missing node_key', 400);
+	}
+
+	if (authKey === null || typeof authKey !== 'string') {
+		throw data('Missing auth_key', 400);
+	}
+
+	await context.db
+		.update(ephemeralNodes)
+		.set({
+			node_key: nodeKey,
+		})
+		.where(eq(ephemeralNodes.auth_key, authKey));
+}
+
 export default function Page() {
+	const submit = useSubmit();
 	const { pause } = useLiveData();
+
 	const [ipn, setIpn] = useState<TsWasmNet | null>(null);
-	const { PreAuthKey, ControlURL, Hostname, ssh } =
-		useLoaderData<typeof loader>();
+	const [nodeKey, setNodeKey] = useState<string | null>(null);
+	const { ipnDetails, sshDetails } = useLoaderData<typeof loader>();
+	const ping = useEventSource('/ssh/ping', { event: 'ping' });
 
 	useEffect(() => {
 		pause();
@@ -127,30 +210,39 @@ export default function Page() {
 		WebAssembly.instantiateStreaming(fetch(wasm), go.importObject).then(
 			(value) => {
 				go.run(value.instance);
-				const handle = TsWasmNet(
-					{
-						PreAuthKey,
-						ControlURL,
-						Hostname,
+				const handle = TsWasmNet(ipnDetails, {
+					NotifyState: (state) => {
+						console.log('State changed:', state);
+						if (state === 'Running') {
+							setIpn(handle);
+						}
 					},
-					{
-						NotifyState: (state) => {
-							console.log('State changed:', state);
-							if (state === 'Running') {
-								setIpn(handle);
-							}
-						},
-						NotifyNetMap: (netmap) => {
-							console.log('NetMap updated:', netmap);
-						},
-						NotifyBrowseToURL: (url) => {
-							console.log('Browse to URL:', url);
-						},
-						NotifyPanicRecover: (message) => {
-							console.error('Panic recover:', message);
-						},
+					NotifyNetMap: (netmap) => {
+						// Only set NodeKey if it is not already set and then
+						// also dispatch that to the backend to track the
+						// ephemeral node.
+						//
+						// We open an SSE connection to the backend
+						// so that when the connection is closed,
+						// the backend can delete the ephemeral node.
+						if (nodeKey === null) {
+							setNodeKey(netmap.NodeKey);
+							submit(
+								{
+									node_key: netmap.NodeKey,
+									auth_key: ipnDetails.PreAuthKey,
+								},
+								{ method: 'POST' },
+							);
+						}
 					},
-				);
+					NotifyBrowseToURL: (url) => {
+						console.log('Browse to URL:', url);
+					},
+					NotifyPanicRecover: (message) => {
+						console.error('Panic recover:', message);
+					},
+				});
 
 				handle.Start();
 			},
@@ -165,7 +257,11 @@ export default function Page() {
 				</div>
 			) : (
 				<div className="flex flex-col h-screen">
-					<XTerm ipn={ipn} username={ssh.username} hostname={ssh.hostname} />
+					<XTerm
+						ipn={ipn}
+						username={sshDetails.username}
+						hostname={sshDetails.hostname}
+					/>
 				</div>
 			)}
 		</div>

app/routes/ssh/hp_ssh.d.ts

@@ -11,11 +11,15 @@ interface TsWasmNetOptions {
 
 interface TsWasmNetCallbacks {
 	NotifyState: (state: IPNState) => void;
-	NotifyNetMap: (netMapJson: string) => void;
+	NotifyNetMap: (netmap: TsWasmNetMap) => void;
 	NotifyBrowseToURL: (url: string) => void;
 	NotifyPanicRecover: (err: string) => void;
 }
 
+interface TsWasmNetMap {
+	NodeKey: string;
+}
+
 interface TsWasmNet {
 	Start: () => void;
 	OpenSSH: (