feat: split browser and serverside pkce generators

Author: zhongliang02

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

@@ -1,21 +1,21 @@
 'use client'
 
 import type { Dispatch, PropsWithChildren, SetStateAction } from 'react'
-import { createContext, useCallback, useContext, useRef, useState } from 'react'
+import { createContext, useContext, useRef, useState } from 'react'
 import { useInterval } from 'usehooks-ts'
 
 import {
-  createPkceChallenge,
-  createPkceVerifier,
-} from '~/server/modules/auth/auth.pkce'
+  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: () => string
+  newChallenge: () => Promise<string>
   clearVerifierMap: () => void
 }
 
@@ -57,18 +57,18 @@ export const SignInWizardProvider = ({
   const [timer, setTimer] = useState(delayForResendSeconds)
 
   const challengeToVerifierMap = useRef(new Map<string, string>())
-  const newChallenge = useCallback(() => {
-    const verifier = createPkceVerifier()
-    const challenge = createPkceChallenge(verifier)
+  const newChallenge = async () => {
+    const verifier = browserCreatePkceVerifier()
+    const challenge = await browserCreatePkceChallenge(verifier)
     challengeToVerifierMap.current.set(challenge, verifier)
     return challenge
-  }, []) // stable reference no deps
-  const getVerifier = useCallback((challenge: string) => {
+  }
+  const getVerifier = (challenge: string) => {
     return challengeToVerifierMap.current.get(challenge)
-  }, []) // stable reference no deps
-  const clearVerifierMap = useCallback(() => {
+  }
+  const clearVerifierMap = () => {
     challengeToVerifierMap.current.clear()
-  }, []) // stable reference no deps
+  }
 
   const resetTimer = () => setTimer(delayForResendSeconds)
 

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'
@@ -15,8 +15,11 @@ import { useSignInWizard } from '../context'
 interface EmailStepProps {
   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.omit({ codeChallenge: true })),
     defaultValues: {
@@ -47,12 +50,23 @@ export const EmailStep = ({ onNext }: EmailStepProps) => {
     }),
   )
 
+  const isPending = loginMutation.isPending || newChallengePending
+
   return (
     <form
       noValidate
-      onSubmit={handleSubmit(({ email }) =>
-        loginMutation.mutate({ email, codeChallenge: newChallenge() }),
-      )}
+      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
@@ -72,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

@@ -29,6 +29,7 @@ export const VerificationStep = () => {
     getVerifier,
     clearVerifierMap,
   } = useSignInWizard()
+  const [newChallengePending, setNewChallengePending] = useState(false)
   const codeVerifier = getVerifier(vfnStepData?.codeChallenge ?? '') ?? ''
 
   useInterval(
@@ -65,28 +66,34 @@ export const VerificationStep = () => {
 
   const resendOtpMutation = useMutation(
     trpc.auth.email.login.mutationOptions({
+      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, codeChallenge: newChallenge() },
-      {
-        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()
-        },
-      },
-    )
+    if (isResendPending) return
+    setNewChallengePending(true)
+    newChallenge()
+      .then((codeChallenge) => {
+        resendOtpMutation.mutate({ email: vfnStepData.email, codeChallenge })
+      })
+      .catch(console.error)
+      .finally(() => {
+        setNewChallengePending(false)
+      })
   }
 
   if (!vfnStepData) return null
@@ -147,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/modules/auth/__tests__/auth.service.spec.ts

@@ -6,8 +6,11 @@ import { mock } from 'vitest-mock-extended'
 
 import { db } from '@acme/db'
 
+import {
+  ssCreatePkceChallenge,
+  ssCreatePkceVerifier,
+} from '~/lib/pkce/server-pkce'
 import * as mailService from '../../mail/mail.service'
-import { createPkceChallenge, createPkceVerifier } from '../auth.pkce'
 import { emailLogin, emailVerifyOtp } from '../auth.service'
 import { createAuthToken, createVfnIdentifier } from '../auth.utils'
 
@@ -88,8 +91,8 @@ describe('auth.service', () => {
   describe('emailVerifyOtp', () => {
     it('should successfully verify a valid OTP', async () => {
       const email = '[email protected]'
-      const codeVerifier = createPkceVerifier()
-      const codeChallenge = createPkceChallenge(codeVerifier)
+      const codeVerifier = ssCreatePkceVerifier()
+      const codeChallenge = ssCreatePkceChallenge(codeVerifier)
 
       // Create a verification token
       const { token } = await emailLogin({ email, codeChallenge })
@@ -109,9 +112,9 @@ describe('auth.service', () => {
 
     it('should reject a correct OTP with wrong codeVerifier', async () => {
       const email = '[email protected]'
-      const correctVerifier = createPkceVerifier()
-      const wrongVerifier = createPkceVerifier()
-      const codeChallenge = createPkceChallenge(correctVerifier)
+      const correctVerifier = ssCreatePkceVerifier()
+      const wrongVerifier = ssCreatePkceVerifier()
+      const codeChallenge = ssCreatePkceChallenge(correctVerifier)
 
       // Create a verification token
       const { token } = await emailLogin({ email, codeChallenge })
@@ -123,7 +126,7 @@ describe('auth.service', () => {
     })
     it('should throw error for non-existent codeChallenge', async () => {
       const email = '[email protected]'
-      const codeVerifier = createPkceVerifier()
+      const codeVerifier = ssCreatePkceVerifier()
       const token = '123456'
 
       await expect(
@@ -133,10 +136,10 @@ describe('auth.service', () => {
 
     it('should reject a wrong OTP with wrong codeVerifier', async () => {
       const email = '[email protected]'
-      const correctVerifier = createPkceVerifier()
-      const correctCodeChallenge = createPkceChallenge(correctVerifier)
+      const correctVerifier = ssCreatePkceVerifier()
+      const correctCodeChallenge = ssCreatePkceChallenge(correctVerifier)
 
-      const wrongVerifier = createPkceVerifier()
+      const wrongVerifier = ssCreatePkceVerifier()
 
       // Create a verification token
       await emailLogin({ email, codeChallenge: correctCodeChallenge })
@@ -155,8 +158,8 @@ describe('auth.service', () => {
 
     it('should reject a wrong OTP with correct codeVerifier', async () => {
       const email = '[email protected]'
-      const codeVerifier = createPkceVerifier()
-      const codeChallenge = createPkceChallenge(codeVerifier)
+      const codeVerifier = ssCreatePkceVerifier()
+      const codeChallenge = ssCreatePkceChallenge(codeVerifier)
       const wrongToken = 'WRONG6'
 
       await emailLogin({ email, codeChallenge: codeChallenge })
@@ -167,8 +170,8 @@ describe('auth.service', () => {
 
     it('should reject an expired OTP with correct codeVerifier', async () => {
       const email = '[email protected]'
-      const codeVerifier = createPkceVerifier()
-      const codeChallenge = createPkceChallenge(codeVerifier)
+      const codeVerifier = ssCreatePkceVerifier()
+      const codeChallenge = ssCreatePkceChallenge(codeVerifier)
 
       const { token, hashedToken } = createAuthToken({
         email,
@@ -196,8 +199,8 @@ describe('auth.service', () => {
 
     it('should increment attempts on each verification try', async () => {
       const email = '[email protected]'
-      const codeVerifier = createPkceVerifier()
-      const codeChallenge = createPkceChallenge(codeVerifier)
+      const codeVerifier = ssCreatePkceVerifier()
+      const codeChallenge = ssCreatePkceChallenge(codeVerifier)
       const wrongToken = 'WRONG6'
 
       await emailLogin({ email, codeChallenge: codeChallenge })
@@ -216,8 +219,8 @@ describe('auth.service', () => {
 
     it('should reject after too many failed attempts (>5)', async () => {
       const email = '[email protected]'
-      const codeVerifier = createPkceVerifier()
-      const codeChallenge = createPkceChallenge(codeVerifier)
+      const codeVerifier = ssCreatePkceVerifier()
+      const codeChallenge = ssCreatePkceChallenge(codeVerifier)
       const token = 'WRONG6'
 
       await emailLogin({ email, codeChallenge: codeChallenge })
@@ -237,8 +240,8 @@ describe('auth.service', () => {
 
     it('should delete verification token after successful verification', async () => {
       const email = '[email protected]'
-      const codeVerifier = createPkceVerifier()
-      const codeChallenge = createPkceChallenge(codeVerifier)
+      const codeVerifier = ssCreatePkceVerifier()
+      const codeChallenge = ssCreatePkceChallenge(codeVerifier)
       const { token } = await emailLogin({ email, codeChallenge })
 
       await emailVerifyOtp({ email, token, codeVerifier })
@@ -256,8 +259,8 @@ describe('auth.service', () => {
 
     it('should prevent token reuse after successful verification', async () => {
       const email = '[email protected]'
-      const codeVerifier = createPkceVerifier()
-      const codeChallenge = createPkceChallenge(codeVerifier)
+      const codeVerifier = ssCreatePkceVerifier()
+      const codeChallenge = ssCreatePkceChallenge(codeVerifier)
       const { token } = await emailLogin({ email, codeChallenge })
 
       // First verification succeeds

apps/web/src/server/modules/auth/__tests__/auth.utils.spec.ts

@@ -82,7 +82,7 @@ describe('auth.utils', () => {
     const testCodeChallenge = 'test-codeChallenge-123'
 
     it('should generate tokens of sufficient entropy', () => {
-      const N = 10000
+      const N = 100
       const tokens = Array.from({ length: N }).map(
         () =>
           createAuthToken({