feat: pkce

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

@@ -47,7 +47,7 @@ function getErrorMessage(code: string) {
 		case 'error_auth_failed':
 			return (
 				<Card.Text>
-					Authentication with the SSO provider failed. Please tray again later.
+					Authentication with the SSO provider failed. Please try again later.
 					Headplane logs may provide more information.
 				</Card.Text>
 			);

app/routes/auth/login/page.tsx

@@ -1,7 +1,6 @@
 import { AlertCircle } from 'lucide-react';
 import { useEffect, useState } from 'react';
 import {
-	data,
 	Form,
 	Link as RemixLink,
 	redirect,
@@ -12,7 +11,6 @@ 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';

app/routes/auth/oidc-callback.ts

@@ -1,18 +1,15 @@
 import { createHash } from 'node:crypto';
 import { count, eq } from 'drizzle-orm';
 import * as oidc from 'openid-client';
-import { data, type LoaderFunctionArgs, redirect } from 'react-router';
+import { data, redirect } from 'react-router';
 import { ulid } from 'ulidx';
-import type { LoadContext } from '~/server';
 import { users } from '~/server/db/schema';
 import { Roles } from '~/server/web/roles';
 import log from '~/utils/log';
 import { createOidcStateCookie } from '~/utils/oidc-state';
+import type { Route } from './+types/oidc-callback';
 
-export async function loader({
-	request,
-	context,
-}: LoaderFunctionArgs<LoadContext>) {
+export async function loader({ request, context }: Route.LoaderArgs) {
 	if (!context.oidcConnector?.isValid) {
 		throw data('OIDC is not enabled or misconfigured', { status: 501 });
 	}
@@ -30,8 +27,8 @@ export async function loader({
 		return redirect('/login?s=error_no_session');
 	}
 
-	const { state, nonce, redirect_uri } = oidcCookieState;
-	if (!state || !nonce || !redirect_uri) {
+	const { state, nonce, redirect_uri, verifier } = oidcCookieState;
+	if (!state || !nonce || !redirect_uri || !verifier) {
 		log.warn('auth', 'OIDC session cookie is missing required fields');
 		return redirect('/login?s=error_invalid_session');
 	}
@@ -47,6 +44,9 @@ export async function loader({
 			{
 				expectedState: state,
 				expectedNonce: nonce,
+				...(context.oidcConnector.usePKCE
+					? { pkceCodeVerifier: verifier }
+					: {}),
 			},
 		);
 

app/routes/auth/oidc-start.ts

@@ -1,13 +1,10 @@
 import * as oidc from 'openid-client';
-import { data, type LoaderFunctionArgs, redirect } from 'react-router';
-import type { LoadContext } from '~/server';
+import { data, redirect } from 'react-router';
 import { HeadplaneConfig } from '~/server/config/config-schema';
 import { createOidcStateCookie } from '~/utils/oidc-state';
+import type { Route } from './+types/oidc-start';
 
-export async function loader({
-	request,
-	context,
-}: LoaderFunctionArgs<LoadContext>) {
+export async function loader({ request, context }: Route.LoaderArgs) {
 	try {
 		await context.sessions.auth(request);
 		return redirect('/');
@@ -21,6 +18,7 @@ export async function loader({
 	const redirect_uri = getRedirectUri(context.config, request);
 
 	const nonce = oidc.randomNonce();
+	const verifier = oidc.randomPKCECodeVerifier();
 	const state = oidc.randomState();
 
 	const url = oidc.buildAuthorizationUrl(context.oidcConnector.client, {
@@ -29,6 +27,12 @@ export async function loader({
 		redirect_uri,
 		state,
 		nonce,
+		...(context.oidcConnector.usePKCE
+			? {
+					code_challenge_method: 'S256',
+					code_challenge: await oidc.calculatePKCECodeChallenge(verifier),
+				}
+			: {}),
 	});
 
 	return redirect(url.href, {
@@ -37,6 +41,7 @@ export async function loader({
 			'Set-Cookie': await cookie.serialize({
 				state,
 				nonce,
+				verifier,
 				redirect_uri,
 			}),
 		},

app/server/config/config-schema.ts

@@ -65,6 +65,7 @@ const oidcConfig = type({
 	client_id: 'string',
 	client_secret: 'string',
 	headscale_api_key: 'string',
+	use_pkce: 'boolean = false',
 	redirect_uri: type('string.url')
 		.pipe((value, ctx) => {
 			log.warn(
@@ -111,6 +112,7 @@ const partialOidcConfig = type({
 	issuer: 'string.url?',
 	client_id: 'string?',
 	client_secret: 'string?',
+	use_pkce: 'boolean?',
 	headscale_api_key: 'string?',
 	redirect_uri: 'string.url?',
 	disable_api_key_login: 'boolean?',

app/server/web/oidc-connector.ts

@@ -25,6 +25,7 @@ export type OidcConnector =
 	| {
 			isValid: true;
 			isExclusive: boolean;
+			usePKCE: boolean;
 			client: oidc.Configuration;
 			apiKey: string;
 			scope: string;
@@ -96,6 +97,7 @@ export async function createOidcConnector(
 	return {
 		isValid: true,
 		isExclusive: config.disable_api_key_login,
+		usePKCE: config.use_pkce,
 		client: oidcClientOrErrors,
 		apiKey: config.headscale_api_key,
 		scope: config.scope,
@@ -120,6 +122,13 @@ async function discoveryCoalesce(
 			config.client_id,
 		);
 		metadata = client.serverMetadata();
+		if (config.use_pkce === true && !client.serverMetadata().supportsPKCE()) {
+			log.warn(
+				'config',
+				'OIDC provider does not support PKCE, but it is enabled in the config',
+			);
+		}
+
 		if (metadata.claims_supported != null) {
 			if (!metadata.claims_supported.includes('sub')) {
 				log.error('config', 'OIDC provider does not support `sub` claim');

app/utils/oidc-state.ts

@@ -4,6 +4,7 @@ import type { HeadplaneConfig } from '~/server/config/config-schema';
 export interface OidcStateCookie {
 	nonce: string;
 	state: string;
+	verifier: string;
 	redirect_uri: string;
 }
 
@@ -31,6 +32,7 @@ export function createOidcStateCookie(config: HeadplaneConfig) {
 				typeof parsed !== 'object' ||
 				typeof parsed.nonce !== 'string' ||
 				typeof parsed.state !== 'string' ||
+				typeof parsed.verifier !== 'string' ||
 				typeof parsed.redirect_uri !== 'string'
 			) {
 				return null;
@@ -39,6 +41,7 @@ export function createOidcStateCookie(config: HeadplaneConfig) {
 			return {
 				nonce: parsed.nonce,
 				state: parsed.state,
+				verifier: parsed.verifier,
 				redirect_uri: parsed.redirect_uri,
 			};
 		},