fix 2fa workflow

Author: kentcdodds

app/routes/_auth+/login.tsx

@@ -9,7 +9,7 @@ import { Spacer } from '~/components/spacer.tsx'
 import { authenticator, requireAnonymous } from '~/utils/auth.server.ts'
 import { commitSession, getSession } from '~/utils/session.server.ts'
 import { InlineLogin } from '../resources+/login.tsx'
-import { unverifiedSessionKey } from '../resources+/verify.tsx'
+import { Verifier, unverifiedSessionKey } from '../resources+/verify.tsx'
 
 export async function loader({ request }: DataFunctionArgs) {
 	await requireAnonymous(request)
@@ -49,11 +49,11 @@ export default function LoginPage() {
 					</p>
 				</div>
 				<Spacer size="xs" />
-				<InlineLogin
-					redirectTo={redirectTo}
-					formError={data.formError}
-					status={data.unverified ? 'unverified' : 'idle'}
-				/>
+				{data.unverified ? (
+					<Verifier redirectTo={redirectTo} />
+				) : (
+					<InlineLogin redirectTo={redirectTo} formError={data.formError} />
+				)}
 			</div>
 		</div>
 	)

app/routes/resources+/login.tsx

@@ -14,7 +14,7 @@ import { commitSession, getSession } from '~/utils/session.server.ts'
 import { passwordSchema, usernameSchema } from '~/utils/user-validation.ts'
 import { checkboxSchema } from '~/utils/zod-extensions.ts'
 import { twoFAVerificationType } from '../settings+/profile.two-factor.tsx'
-import { Verifier, unverifiedSessionKey } from './verify.tsx'
+import { unverifiedSessionKey } from './verify.tsx'
 
 const ROUTE_PATH = '/resources/login'
 
@@ -93,29 +93,25 @@ export async function action({ request }: DataFunctionArgs) {
 			}),
 		},
 	}
-	if (user2FA) {
-		return json({ status: 'unverified', submission } as const, responseInit)
-	} else {
-		if (redirectTo) {
-			throw redirect(safeRedirect(redirectTo), responseInit)
-		}
+	if (user2FA || !redirectTo) {
 		return json({ status: 'success', submission } as const, responseInit)
+	} else {
+		throw redirect(safeRedirect(redirectTo), responseInit)
 	}
 }
 
 export function InlineLogin({
 	redirectTo,
 	formError,
-	status = 'idle',
 }: {
-	status?: 'idle' | 'error' | 'unverified'
 	redirectTo?: string
 	formError?: string | null
 }) {
 	const loginFetcher = useFetcher<typeof action>()
 
 	const [form, fields] = useForm({
 		id: 'inline-login',
+		defaultValue: { redirectTo },
 		constraint: getFieldsetConstraint(loginFormSchema),
 		lastSubmission: loginFetcher.data?.submission,
 		onValidate({ formData }) {
@@ -124,12 +120,6 @@ export function InlineLogin({
 		shouldRevalidate: 'onBlur',
 	})
 
-	const fetcherStatus = loginFetcher.data?.status ?? status
-
-	if (fetcherStatus === 'unverified') {
-		return <Verifier redirectTo={redirectTo} />
-	}
-
 	return (
 		<div>
 			<div className="mx-auto w-full max-w-md px-8">
@@ -144,7 +134,7 @@ export function InlineLogin({
 							htmlFor: fields.username.id,
 							children: 'Username',
 						}}
-						inputProps={conform.input(fields.username)}
+						inputProps={{ ...conform.input(fields.username), autoFocus: true }}
 						errors={fields.username.errors}
 					/>
 
@@ -177,11 +167,7 @@ export function InlineLogin({
 						</div>
 					</div>
 
-					<input
-						value={redirectTo}
-						{...conform.input(fields.redirectTo)}
-						type="hidden"
-					/>
+					<input {...conform.input(fields.redirectTo)} type="hidden" />
 					<ErrorList errors={[...form.errors, formError]} id={form.errorId} />
 
 					<div className="flex items-center justify-between gap-6 pt-3">
@@ -190,7 +176,9 @@ export function InlineLogin({
 							size="md"
 							variant="primary"
 							status={
-								loginFetcher.state === 'submitting' ? 'pending' : fetcherStatus
+								loginFetcher.state === 'submitting'
+									? 'pending'
+									: loginFetcher.data?.status
 							}
 							type="submit"
 							disabled={loginFetcher.state !== 'idle'}

app/routes/resources+/verify.tsx

@@ -2,26 +2,29 @@ import { conform, useForm } from '@conform-to/react'
 import { getFieldsetConstraint, parse } from '@conform-to/zod'
 import { json, redirect, type DataFunctionArgs } from '@remix-run/node'
 import { useFetcher } from '@remix-run/react'
+import invariant from 'tiny-invariant'
 import { z } from 'zod'
 import { twoFAVerificationType } from '~/routes/settings+/profile.two-factor.tsx'
 import { authenticator } from '~/utils/auth.server.ts'
 import { prisma } from '~/utils/db.server.ts'
 import { Button, ErrorList, Field } from '~/utils/forms.tsx'
 import { safeRedirect } from '~/utils/misc.ts'
-import {
-	commitSession,
-	destroySession,
-	getSession,
-} from '~/utils/session.server.ts'
+import { commitSession, getSession } from '~/utils/session.server.ts'
 import { verifyTOTP } from '~/utils/totp.server.ts'
 
 const ROUTE_PATH = '/resources/verify'
 export const unverifiedSessionKey = 'unverified-sessionId'
 
-const verifySchema = z.object({
-	code: z.string().min(6).max(6),
-	redirectTo: z.string().optional(),
-})
+const verifySchema = z.union([
+	z.object({
+		intent: z.literal('cancel'),
+	}),
+	z.object({
+		intent: z.literal('confirm'),
+		code: z.string().min(6).max(6),
+		redirectTo: z.string().optional(),
+	}),
+])
 
 export async function action({ request }: DataFunctionArgs) {
 	const form = await request.formData()
@@ -36,24 +39,28 @@ export async function action({ request }: DataFunctionArgs) {
 		where: { id: sessionId },
 		select: { userId: true },
 	})
+
 	// if there's no session for their session ID, something is *way* wrong.
 	// destory it and let them try over again... Or we'll do that if they wanna cancel.
 	if (!session || form.get('intent') === 'cancel') {
 		const redirectTo = form.get('redirectTo')
 		const params =
-			typeof redirectTo === 'string' && redirectTo
+			typeof redirectTo === 'string' && redirectTo && redirectTo !== '/'
 				? new URLSearchParams({ redirectTo })
 				: null
+		cookieSession.unset(unverifiedSessionKey)
 		throw redirect(['/login', params?.toString()].filter(Boolean).join('?'), {
 			headers: {
-				'Set-Cookie': await destroySession(cookieSession),
+				'Set-Cookie': await commitSession(cookieSession),
 			},
 		})
 	}
 
 	const submission = await parse(form, {
 		schema: () =>
 			verifySchema.superRefine(async (data, ctx) => {
+				if (data.intent === 'cancel') return
+
 				const verification = await prisma.verification.findFirst({
 					where: {
 						type: twoFAVerificationType,
@@ -105,6 +112,8 @@ export async function action({ request }: DataFunctionArgs) {
 		)
 	}
 
+	invariant(submission.value.intent === 'confirm', 'invalid intent')
+
 	const { redirectTo } = submission.value
 	cookieSession.unset(unverifiedSessionKey)
 	cookieSession.set(authenticator.sessionKey, sessionId)
@@ -143,7 +152,7 @@ export function Verifier({
 			<input type="hidden" name="redirectTo" value={redirectTo} />
 			<Field
 				labelProps={{ children: '2FA Code' }}
-				inputProps={conform.input(fields.code)}
+				inputProps={{ ...conform.input(fields.code), autoFocus: true }}
 				errors={fields.code.errors}
 			/>
 			<div className="flex flex-row-reverse justify-between">

playwright.config.ts

@@ -5,9 +5,9 @@ const PORT = process.env.PORT || '3000'
 
 export default {
 	testDir: './tests/e2e',
-	timeout: 30 * 1000,
+	timeout: 15 * 1000,
 	expect: {
-		timeout: 10 * 1000,
+		timeout: 5 * 1000,
 	},
 	fullyParallel: true,
 	forbidOnly: !!process.env.CI,

tests/e2e/2fa.test.ts

@@ -40,6 +40,7 @@ test('Users can add 2FA to their account and use it when logging in', async ({
 	await page.getByRole('menuitem', { name: /logout/i }).click()
 
 	await page.goto('/login')
+	await expect(page).toHaveURL(`/login`)
 	await page.getByRole('textbox', { name: /username/i }).fill(user.username)
 	await page.getByLabel(/^password$/i).fill(password)
 	await page.getByRole('button', { name: /log in/i }).click()