feat: add pkce-like session binding and remove nonce impl

apps/web/src/app/(public)/sign-in/_components/wizard/context.tsx

@@ -1,14 +1,22 @@
 'use client'
 
 import type { Dispatch, PropsWithChildren, SetStateAction } from 'react'
-import { createContext, useContext, useState } from 'react'
+import { createContext, useContext, useRef, useState } from 'react'
 import { useInterval } from 'usehooks-ts'
 
+import {
+  browserCreatePkceChallenge,
+  browserCreatePkceVerifier,
+} from '~/lib/pkce/browser-pkce'
+
 interface SignInState {
   timer: number
   resetTimer: () => void
   vfnStepData: VfnStepData | undefined
   setVfnStepData: Dispatch<SetStateAction<VfnStepData | undefined>>
+  getVerifier: (challenge: string) => string | undefined
+  newChallenge: () => Promise<string>
+  clearVerifierMap: () => void
 }
 
 export const SignInWizardContext = createContext<SignInState | undefined>(
@@ -38,6 +46,7 @@ interface SignInWizardProviderProps {
 export interface VfnStepData {
   email: string
   otpPrefix: string
+  codeChallenge: string
 }
 
 export const SignInWizardProvider = ({
@@ -47,6 +56,20 @@ export const SignInWizardProvider = ({
   const [vfnStepData, setVfnStepData] = useState<VfnStepData>()
   const [timer, setTimer] = useState(delayForResendSeconds)
 
+  const challengeToVerifierMap = useRef(new Map<string, string>())
+  const newChallenge = async () => {
+    const verifier = browserCreatePkceVerifier()
+    const challenge = await browserCreatePkceChallenge(verifier)
+    challengeToVerifierMap.current.set(challenge, verifier)
+    return challenge
+  }
+  const getVerifier = (challenge: string) => {
+    return challengeToVerifierMap.current.get(challenge)
+  }
+  const clearVerifierMap = () => {
+    challengeToVerifierMap.current.clear()
+  }
+
   const resetTimer = () => setTimer(delayForResendSeconds)
 
   // Start the resend timer once in the vfn step.
@@ -63,6 +86,9 @@ export const SignInWizardProvider = ({
         setVfnStepData,
         timer,
         resetTimer,
+        newChallenge,
+        getVerifier,
+        clearVerifierMap,
       }}
     >
       {children}

apps/web/src/app/(public)/sign-in/_components/wizard/email/email-flow.tsx

@@ -8,8 +8,12 @@ import { VerificationStep } from './verification-step'
 export const EmailFlow = () => {
   const { setVfnStepData, vfnStepData } = useSignInWizard()
 
-  const handleOnSuccessEmail = ({ email, otpPrefix }: VfnStepData) => {
-    setVfnStepData({ email, otpPrefix })
+  const handleOnSuccessEmail = ({
+    email,
+    otpPrefix,
+    codeChallenge,
+  }: VfnStepData) => {
+    setVfnStepData({ email, otpPrefix, codeChallenge })
   }
 
   if (vfnStepData?.email) {

apps/web/src/app/(public)/sign-in/_components/wizard/email/email-step.tsx

@@ -1,4 +1,4 @@
-import { useEffect } from 'react'
+import { useEffect, useState } from 'react'
 import { zodResolver } from '@hookform/resolvers/zod'
 import { Button } from '@opengovsg/oui/button'
 import { useMutation } from '@tanstack/react-query'
@@ -10,13 +10,18 @@ import { TextField } from '@acme/ui/text-field'
 import type { VfnStepData } from '../context'
 import { useTRPC } from '~/trpc/react'
 import { emailSignInSchema } from '~/validators/auth'
+import { useSignInWizard } from '../context'
 
 interface EmailStepProps {
-  onNext: ({ email, otpPrefix }: VfnStepData) => void
+  onNext: ({ email, otpPrefix, codeChallenge }: VfnStepData) => void
 }
+
 export const EmailStep = ({ onNext }: EmailStepProps) => {
+  const { newChallenge } = useSignInWizard()
+  const [newChallengePending, setNewChallengePending] = useState(false)
+
   const { handleSubmit, setError, control } = useForm({
-    resolver: zodResolver(emailSignInSchema),
+    resolver: zodResolver(emailSignInSchema.omit({ codeChallenge: true })),
     defaultValues: {
       email: '',
     },
@@ -34,21 +39,34 @@ export const EmailStep = ({ onNext }: EmailStepProps) => {
 
   const loginMutation = useMutation(
     trpc.auth.email.login.mutationOptions({
-      onSuccess: onNext,
-      onError: (error) => setError('email', { message: error.message }),
-      trpc: {
-        context: {
-          // Need to set session data for nonce
-          skipStreaming: true,
-        },
+      onSuccess: (res, req) => {
+        return onNext({
+          email: res.email,
+          otpPrefix: res.otpPrefix,
+          codeChallenge: req.codeChallenge,
+        })
       },
+      onError: (error) => setError('email', { message: error.message }),
     }),
   )
 
+  const isPending = loginMutation.isPending || newChallengePending
+
   return (
     <form
       noValidate
-      onSubmit={handleSubmit(({ email }) => loginMutation.mutate({ email }))}
+      onSubmit={handleSubmit(({ email }) => {
+        if (isPending) return
+        setNewChallengePending(true)
+        newChallenge()
+          .then((codeChallenge) => {
+            loginMutation.mutate({ email, codeChallenge })
+          })
+          .catch(console.error)
+          .finally(() => {
+            setNewChallengePending(false)
+          })
+      })}
       className="flex flex-1 flex-col gap-4"
     >
       <Controller
@@ -68,7 +86,7 @@ export const EmailStep = ({ onNext }: EmailStepProps) => {
           />
         )}
       />
-      <Button size="sm" isPending={loginMutation.isPending} type="submit">
+      <Button size="sm" isPending={isPending} type="submit">
         Get OTP
       </Button>
     </form>

apps/web/src/app/(public)/sign-in/_components/wizard/email/verification-step.tsx

@@ -20,7 +20,17 @@ export const VerificationStep = () => {
   const [showOtpDelayMessage, setShowOtpDelayMessage] = useState(false)
   const trpc = useTRPC()
 
-  const { vfnStepData, timer, setVfnStepData, resetTimer } = useSignInWizard()
+  const {
+    vfnStepData,
+    timer,
+    setVfnStepData,
+    resetTimer,
+    newChallenge,
+    getVerifier,
+    clearVerifierMap,
+  } = useSignInWizard()
+  const [newChallengePending, setNewChallengePending] = useState(false)
+  const codeVerifier = getVerifier(vfnStepData?.codeChallenge ?? '') ?? ''
 
   useInterval(
     () => setShowOtpDelayMessage(true),
@@ -29,7 +39,7 @@ export const VerificationStep = () => {
   )
 
   const { control, handleSubmit, resetField, setFocus, setError } = useForm({
-    resolver: zodResolver(emailVerifyOtpSchema),
+    resolver: zodResolver(emailVerifyOtpSchema.omit({ codeVerifier: true })),
     defaultValues: {
       email: vfnStepData?.email ?? '',
       token: '',
@@ -39,6 +49,7 @@ export const VerificationStep = () => {
   const verifyOtpMutation = useMutation(
     trpc.auth.email.verifyOtp.mutationOptions({
       onSuccess: () => {
+        clearVerifierMap()
         router.refresh()
       },
       onError: (error) => {
@@ -55,38 +66,44 @@ export const VerificationStep = () => {
 
   const resendOtpMutation = useMutation(
     trpc.auth.email.login.mutationOptions({
-      onError: (error) => setError('token', { message: error.message }),
-      trpc: {
-        context: {
-          // Need to set session data for nonce
-          skipStreaming: true,
-        },
+      onSuccess: (res, req) => {
+        setVfnStepData({
+          email: res.email,
+          otpPrefix: res.otpPrefix,
+          codeChallenge: req.codeChallenge,
+        })
+        resetField('token')
+        setFocus('token')
+        // On success, restart the timer before this can be called again.
+        resetTimer()
       },
+      onError: (error) => setError('token', { message: error.message }),
     }),
   )
 
+  const isResendPending = resendOtpMutation.isPending || newChallengePending
   const handleResendOtp = () => {
     if (timer > 0 || !vfnStepData?.email) return
-    return resendOtpMutation.mutate(
-      { email: vfnStepData.email },
-      {
-        onSuccess: ({ email, otpPrefix }) => {
-          setVfnStepData({ email, otpPrefix })
-          resetField('token')
-          setFocus('token')
-          // On success, restart the timer before this can be called again.
-          resetTimer()
-        },
-      },
-    )
+    if (isResendPending) return
+    setNewChallengePending(true)
+    newChallenge()
+      .then((codeChallenge) => {
+        resendOtpMutation.mutate({ email: vfnStepData.email, codeChallenge })
+      })
+      .catch(console.error)
+      .finally(() => {
+        setNewChallengePending(false)
+      })
   }
 
   if (!vfnStepData) return null
 
   return (
     <form
       noValidate
-      onSubmit={handleSubmit((values) => verifyOtpMutation.mutate(values))}
+      onSubmit={handleSubmit(({ email, token }) =>
+        verifyOtpMutation.mutate({ email, token, codeVerifier }),
+      )}
       className="flex flex-1 flex-col gap-4"
     >
       <Controller
@@ -137,7 +154,7 @@ export const VerificationStep = () => {
           )}
           size="xs"
           onPress={handleResendOtp}
-          isPending={resendOtpMutation.isPending}
+          isPending={isResendPending}
           // isPending
           isDisabled={timer > 0}
           spinner={

apps/web/src/lib/pkce/browser-pkce.ts

@@ -0,0 +1,24 @@
+import { pkceVerifierGenerator } from '~/lib/pkce/pkce'
+
+// This file hosts functions pertaining to PKCE (Proof Key for Code Exchange) as per RFC 7636
+// If you need to understand PKCE, read https://datatracker.ietf.org/doc/html/rfc7636
+// Do not use this for actual OAuth flows, please use a tested library instead.
+
+// We need to split up server and browser implementations due to using crypto APIs
+
+export function browserCreatePkceVerifier(): string {
+  return pkceVerifierGenerator()
+}
+
+export async function browserCreatePkceChallenge(
+  codeVerifier: string,
+): Promise<string> {
+  const encoder = new TextEncoder()
+  const data = encoder.encode(codeVerifier)
+  const hashBuffer = await window.crypto.subtle.digest('SHA-256', data)
+
+  const hashArray = Array.from(new Uint8Array(hashBuffer))
+  const base64 = btoa(String.fromCharCode(...hashArray))
+
+  return base64.replace(/=/g, '').replace(/\+/g, '-').replace(/\//g, '_')
+}

apps/web/src/lib/pkce/pkce.ts

@@ -0,0 +1,8 @@
+import { customAlphabet } from 'nanoid'
+
+import { PKCE_LENGTH, PKCE_VERIFIER_ALPHABET } from '~/validators/auth'
+
+export const pkceVerifierGenerator = customAlphabet(
+  PKCE_VERIFIER_ALPHABET,
+  PKCE_LENGTH,
+)

apps/web/src/lib/pkce/server-pkce.ts

@@ -0,0 +1,17 @@
+import { createHash } from 'crypto'
+
+import { pkceVerifierGenerator } from '~/lib/pkce/pkce'
+
+// This file hosts functions pertaining to PKCE (Proof Key for Code Exchange) as per RFC 7636
+// If you need to understand PKCE, read https://datatracker.ietf.org/doc/html/rfc7636
+// Do not use this for actual OAuth flows, please use a tested library instead.
+
+// We need to split up server and browser implementations due to using crypto APIs
+
+export function ssCreatePkceVerifier(): string {
+  return pkceVerifierGenerator()
+}
+
+export function ssCreatePkceChallenge(codeVerifier: string): string {
+  return createHash('sha256').update(codeVerifier).digest('base64url')
+}

apps/web/src/server/api/routers/auth/auth.email.router.ts

@@ -1,5 +1,3 @@
-import { randomUUID } from 'node:crypto'
-import { TRPCError } from '@trpc/server'
 import z from 'zod'
 
 import { emailLogin, emailVerifyOtp } from '~/server/modules/auth/auth.service'
@@ -21,36 +19,23 @@ export const emailAuthRouter = createTRPCRouter({
         otpPrefix: z.string().length(OTP_PREFIX_LENGTH),
       }),
     )
-    .mutation(async ({ input, ctx }) => {
-      const nonce = randomUUID()
-      const { email, otpPrefix } = await emailLogin({
-        email: input.email,
-        nonce,
+    .mutation(async ({ input: { email, codeChallenge } }) => {
+      // returnedEmail may differ from input email
+      const { email: returnedEmail, otpPrefix } = await emailLogin({
+        email,
+        codeChallenge,
       })
 
-      ctx.session.nonce = nonce
-      await ctx.session.save()
       return {
-        email,
+        email: returnedEmail,
         otpPrefix,
       }
     }),
   verifyOtp: publicProcedure
     .input(emailVerifyOtpSchema)
-    .mutation(async ({ input: { email, token }, ctx }) => {
-      const nonce = ctx.session.nonce
-      // Ensure nonce exists in session
-      if (!nonce) {
-        throw new TRPCError({
-          code: 'BAD_REQUEST',
-          message:
-            'Something went wrong. Please request for a new OTP before retrying.',
-        })
-      }
-
-      await emailVerifyOtp({ email, token, nonce })
+    .mutation(async ({ input: { email, token, codeVerifier }, ctx }) => {
+      await emailVerifyOtp({ email, token, codeVerifier })
       const user = await upsertUserAndAccountByEmail(email)
-      ctx.session.nonce = undefined
       ctx.session.userId = user.id
       await ctx.session.save()
       return user