feat: cleanup oidc logic and surface errors better

app/routes/auth/login/config-error.tsx

@@ -0,0 +1,123 @@
+import { AlertCircle } from 'lucide-react';
+import Card from '~/components/Card';
+import Code from '~/components/Code';
+import Link from '~/components/Link';
+import { OidcConnectorError } from '~/server/web/oidc-connector';
+
+export function OidcConfigErrorNotice({
+	errors,
+}: {
+	errors: OidcConnectorError[];
+}) {
+	return (
+		<Card className="max-w-md m-4 sm:m-0 mb-4 sm:mb-4 border border-red-500">
+			<div className="flex items-center justify-between gap-4">
+				<Card.Title className="text-red-500">Authentication Error</Card.Title>
+				<AlertCircle className="w-6 h-6 mb-2 text-red-500" />
+			</div>
+			<Card.Text className="text-sm">
+				The OpenID Connect (OIDC) Single Sign-On (SSO) configuration has issues:{' '}
+				<ul className="list-disc list-inside mt-2 mb-1">
+					{mapOidcErrorsToMessages(errors).map((code) => (
+						<li key={code.key}>{code.node}</li>
+					))}
+				</ul>{' '}
+				<Link
+					name="Headplane OIDC Issues"
+					to="https://headplane.net/configuration/sso#help"
+				>
+					Learn more
+				</Link>
+			</Card.Text>
+		</Card>
+	);
+}
+
+function mapOidcErrorsToMessages(errors: OidcConnectorError[]) {
+	const messages: {
+		key: string;
+		node: React.ReactNode;
+	}[] = [];
+
+	for (const error of errors) {
+		switch (error) {
+			case 'INVALID_API_KEY':
+				messages.push({
+					key: error,
+					node: (
+						<Card.Text className="inline">
+							The provided API key for OIDC authentication is invalid. Ensure
+							that <Code>oidc.headscale_api_key</Code> is a valid API key.
+						</Card.Text>
+					),
+				});
+				break;
+
+			case 'MISSING_AUTHORIZATION_ENDPOINT':
+				messages.push({
+					key: error,
+					node: (
+						<Card.Text className="inline">
+							The OIDC provided does not have a configured{' '}
+							<Code>authorization_endpoint</Code>. Ensure discovery URL or
+							manual configuration is correct.
+						</Card.Text>
+					),
+				});
+				break;
+
+			case 'MISSING_TOKEN_ENDPOINT':
+				messages.push({
+					key: error,
+					node: (
+						<Card.Text className="inline">
+							The OIDC provided does not have a configured{' '}
+							<Code>token_endpoint</Code>. Ensure discovery URL or manual
+							configuration is correct.
+						</Card.Text>
+					),
+				});
+				break;
+
+			case 'MISSING_USERINFO_ENDPOINT':
+				messages.push({
+					key: error,
+					node: (
+						<Card.Text className="inline">
+							The OIDC provided does not have a configured{' '}
+							<Code>user_endpoint</Code>. Ensure discovery URL or manual
+							configuration is correct.
+						</Card.Text>
+					),
+				});
+				break;
+
+			case 'MISSING_REQUIRED_CLAIMS':
+				messages.push({
+					key: error,
+					node: (
+						<Card.Text className="inline">
+							The OIDC provider does not support the <Code>sub</Code> claim,
+							which is required for authentication. Your OIDC provider may be
+							misconfigured.
+						</Card.Text>
+					),
+				});
+				break;
+
+			case 'UNKNOWN_ERROR':
+				messages.push({
+					key: error,
+					node: (
+						<Card.Text className="inline">
+							An unknown error occurred during OIDC configuration. Please check
+							the Headplane logs for more information.
+						</Card.Text>
+					),
+				});
+				break;
+		}
+	}
+
+	return messages;
+}

app/routes/auth/login/oidc-error.tsx

@@ -0,0 +1,63 @@
+import { AlertCircle } from 'lucide-react';
+import Card from '~/components/Card';
+import Code from '~/components/Code';
+
+export function OidcErrorNotice({ code }: { code: string }) {
+	return (
+		<Card className="max-w-md m-4 sm:m-0 mb-4 sm:mb-4 border border-red-500">
+			<div className="flex items-center justify-between gap-4">
+				<Card.Title className="text-red-500">Configuration Issue(s)</Card.Title>
+				<AlertCircle className="w-6 h-6 mb-2 text-red-500" />
+			</div>
+			{getErrorMessage(code)}
+		</Card>
+	);
+}
+
+function getErrorMessage(code: string) {
+	switch (code) {
+		case 'error_no_query':
+			return (
+				<Card.Text>
+					The SSO provider did not correctly redirect back to Headplane with the
+					required parameters. Please ensure your SSO provider is configured
+					correctly.
+				</Card.Text>
+			);
+
+		case 'error_no_session':
+		case 'error_invalid_session':
+			return (
+				<Card.Text>
+					Unable to complete SSO login due to missing or invalid session data.
+					Ensure that your Headplane cookie configuration is correct and that
+					your browser is accepting cookies.
+				</Card.Text>
+			);
+
+		case 'error_no_sub':
+			return (
+				<Card.Text>
+					The SSO provider did not return a valid user identifier. Please ensure
+					your SSO provider is correctly configured to provide the{' '}
+					<Code>sub</Code> claim.
+				</Card.Text>
+			);
+
+		case 'error_auth_failed':
+			return (
+				<Card.Text>
+					Authentication with the SSO provider failed. Please tray again later.
+					Headplane logs may provide more information.
+				</Card.Text>
+			);
+
+		default:
+			return (
+				<Card.Text>
+					An unknown error occurred during OIDC authentication. Please try again
+					later.
+				</Card.Text>
+			);
+	}
+}

app/routes/auth/login/page.tsx

@@ -12,10 +12,13 @@ import Card from '~/components/Card';
 import Code from '~/components/Code';
 import Input from '~/components/Input';
 import Link from '~/components/Link';
+import { OidcConnectorError } from '~/server/web/oidc-connector';
 import { useLiveData } from '~/utils/live-data';
 import type { Route } from './+types/page';
 import { loginAction } from './action';
+import { OidcConfigErrorNotice } from './config-error';
 import Logout from './logout';
+import { OidcErrorNotice } from './oidc-error';
 
 export async function loader({ request, context }: Route.LoaderArgs) {
 	try {
@@ -26,28 +29,21 @@ export async function loader({ request, context }: Route.LoaderArgs) {
 	const qp = new URL(request.url).searchParams;
 	const urlState = qp.get('s') ?? undefined;
 
-	// OIDC config cannot be undefined if an OIDC client is set
-	// Also check if we are in a logout state and skip redirect if we are
-	const ssoOnly = context.config.oidc?.disable_api_key_login;
-	if (urlState !== 'logout' && ssoOnly) {
-		// This shouldn't be possible, but still a safe sanity check
-		if (!context.oidc || typeof context.oidc === 'string') {
-			throw data(
-				'`oidc.disable_api_key_login` was set without a valid OIDC configuration',
-				{
-					status: 400,
-				},
-			);
-		}
-
+	// MARK: This works because the OIDC connector will always return false
+	// for `isExclusive` if the OIDC config isn't usable.
+	if (context.oidcConnector?.isExclusive && urlState !== 'logout') {
 		return redirect('/oidc/start');
 	}
 
+	const isOidcConnectorEnabled = context.oidcConnector?.isValid;
+	const oidcErrorCodes = !isOidcConnectorEnabled
+		? context.oidcConnector!.errors
+		: [];
+
 	return {
 		isCookieSecureEnabled: context.config.server.cookie_secure,
-		isOidcEnabled: context.oidc !== undefined,
-		oidcErrorMessage:
-			typeof context.oidc === 'string' ? context.oidc : undefined,
+		isOidcConnectorEnabled,
+		oidcErrorCodes,
 		urlState,
 	};
 }
@@ -55,8 +51,13 @@ export async function loader({ request, context }: Route.LoaderArgs) {
 export const action = loginAction;
 
 export default function Page({ loaderData, actionData }: Route.ComponentProps) {
-	const { isOidcEnabled, isCookieSecureEnabled, oidcErrorMessage, urlState } =
-		loaderData;
+	const {
+		isCookieSecureEnabled,
+		isOidcConnectorEnabled,
+		oidcErrorCodes,
+		urlState,
+	} = loaderData;
+
 	const [showCookieWarning, setShowCookieWarning] = useState(false);
 	const [params] = useSearchParams();
 	const { pause } = useLiveData();
@@ -95,40 +96,31 @@ export default function Page({ loaderData, actionData }: Route.ComponentProps) {
 	return (
 		<div className="flex w-screen h-screen items-center justify-center">
 			<div>
-				{showCookieWarning || oidcErrorMessage ? (
+				{urlState?.startsWith('error_') ? (
+					<OidcErrorNotice code={urlState} />
+				) : oidcErrorCodes.length > 0 ? (
+					<OidcConfigErrorNotice errors={oidcErrorCodes} />
+				) : showCookieWarning ? (
 					<Card className="max-w-md m-4 sm:m-0 mb-4 sm:mb-4 border border-red-500">
 						<div className="flex items-center justify-between gap-4">
 							<Card.Title className="text-red-500">
-								Configuration Issue(s)
+								Configuration Issue
 							</Card.Title>
 							<AlertCircle className="w-6 h-6 mb-2 text-red-500" />
 						</div>
-						<div className="flex flex-col gap-2">
-							{showCookieWarning ? (
-								<Card.Text className="text-sm text-red-600 dark:text-red-400">
-									Headplane is configured to use secure cookies, but this site
-									is being served over an insecure connection and login will not
-									work correctly.{' '}
-									<Link
-										name="Headplane Common Issues"
-										to="https://headplane.net/configuration/common-issues#issue-logging-in-does-not-do-anything"
-									>
-										Learn more.
-									</Link>
-								</Card.Text>
-							) : undefined}
-							{oidcErrorMessage ? (
-								<Card.Text className="text-sm text-red-600 dark:text-red-400">
-									{oidcErrorMessage}{' '}
-									<Link
-										name="Headplane OIDC Issues"
-										to="https://headplane.net/configuration/sso#help"
-									>
-										Learn more.
-									</Link>
-								</Card.Text>
-							) : undefined}
-						</div>
+						{showCookieWarning ? (
+							<Card.Text className="text-sm">
+								Headplane is configured to use secure cookies, but this site is
+								being served over an insecure connection and login will not work
+								correctly.{' '}
+								<Link
+									name="Headplane Common Issues"
+									to="https://headplane.net/configuration/common-issues#issue-logging-in-does-not-do-anything"
+								>
+									Learn more.
+								</Link>
+							</Card.Text>
+						) : undefined}
 					</Card>
 				) : undefined}
 				<Card className="max-w-md m-4 sm:m-0">
@@ -157,11 +149,11 @@ export default function Page({ loaderData, actionData }: Route.ComponentProps) {
 							Sign In
 						</Button>
 					</Form>
-					{isOidcEnabled ? (
+					{isOidcConnectorEnabled ? (
 						<RemixLink to="/oidc/start">
 							<Button
 								className="w-full mt-2"
-								isDisabled={oidcErrorMessage !== undefined}
+								isDisabled={oidcErrorCodes.length > 0}
 								variant="light"
 							>
 								Single Sign-On

app/routes/auth/oidc-callback.ts

@@ -1,75 +1,99 @@
 import { createHash } from 'node:crypto';
 import { count, eq } from 'drizzle-orm';
-import { createCookie, type LoaderFunctionArgs, redirect } from 'react-router';
+import * as oidc from 'openid-client';
+import {
+	createCookie,
+	data,
+	type LoaderFunctionArgs,
+	redirect,
+} from 'react-router';
 import { ulid } from 'ulidx';
 import type { LoadContext } from '~/server';
-import { HeadplaneConfig } from '~/server/config/config-schema';
 import { users } from '~/server/db/schema';
 import { Roles } from '~/server/web/roles';
-import { FlowUser, finishAuthFlow, formatError } from '~/utils/oidc';
-import { send } from '~/utils/res';
-
-interface OidcFlowSession {
-	state: string;
-	nonce: string;
-	code_verifier: string;
-	redirect_uri: string;
-}
+import log from '~/utils/log';
+import type { OidcCookieState } from './oidc-start';
 
 export async function loader({
 	request,
 	context,
 }: LoaderFunctionArgs<LoadContext>) {
-	if (!context.oidc || typeof context.oidc === 'string') {
-		throw new Error('OIDC is not enabled');
+	if (!context.oidcConnector?.isValid) {
+		throw data('OIDC is not enabled or misconfigured', { status: 501 });
 	}
 
-	// Check if we have 0 query parameters
 	const url = new URL(request.url);
 	if (url.searchParams.toString().length === 0) {
-		return redirect('/login');
+		return redirect('/login?s=error_no_query');
 	}
 
 	const cookie = createCookie('__oidc_auth_flow', {
 		httpOnly: true,
-		maxAge: 300, // 5 minutes
+		maxAge: 300,
+		secure: context.config.server.cookie_secure,
+		domain: context.config.server.cookie_domain,
 	});
 
-	const data: OidcFlowSession | null = await cookie.parse(
+	const oidcCookieState: OidcCookieState | null = await cookie.parse(
 		request.headers.get('Cookie'),
 	);
 
-	if (data === null) {
-		console.warn('OIDC flow session not found');
-		return redirect('/login');
+	if (oidcCookieState == null || typeof oidcCookieState !== 'object') {
+		log.warn('auth', 'Called OIDC callback without session cookie');
+		return redirect('/login?s=error_no_session');
 	}
 
-	const { code_verifier, state, nonce, redirect_uri } = data;
-	if (!code_verifier || !state || !nonce || !redirect_uri) {
-		return send({ error: 'Missing OIDC state' }, { status: 400 });
+	const { state, nonce } = oidcCookieState;
+	if (!state || !nonce) {
+		log.warn('auth', 'OIDC session cookie is missing required fields');
+		return redirect('/login?s=error_invalid_session');
 	}
 
-	// Reconstruct the redirect URI using the query parameters
-	// and the one we saved in the session
-	const flowRedirectUri = new URL(redirect_uri);
-	flowRedirectUri.search = url.search;
-
-	const flowOptions = {
-		redirect_uri: flowRedirectUri.toString(),
-		code_verifier,
-		state,
-		nonce: nonce === '<none>' ? undefined : nonce,
-	};
-
 	try {
-		let user = await finishAuthFlow(context.oidc, flowOptions);
-		user = {
-			...user,
-			picture: setOidcPictureForSource(
-				user,
-				context.config.oidc?.profile_picture_source ?? 'oidc',
-			),
-		};
+		const tokens = await oidc.authorizationCodeGrant(
+			context.oidcConnector.client,
+			request,
+			{
+				expectedState: state,
+				expectedNonce: nonce,
+			},
+		);
+
+		const claims = tokens.claims();
+		if (claims?.sub == null) {
+			log.warn('auth', 'No subject found in OIDC claims');
+			return redirect('/login?s=error_no_sub');
+		}
+
+		const userInfo = await oidc.fetchUserInfo(
+			context.oidcConnector.client,
+			tokens.access_token,
+			claims.sub,
+		);
+
+		// We have defaults that closely follow what Headscale uses, maybe we
+		// can make it configurable in the future, but for now we only need the
+		// `sub` claim.
+		const username =
+			userInfo.preferred_username ?? userInfo.email?.split('@')[0] ?? 'user';
+		const name =
+			userInfo.name ??
+			(userInfo.given_name && userInfo.family_name
+				? `${userInfo.given_name} ${userInfo.family_name}`
+				: (userInfo.preferred_username ?? 'SSO User'));
+
+		const picture =
+			context.config.oidc?.profile_picture_source === 'gravatar'
+				? (() => {
+						if (!userInfo.email) {
+							return undefined;
+						}
+
+						const emailHash = userInfo.email.trim().toLowerCase();
+						const hash = createHash('sha256').update(emailHash).digest('hex');
+						return `https://www.gravatar.com/avatar/${hash}?s=200&d=identicon&r=x`;
+					})()
+				: userInfo.picture;
 
 		const [{ count: userCount }] = await context.db
 			.select({ count: count() })
@@ -80,52 +104,46 @@ export async function loader({
 			.insert(users)
 			.values({
 				id: ulid(),
-				sub: user.subject,
+				sub: claims.sub,
 				caps: userCount === 0 ? Roles.owner : Roles.member,
 			})
 			.onConflictDoNothing();
 
-		return redirect('/machines', {
+		return redirect('/', {
 			headers: {
 				'Set-Cookie': await context.sessions.createSession({
-					// 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.
-					api_key: context.config.oidc!.headscale_api_key,
-					user,
+					api_key: context.oidcConnector.apiKey,
+					user: {
+						subject: claims.sub,
+						username,
+						name,
+						email: userInfo.email,
+						picture,
+					},
 				}),
 			},
 		});
 	} catch (error) {
-		return new Response(JSON.stringify(formatError(error)), {
-			status: 500,
-			headers: {
-				'Content-Type': 'application/json',
-			},
-		});
-	}
-}
-
-type PictureSource = NonNullable<
-	HeadplaneConfig['oidc']
->['profile_picture_source'];
-
-function setOidcPictureForSource(user: FlowUser, source: PictureSource) {
-	// Already set by default in the callback, so we can just return it
-	if (source === 'oidc') {
-		return user.picture;
-	}
-
-	if (source === 'gravatar') {
-		if (!user.email) {
-			return undefined;
+		if (error instanceof oidc.ResponseBodyError) {
+			log.error(
+				'auth',
+				'Got an OIDC response error body: %s',
+				JSON.stringify(error.cause),
+			);
 		}
 
-		const emailHash = user.email.trim().toLowerCase();
-		const hash = createHash('sha256').update(emailHash).digest('hex');
-		return `https://www.gravatar.com/avatar/${hash}?s=200&d=identicon&r=x`;
-	}
+		if (error instanceof oidc.AuthorizationResponseError) {
+			log.error(
+				'auth',
+				'Got an OIDC authorization response error: %s',
+				error.error,
+			);
+		}
 
-	return undefined;
+		if (error instanceof oidc.WWWAuthenticateChallengeError) {
+			log.error('auth', 'Got an OIDC WWW-Authenticate challenge error');
+		}
+
+		return redirect('/login?s=error_auth_failed');
+	}
 }

app/routes/auth/oidc-start.ts

@@ -1,6 +1,16 @@
-import { createCookie, type LoaderFunctionArgs, redirect } from 'react-router';
+import * as oidc from 'openid-client';
+import {
+	createCookie,
+	data,
+	type LoaderFunctionArgs,
+	redirect,
+} from 'react-router';
 import type { LoadContext } from '~/server';
-import { beginAuthFlow, getRedirectUri } from '~/utils/oidc';
+
+export interface OidcCookieState {
+	nonce: string;
+	state: string;
+}
 
 export async function loader({
 	request,
@@ -8,40 +18,63 @@ export async function loader({
 }: LoaderFunctionArgs<LoadContext>) {
 	try {
 		await context.sessions.auth(request);
-		return redirect('/machines');
+		return redirect('/');
 	} catch {}
 
-	if (
-		!context.oidc ||
-		typeof context.oidc === 'string' ||
-		!context.config.oidc
-	) {
-		throw new Error('OIDC is not enabled');
+	if (!context.oidcConnector?.isValid) {
+		throw data('OIDC is not enabled or misconfigured', { status: 501 });
 	}
 
 	const cookie = createCookie('__oidc_auth_flow', {
 		httpOnly: true,
-		maxAge: 300, // 5 minutes
+		maxAge: 300,
+		secure: context.config.server.cookie_secure,
+		domain: context.config.server.cookie_domain,
 	});
 
 	const redirectUri =
 		context.config.oidc?.redirect_uri ?? getRedirectUri(request);
-	const data = await beginAuthFlow(
-		context.oidc,
-		redirectUri,
-		context.config.oidc.scope,
-		context.config.oidc.extra_params,
-	);
 
-	return redirect(data.url, {
+	const nonce = oidc.randomNonce();
+	const state = oidc.randomState();
+
+	const url = oidc.buildAuthorizationUrl(context.oidcConnector.client, {
+		...(context.oidcConnector.extraParams ?? {}),
+		scope: context.oidcConnector.scope,
+		redirect_uri: redirectUri,
+		state,
+		nonce,
+	});
+
+	return redirect(url.href, {
 		status: 302,
 		headers: {
 			'Set-Cookie': await cookie.serialize({
-				state: data.state,
-				nonce: data.nonce,
-				code_verifier: data.codeVerifier,
-				redirect_uri: redirectUri,
-			}),
+				state,
+				nonce,
+			} satisfies OidcCookieState),
 		},
 	});
 }
+
+function getRedirectUri(req: Request) {
+	const url = new URL(`${__PREFIX__}/oidc/callback`, req.url);
+	let host = req.headers.get('Host');
+	if (!host) {
+		host = req.headers.get('X-Forwarded-Host');
+	}
+
+	if (!host) {
+		throw data(
+			'Cannot determine redirect URI: no Host or X-Forwarded-Host header',
+			{
+				status: 500,
+			},
+		);
+	}
+
+	const proto = req.headers.get('X-Forwarded-Proto');
+	url.protocol = proto ?? 'http:';
+	url.host = host;
+	return url.href;
+}

app/routes/settings/overview.tsx

@@ -1,26 +1,18 @@
 import { ArrowRight } from 'lucide-react';
-import {
-	LoaderFunctionArgs,
-	Link as RemixLink,
-	useLoaderData,
-} from 'react-router';
+import { Link as RemixLink } from 'react-router';
 import Link from '~/components/Link';
-import { LoadContext } from '~/server';
+import type { Route } from './+types/overview';
 
-export async function loader({ context }: LoaderFunctionArgs<LoadContext>) {
+export async function loader({ context }: Route.LoaderArgs) {
 	return {
 		config: context.hs.writable(),
-		oidc: context.oidc
-			? typeof context.oidc === 'string'
-				? undefined
-				: context.oidc
-			: undefined,
+		isOidcEnabled: context.oidcConnector?.isValid ?? false,
 	};
 }
 
-export default function Page() {
-	const { config, oidc } = useLoaderData<typeof loader>();
-
+export default function Page({
+	loaderData: { config, isOidcEnabled },
+}: Route.ComponentProps) {
 	return (
 		<div className="flex flex-col gap-8 max-w-(--breakpoint-lg)">
 			<div className="flex flex-col w-2/3">
@@ -51,7 +43,7 @@ export default function Page() {
 					<ArrowRight className="w-5 h-5 ml-2" />
 				</div>
 			</RemixLink>
-			{config && oidc ? (
+			{config && isOidcEnabled ? (
 				<>
 					<div className="flex flex-col w-2/3">
 						<h1 className="text-2xl font-medium mb-4">