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

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

@@ -1,14 +1,19 @@
 'use client'
 
-import type { Dispatch, PropsWithChildren, SetStateAction } from 'react'
+import type { Dispatch, PropsWithChildren, SetStateAction } from 'react';
+import { useCallback, useRef} from 'react'
 import { createContext, useContext, useState } from 'react'
 import { useInterval } from 'usehooks-ts'
+import { createPkceChallenge, createPkceVerifier } from "~/server/modules/auth/auth.pkce";
 
 interface SignInState {
   timer: number
   resetTimer: () => void
   vfnStepData: VfnStepData | undefined
   setVfnStepData: Dispatch<SetStateAction<VfnStepData | undefined>>
+  getVerifier: (challenge: string) => string | undefined
+  newChallenge: () => string
+  clearVerifierMap: () => void
 }
 
 export const SignInWizardContext = createContext<SignInState | undefined>(
@@ -38,6 +43,7 @@ interface SignInWizardProviderProps {
 export interface VfnStepData {
   email: string
   otpPrefix: string
+  codeChallenge: string
 }
 
 export const SignInWizardProvider = ({
@@ -47,6 +53,20 @@ export const SignInWizardProvider = ({
   const [vfnStepData, setVfnStepData] = useState<VfnStepData>()
   const [timer, setTimer] = useState(delayForResendSeconds)
 
+  const challengeToVerifierMap = useRef(new Map<string, string>())
+  const newChallenge = useCallback(()=>{
+    const verifier = createPkceVerifier()
+    const challenge = createPkceChallenge(verifier)
+    challengeToVerifierMap.current.set(challenge, verifier)
+    return challenge
+  }, []) // stable reference no deps
+  const getVerifier = useCallback((challenge: string)=>{
+    return challengeToVerifierMap.current.get(challenge)
+  }, []) // stable reference no deps
+  const clearVerifierMap = useCallback(()=>{
+    challengeToVerifierMap.current.clear()
+  }, []) // stable reference no deps
+
   const resetTimer = () => setTimer(delayForResendSeconds)
 
   // Start the resend timer once in the vfn step.
@@ -63,6 +83,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,8 @@ 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

@@ -7,19 +7,21 @@ import { Controller, useForm } from 'react-hook-form'
 
 import { TextField } from '@acme/ui/text-field'
 
-import type { VfnStepData } from '../context'
+import type { VfnStepData } from '../context';
+import { useSignInWizard } from '../context'
 import { useTRPC } from '~/trpc/react'
 import { emailSignInSchema } from '~/validators/auth'
 
 interface EmailStepProps {
-  onNext: ({ email, otpPrefix }: VfnStepData) => void
+  onNext: ({ email, otpPrefix, codeChallenge }: VfnStepData) => void
 }
 export const EmailStep = ({ onNext }: EmailStepProps) => {
+  const { newChallenge } = useSignInWizard()
   const { handleSubmit, setError, control } = useForm({
-    resolver: zodResolver(emailSignInSchema),
+    resolver: zodResolver(emailSignInSchema.omit({codeChallenge: true})),
     defaultValues: {
       email: '',
-    },
+    }
   })
 
   const [queryError] = useQueryState('error', { defaultValue: '' })
@@ -34,21 +36,21 @@ 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 }),
     }),
   )
 
   return (
     <form
       noValidate
-      onSubmit={handleSubmit(({ email }) => loginMutation.mutate({ email }))}
+      onSubmit={handleSubmit(({ email }) => loginMutation.mutate({ email, codeChallenge: newChallenge() }))}
       className="flex flex-1 flex-col gap-4"
     >
       <Controller

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

@@ -20,7 +20,8 @@ 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 codeVerifier = getVerifier(vfnStepData?.codeChallenge ?? '') ?? ''
 
   useInterval(
     () => setShowOtpDelayMessage(true),
@@ -29,7 +30,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 +40,7 @@ export const VerificationStep = () => {
   const verifyOtpMutation = useMutation(
     trpc.auth.email.verifyOtp.mutationOptions({
       onSuccess: () => {
+        clearVerifierMap()
         router.refresh()
       },
       onError: (error) => {
@@ -56,22 +58,20 @@ 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,
-        },
-      },
     }),
   )
 
   const handleResendOtp = () => {
     if (timer > 0 || !vfnStepData?.email) return
     return resendOtpMutation.mutate(
-      { email: vfnStepData.email },
+      { email: vfnStepData.email, codeChallenge: newChallenge() },
       {
-        onSuccess: ({ email, otpPrefix }) => {
-          setVfnStepData({ email, otpPrefix })
+        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.
@@ -86,7 +86,7 @@ export const VerificationStep = () => {
   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

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