feat: rework oidc to be more resilient

This includes setting a custom redirect URI, handling errors, and using a better library.
As an API decision I've also disabled per session API keys as it clutters up too much.

app/routes/auth/oidc-callback.ts

@@ -1,17 +1,77 @@
-import { type LoaderFunctionArgs, data } from 'react-router';
+import { type LoaderFunctionArgs, redirect } from 'react-router';
 import { loadContext } from '~/utils/config/headplane';
-import { finishOidc } from '~/utils/oidc';
+import { getSession, commitSession } from '~/utils/sessions.server';
+import { finishAuthFlow, getRedirectUri, formatError } from '~/utils/oidc';
+import { send } from '~/utils/res';
 
 export async function loader({ request }: LoaderFunctionArgs) {
-	try {
-		const context = await loadContext();
-		if (!context.oidc) {
-			throw new Error('An invalid OIDC configuration was provided');
-		}
+	// Check if we have 0 query parameters
+	const url = new URL(request.url);
+	if (url.searchParams.toString().length === 0) {
+		return redirect('/machines');
+	}
 
-		return finishOidc(context.oidc, request);
+	const session = await getSession(request.headers.get('Cookie'));
+	if (session.has('hsApiKey')) {
+		return redirect('/machines')
+	}
+
+	// This is a hold-over from the old code
+	// TODO: Rewrite checkOIDC in the context loader
+	const { oidc } = await loadContext();
+	if (!oidc) {
+		throw new Error('An invalid OIDC configuration was provided');
+	}
+
+	const oidcConfig = {
+		issuer: oidc.issuer,
+		clientId: oidc.client,
+		clientSecret: oidc.secret,
+		redirectUri: oidc.redirectUri,
+		tokenEndpointAuthMethod: oidc.method,
+	}
+
+	const codeVerifier = session.get('oidc_code_verif');
+	const state = session.get('oidc_state');
+	const nonce = session.get('oidc_nonce');
+
+	if (!codeVerifier || !state || !nonce) {
+		return send({ error: 'Missing OIDC state' }, { status: 400 });
+	}
+
+	const flowOptions = {
+		redirect_uri: request.url,
+		codeVerifier,
+		state,
+		nonce: nonce === '<none>' ? undefined : nonce,
+	}
+
+	try {
+		const user = await finishAuthFlow(oidcConfig, flowOptions);
+		session.set('user', user);
+		session.unset('oidc_code_verif');
+		session.unset('oidc_state');
+		session.unset('oidc_nonce');
+
+		// TODO: This is breaking, to stop the "over-generation" of API
+		// keys because they are currently non-deletable in the headscale
+		// database. Look at this in the future once we have a solution
+		// or we have permissioned API keys.
+		session.set('hsApiKey', oidc.rootKey);
+		return redirect('/machines', {
+			headers: {
+				'Set-Cookie': await commitSession(session),
+			},
+		});
 	} catch (error) {
-		// Gracefully present OIDC errors
-		return data({ error }, { status: 500 });
+		return new Response(
+			JSON.stringify(formatError(error)),
+			{
+				status: 500,
+				headers: {
+					'Content-Type': 'application/json',
+				},
+			}
+		);
 	}
 }