headplane/app/routes/auth/oidc-callback.ts

import { createHash } from 'node:crypto';
import { count, eq } from 'drizzle-orm';
import { createCookie, type LoaderFunctionArgs, redirect } from 'react-router';
import { ulid } from 'ulidx';
import type { LoadContext } from '~/server';
import { HeadplaneConfig } from '~/server/config/schema';
import { users } from '~/server/db/schema';
import { mapOidcGroupsToRole, 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;
}

export async function loader({
	request,
	context,
}: LoaderFunctionArgs<LoadContext>) {
	if (!context.oidc) {
		throw new Error('OIDC is not enabled');
	}

	// Check if we have 0 query parameters
	const url = new URL(request.url);
	if (url.searchParams.toString().length === 0) {
		return redirect('/login');
	}

	const cookie = createCookie('__oidc_auth_flow', {
		httpOnly: true,
		maxAge: 300, // 5 minutes
	});

	const data: OidcFlowSession | null = await cookie.parse(
		request.headers.get('Cookie'),
	);

	if (data === null) {
		console.warn('OIDC flow session not found');
		return redirect('/login');
	}

	const { code_verifier, state, nonce, redirect_uri } = data;
	if (!code_verifier || !state || !nonce || !redirect_uri) {
		return send({ error: 'Missing OIDC state' }, { status: 400 });
	}

	// 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 [{ count: userCount }] = await context.db
			.select({ count: count() })
			.from(users)
			.where(eq(users.caps, Roles.owner));

		// Determine role from OIDC groups with smart mapping, but ensure first user becomes owner
		const mappedRole = mapOidcGroupsToRole(user.groups);
		const finalRole = userCount === 0 ? 'owner' : mappedRole;
		const capabilities = Roles[finalRole];

		// Log helpful information for debugging
		console.log(
			`✓ OIDC Authentication successful for ${user.email || user.subject}`,
		);
		console.log(
			`  Groups found: ${user.groups.length > 0 ? user.groups.join(', ') : 'none'}`,
		);
		console.log(
			`  Assigned role: ${finalRole}${userCount === 0 ? ' (first user = owner)' : ''}`,
		);

		if (user.groups.length === 0) {
			console.log(`ℹ️  No groups found for user. Using default 'member' role.`);
			console.log(
				`   To assign admin role, try: HEADPLANE_ADMIN_GROUPS="${user.email?.split('@')[0] || 'admin'}"`,
			);
		}

		// Insert or update user with groups and role-based capabilities
		await context.db
			.insert(users)
			.values({
				id: ulid(),
				sub: user.subject,
				caps: capabilities,
				groups: user.groups,
			})
			.onConflictDoUpdate({
				target: users.sub,
				set: {
					caps: capabilities,
					groups: user.groups,
				},
			});

		return redirect('/machines', {
			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,
				}),
			},
		});
	} catch (error) {
		// Enhanced error logging with helpful guidance
		console.error('❌ OIDC Authentication failed:', error);

		// Provide helpful troubleshooting suggestions
		if (error instanceof Error) {
			if (error.message.includes('invalid_grant')) {
				console.error('💡 Common fixes for invalid_grant:');
				console.error('   - Check client_secret is correct');
				console.error(
					'   - Verify redirect_uri matches exactly in identity provider',
				);
				console.error('   - Ensure system time is synchronized');
			} else if (error.message.includes('invalid_scope')) {
				console.error('💡 Scope issue detected:');
				console.error(
					'   - Remove "groups" from scope if provider doesn\'t support it',
				);
				console.error(
					'   - Try environment variables: HEADPLANE_ADMIN_GROUPS="admin,managers"',
				);
			} else if (error.message.includes('userinfo')) {
				console.error('💡 UserInfo endpoint issue:');
				console.error('   - Provider may not support userinfo endpoint');
				console.error('   - Check if user has required permissions');
			}
		}

		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;
		}

		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`;
	}

	return undefined;
}