From 1b45b0917fd986b18bf6a334f44d78c177b169c1 Mon Sep 17 00:00:00 2001
From: Aarnav Tale <aarnavtale@icloud.com>
Date: Fri, 17 Jan 2025 11:45:36 +0000
Subject: [PATCH] fix: handle headscale unavailability gracefully

---
 app/components/Error.tsx         |  5 +-
 app/components/Footer.tsx        |  2 +-
 app/components/Header.tsx        | 12 ++--
 app/layouts/dashboard.tsx        | 98 ++++++++++++++++----------------
 app/layouts/shell.tsx            | 33 +++++++++++
 app/root.tsx                     | 25 +++++++-
 app/routes.ts                    | 26 +++++----
 app/routes/auth/login.tsx        |  7 +--
 app/routes/machines/overview.tsx |  7 +++
 9 files changed, 135 insertions(+), 80 deletions(-)
 create mode 100644 app/layouts/shell.tsx

diff --git a/app/components/Error.tsx b/app/components/Error.tsx
index 98756b9..e52b3c4 100644
--- a/app/components/Error.tsx
+++ b/app/components/Error.tsx
@@ -8,7 +8,6 @@ interface Props {
 	type?: 'full' | 'embedded';
 }
 
-
 function getMessage(error: Error | unknown) {
 	if (!(error instanceof Error)) {
 		return "An unknown error occurred";
@@ -18,9 +17,9 @@ function getMessage(error: Error | unknown) {
 
 	// Traverse the error chain to find the root cause
 	if (error.cause) {
-		rootError = error.cause;
+		rootError = error.cause as Error;
 		while (rootError.cause) {
-			rootError = rootError.cause;
+			rootError = rootError.cause as Error;
 		}
 	}
 
diff --git a/app/components/Footer.tsx b/app/components/Footer.tsx
index 9f4d118..703f221 100644
--- a/app/components/Footer.tsx
+++ b/app/components/Footer.tsx
@@ -15,7 +15,7 @@ export default function Footer({ url, debug }: FooterProps) {
 	return (
 		<footer
 			className={cn(
-				'fixed bottom-0 left-0 z-50 w-full h-14',
+				'fixed bottom-0 left-0 z-40 w-full h-14',
 				'bg-ui-100 dark:bg-ui-900 text-ui-500',
 				'flex flex-col justify-center gap-1',
 				'border-t border-ui-200 dark:border-ui-800',
diff --git a/app/components/Header.tsx b/app/components/Header.tsx
index 6cf8ab2..05536e6 100644
--- a/app/components/Header.tsx
+++ b/app/components/Header.tsx
@@ -17,10 +17,8 @@ import Menu from './Menu';
 import TabLink from './TabLink';
 
 interface Props {
-	data?: {
-		config: HeadplaneContext['config'];
-		user?: SessionData['user'];
-	};
+	config: HeadplaneContext['config'];
+	user?: SessionData['user'];
 }
 
 interface LinkProps {
@@ -45,7 +43,7 @@ function Link({ href, text, isMenu }: LinkProps) {
 	);
 }
 
-export default function Header({ data }: Props) {
+export default function Header(data: Props) {
 	return (
 		<header className="bg-main-700 dark:bg-main-800 text-ui-50">
 			<div className="container flex items-center justify-between py-4">
@@ -57,7 +55,7 @@ export default function Header({ data }: Props) {
 					<Link href="https://tailscale.com/download" text="Download" />
 					<Link href="https://github.com/tale/headplane" text="GitHub" />
 					<Link href="https://github.com/juanfont/headscale" text="Headscale" />
-					{data?.user ? (
+					{data.user ? (
 						<Menu>
 							<Menu.Button
 								className={cn(
@@ -122,7 +120,7 @@ export default function Header({ data }: Props) {
 					name="Access Control"
 					icon={<LockIcon className="w-4 h-4" />}
 				/>
-				{data?.config.read ? (
+				{data.config.read ? (
 					<>
 						<TabLink
 							to="/dns"
diff --git a/app/layouts/dashboard.tsx b/app/layouts/dashboard.tsx
index f6546b4..7e57389 100644
--- a/app/layouts/dashboard.tsx
+++ b/app/layouts/dashboard.tsx
@@ -1,77 +1,77 @@
 import { type LoaderFunctionArgs, redirect } from 'react-router';
-import { Outlet, useLoaderData, useNavigation } from 'react-router';
-import { ProgressBar } from 'react-aria-components';
+import { Outlet, useLoaderData } from 'react-router';
+import { useEffect } from 'react'
 
 import { ErrorPopup } from '~/components/Error';
 import Header from '~/components/Header';
+import { toast } from '~/components/Toaster';
 import Footer from '~/components/Footer';
 import Link from '~/components/Link';
+import { useLiveData } from '~/utils/useLiveData'
 import { cn } from '~/utils/cn';
 import { loadContext } from '~/utils/config/headplane';
-import { HeadscaleError, pull } from '~/utils/headscale';
+import { HeadscaleError, pull, healthcheck } from '~/utils/headscale';
 import { destroySession, getSession } from '~/utils/sessions.server';
+import { XCircleFillIcon } from '@primer/octicons-react';
+import log from '~/utils/log';
 
 export async function loader({ request }: LoaderFunctionArgs) {
-	const session = await getSession(request.headers.get('Cookie'));
-	if (!session.has('hsApiKey')) {
-		return redirect('/login');
-	}
-
+	let healthy = false;
 	try {
-		await pull('v1/apikey', session.get('hsApiKey')!);
+		healthy = await healthcheck();
 	} catch (error) {
-		if (error instanceof HeadscaleError) {
-			// Safest to just redirect to login if we can't pull
-			return redirect('/login', {
-				headers: {
-					'Set-Cookie': await destroySession(session),
-				},
-			});
-		}
-
-		// Otherwise propagate to boundary
-		throw error;
+		log.debug('APIC', 'Healthcheck failed %o', error);
+	}
+
+	// We shouldn't session invalidate if Headscale is down
+	if (healthy) {
+		// We can assert because shell ensures this is set
+		const session = await getSession(request.headers.get('Cookie'));
+		const apiKey = session.get('hsApiKey')!;
+
+		try {
+			await pull('v1/apikey', apiKey);
+		} catch (error) {
+			if (error instanceof HeadscaleError) {
+				log.debug('APIC', 'API Key validation failed %o', error);
+				return redirect('/login', {
+					headers: {
+						'Set-Cookie': await destroySession(session),
+					},
+				});
+			}
+		}
 	}
 
-	const context = await loadContext();
 	return {
-		config: context.config,
-		url: context.headscalePublicUrl ?? context.headscaleUrl,
-		debug: context.debug,
-		user: session.get('user'),
-	};
+		healthy,
+	}
 }
 
 export default function Layout() {
-	const data = useLoaderData<typeof loader>();
-	const nav = useNavigation();
+	useLiveData({ interval: 3000 });
+	const { healthy } = useLoaderData<typeof loader>()
 
 	return (
 		<>
-			<ProgressBar aria-label="Loading...">
-				<div
-					className={cn(
-						'fixed top-0 left-0 z-50 w-1/2 h-1',
-						'bg-blue-500 dark:bg-blue-400 opacity-0',
-						nav.state === 'loading' && 'animate-loading opacity-100',
-					)}
-				/>
-			</ProgressBar>
-			<Header data={data} />
+			{!healthy ? (
+				<div className={cn(
+					'fixed bottom-0 right-0 z-50 w-fit h-14',
+					'flex flex-col justify-center gap-1',
+				)}>
+					<div className={cn(
+						'flex items-center gap-1.5 mr-1.5 py-2 px-1.5',
+						'border rounded-lg text-white bg-red-500',
+						'border-red-600 dark:border-red-400 shadow-sm',
+					)}>
+						<XCircleFillIcon className="w-4 h-4 text-white" />
+						Headscale is unreachable
+					</div>
+				</div>
+			) : undefined}
 			<main className="container mx-auto overscroll-contain mt-4 mb-24">
 				<Outlet />
 			</main>
-			<Footer {...data} />
-		</>
-	);
-}
-
-export function ErrorBoundary() {
-	return (
-		<>
-			<Header />
-			<ErrorPopup type="embedded" />
-			<Footer url="Unknown" debug={false} />
 		</>
 	);
 }
diff --git a/app/layouts/shell.tsx b/app/layouts/shell.tsx
new file mode 100644
index 0000000..6b7e2d9
--- /dev/null
+++ b/app/layouts/shell.tsx
@@ -0,0 +1,33 @@
+import Header from '~/components/Header';
+import Footer from '~/components/Footer';
+import { getSession } from '~/utils/sessions.server';
+import { loadContext } from '~/utils/config/headplane';
+import { useLoaderData, LoaderFunctionArgs, Outlet, redirect } from 'react-router';
+
+// This loads the bare minimum for the application to function
+// So we know that if context fails to load then well, oops?
+export async function loader({ request }: LoaderFunctionArgs) {
+	const session = await getSession(request.headers.get('Cookie'));
+	if (!session.has('hsApiKey')) {
+		return redirect('/login');
+	}
+
+	const context = await loadContext();
+	return {
+		config: context.config,
+		url: context.headscalePublicUrl ?? context.headscaleUrl,
+		debug: context.debug,
+		user: session.get('user'),
+	};
+}
+
+export default function Shell() {
+	const data = useLoaderData<typeof loader>();
+	return (
+		<>
+			<Header {...data} />
+			<Outlet />
+			<Footer {...data} />
+		</>
+	)
+}
diff --git a/app/root.tsx b/app/root.tsx
index 97eb2c8..70d49db 100644
--- a/app/root.tsx
+++ b/app/root.tsx
@@ -1,9 +1,13 @@
-import type { LinksFunction, MetaFunction } from 'react-router';
-import { Links, Meta, Outlet, Scripts, ScrollRestoration } from 'react-router';
+import type { LoaderFunctionArgs, LinksFunction, MetaFunction } from 'react-router';
+import { Links, Meta, Outlet, Scripts, ScrollRestoration, useNavigation } from 'react-router';
+import { loadContext } from '~/utils/config/headplane';
 
+import { ProgressBar } from 'react-aria-components';
 import { ErrorPopup } from '~/components/Error';
+// TODO: Make this a default export
 import { Toaster } from '~/components/Toaster';
 import stylesheet from '~/tailwind.css?url';
+import { cn } from '~/utils/cn';
 
 export const meta: MetaFunction = () => [
 	{ title: 'Headplane' },
@@ -41,5 +45,20 @@ export function ErrorBoundary() {
 }
 
 export default function App() {
-	return <Outlet />;
+	const nav = useNavigation();
+
+	return (
+		<>
+			<ProgressBar aria-label="Loading...">
+				<div
+					className={cn(
+						'fixed top-0 left-0 z-50 w-1/2 h-1',
+						'bg-blue-500 dark:bg-blue-400 opacity-0',
+						nav.state === 'loading' && 'animate-loading opacity-100',
+					)}
+				/>
+			</ProgressBar>
+			<Outlet />
+		</>
+	)
 }
diff --git a/app/routes.ts b/app/routes.ts
index 499f98f..2f0cf34 100644
--- a/app/routes.ts
+++ b/app/routes.ts
@@ -12,19 +12,23 @@ export default [
 	route('/oidc/start', 'routes/auth/oidc-start.ts'),
 
 	// All the main logged-in dashboard routes
-	layout('layouts/dashboard.tsx', [
-		...prefix('/machines', [
-			index('routes/machines/overview.tsx'),
-			route('/:id', 'routes/machines/machine.tsx'),
-		]),
+	// Double nested to separate error propagations
+	layout('layouts/shell.tsx', [
+		layout('layouts/dashboard.tsx', [
+			...prefix('/machines', [
+				index('routes/machines/overview.tsx'),
+				route('/:id', 'routes/machines/machine.tsx'),
+			]),
 
-		route('/users', 'routes/users/overview.tsx'),
-		route('/acls', 'routes/acls/editor.tsx'),
-		route('/dns', 'routes/dns/overview.tsx'),
+			route('/users', 'routes/users/overview.tsx'),
+			route('/acls', 'routes/acls/editor.tsx'),
+			route('/dns', 'routes/dns/overview.tsx'),
 
-		...prefix('/settings', [
-			index('routes/settings/overview.tsx'),
-			route('/auth-keys', 'routes/settings/auth-keys.tsx'),
+			...prefix('/settings', [
+				index('routes/settings/overview.tsx'),
+				route('/auth-keys', 'routes/settings/auth-keys.tsx'),
+				route('/local-agent', 'routes/settings/local-agent.tsx'),
+			]),
 		]),
 	]),
 ];
diff --git a/app/routes/auth/login.tsx b/app/routes/auth/login.tsx
index 535efa0..e593f74 100644
--- a/app/routes/auth/login.tsx
+++ b/app/routes/auth/login.tsx
@@ -13,15 +13,10 @@ import TextField from '~/components/TextField';
 import type { Key } from '~/types';
 import { loadContext } from '~/utils/config/headplane';
 import { pull } from '~/utils/headscale';
-import {
-	startOidc,
-	beginAuthFlow,
-	getRedirectUri
-} from '~/utils/oidc';
+import { beginAuthFlow, getRedirectUri } from '~/utils/oidc';
 import { commitSession, getSession } from '~/utils/sessions.server';
 
 export async function loader({ request }: LoaderFunctionArgs) {
-
 	const session = await getSession(request.headers.get('Cookie'));
 	if (session.has('hsApiKey')) {
 		return redirect('/machines', {
diff --git a/app/routes/machines/overview.tsx b/app/routes/machines/overview.tsx
index adf22d4..027c748 100644
--- a/app/routes/machines/overview.tsx
+++ b/app/routes/machines/overview.tsx
@@ -13,6 +13,7 @@ import { getSession } from '~/utils/sessions.server';
 import { useLiveData } from '~/utils/useLiveData';
 import type { Machine, Route, User } from '~/types';
 import { queryAgent, initAgentSocket } from '~/utils/ws-agent';
+import { ErrorPopup } from '~/components/Error'
 
 import { menuAction } from './action';
 import MachineRow from './components/machine';
@@ -139,3 +140,9 @@ export default function Page() {
 		</>
 	);
 }
+
+export function ErrorBoundary() {
+	return (
+		<ErrorPopup type="embedded" />
+	)
+}