Small QoL improvements to MFA (#43861)

## Context

Resolves FE-2794

Just addressing some friction points i ran into when setting up MFA for
my account under account preferences

## Changes involved

- Hitting enter at each step of setting up MFA will submit and proceed
(Previously didn't, had to click buttons)
- When logging in with MFA, automatically submit after entering 6 digits
- Refactored to use react hook form for `AddNewFactorModal` in
`FirstStep` and `SecondStep` + simplified the logic a little

apps/studio/components/interfaces/Account/TOTPFactors/AddNewFactorModal.tsx

@@ -1,7 +1,5 @@
+import { zodResolver } from '@hookform/resolvers/zod'
 import { useQueryClient } from '@tanstack/react-query'
-import { Dispatch, SetStateAction, useEffect, useState } from 'react'
-import { toast } from 'sonner'
-
 import { LOCAL_STORAGE_KEYS } from 'common'
 import InformationBox from 'components/ui/InformationBox'
 import { organizationKeys } from 'data/organizations/keys'
@@ -9,9 +7,16 @@ import { useMfaChallengeAndVerifyMutation } from 'data/profile/mfa-challenge-and
 import { useMfaEnrollMutation } from 'data/profile/mfa-enroll-mutation'
 import { useMfaUnenrollMutation } from 'data/profile/mfa-unenroll-mutation'
 import { useLocalStorageQuery } from 'hooks/misc/useLocalStorage'
-import { Input } from 'ui'
+import { useEffect, useState } from 'react'
+import { useForm, type SubmitHandler } from 'react-hook-form'
+import { toast } from 'sonner'
+import { Form_Shadcn_, FormControl_Shadcn_, FormField_Shadcn_, Input, Input_Shadcn_ } from 'ui'
 import ConfirmationModal from 'ui-patterns/Dialogs/ConfirmationModal'
+import { FormItemLayout } from 'ui-patterns/form/FormItemLayout/FormItemLayout'
 import { GenericSkeletonLoader } from 'ui-patterns/ShimmeringLoader'
+import { z } from 'zod'
+
+type TOTP = { qr_code: string; secret: string; uri: string }
 
 interface AddNewFactorModalProps {
   visible: boolean
@@ -19,32 +24,24 @@ interface AddNewFactorModalProps {
 }
 
 export const AddNewFactorModal = ({ visible, onClose }: AddNewFactorModalProps) => {
-  // Generate a name with a number between 0 and 1000
-  const [name, setName] = useState(`App ${Math.floor(Math.random() * 1000)}`)
   const { data, mutate: enroll, isPending: isEnrolling, reset } = useMfaEnrollMutation()
 
   useEffect(() => {
-    // reset has to be called because the state is kept between if the process is canceled during
-    // the second step.
-    if (!visible) {
-      setName(`App ${Math.floor(Math.random() * 1000)}`)
-      reset()
-    }
+    if (!visible) reset()
   }, [reset, visible])
 
   return (
     <>
       <FirstStep
         visible={visible && !Boolean(data)}
-        name={name}
-        setName={setName}
-        enroll={enroll}
         isEnrolling={isEnrolling}
+        enroll={enroll}
+        reset={reset}
         onClose={onClose}
       />
       <SecondStep
         visible={visible && Boolean(data)}
-        factorName={name}
+        factorName={data?.friendly_name ?? ''}
         factor={data as Extract<typeof data, { type: 'totp' }>}
         isLoading={isEnrolling}
         onClose={onClose}
@@ -55,14 +52,33 @@ export const AddNewFactorModal = ({ visible, onClose }: AddNewFactorModalProps)
 
 interface FirstStepProps {
   visible: boolean
-  name: string
-  setName: Dispatch<SetStateAction<string>>
-  enroll: (params: { factorType: 'totp'; friendlyName?: string }) => void
   isEnrolling: boolean
+  reset: () => void
+  enroll: (params: { factorType: 'totp'; friendlyName?: string }) => void
   onClose: () => void
 }
 
-const FirstStep = ({ visible, name, enroll, setName, isEnrolling, onClose }: FirstStepProps) => {
+const FirstStep = ({ visible, isEnrolling, reset, enroll, onClose }: FirstStepProps) => {
+  const FormSchema = z.object({
+    name: z.string().min(1, 'Please provide a name to identify this app'),
+  })
+  const form = useForm<z.infer<typeof FormSchema>>({
+    resolver: zodResolver(FormSchema),
+    defaultValues: { name: '' },
+    mode: 'onChange',
+  })
+
+  const onSubmit: SubmitHandler<z.infer<typeof FormSchema>> = async (values) => {
+    enroll({ factorType: 'totp', friendlyName: values.name })
+  }
+
+  useEffect(() => {
+    if (!visible) {
+      // Generate a name with a number between 0 and 1000
+      form.reset({ name: `App ${Math.floor(Math.random() * 1000)}` })
+    }
+  }, [form, visible])
+
   return (
     <ConfirmationModal
       size="medium"
@@ -70,22 +86,34 @@ const FirstStep = ({ visible, name, enroll, setName, isEnrolling, onClose }: Fir
       title="Add a new authenticator app as a factor"
       confirmLabel="Generate QR"
       confirmLabelLoading="Generating QR"
-      disabled={name.length === 0}
       loading={isEnrolling}
       onCancel={onClose}
-      onConfirm={() => {
-        enroll({
-          factorType: 'totp',
-          friendlyName: name,
-        })
-      }}
+      onConfirm={form.handleSubmit(onSubmit)}
     >
-      <Input
-        label="Provide a name to identify this app"
-        descriptionText="A string will be randomly generated if a name is not provided"
-        value={name}
-        onChange={(e) => setName(e.target.value)}
-      />
+      <Form_Shadcn_ {...form}>
+        <form
+          id="verify-otp-form"
+          className="flex flex-col gap-4"
+          onSubmit={form.handleSubmit(onSubmit)}
+        >
+          <FormField_Shadcn_
+            key="name"
+            name="name"
+            control={form.control}
+            render={({ field }) => (
+              <FormItemLayout
+                name="name"
+                label="Provide a name to identify this app"
+                description="A string will be randomly generated if a name is not provided"
+              >
+                <FormControl_Shadcn_>
+                  <Input_Shadcn_ id="name" {...field} />
+                </FormControl_Shadcn_>
+              </FormItemLayout>
+            )}
+          />
+        </form>
+      </Form_Shadcn_>
     </ConfirmationModal>
   )
 }
@@ -96,11 +124,7 @@ interface SecondStepProps {
   factor?: {
     id: string
     type: 'totp'
-    totp: {
-      qr_code: string
-      secret: string
-      uri: string
-    }
+    totp: TOTP
   }
   isLoading: boolean
   onClose: () => void
@@ -113,14 +137,23 @@ const SecondStep = ({
   isLoading,
   onClose,
 }: SecondStepProps) => {
-  const [code, setCode] = useState('')
   const queryClient = useQueryClient()
-
   const [lastVisitedOrganization] = useLocalStorageQuery(
     LOCAL_STORAGE_KEYS.LAST_VISITED_ORGANIZATION,
     ''
   )
 
+  const FormSchema = z.object({
+    code: z.string().min(1, 'Please provide a code from your authenticator app'),
+  })
+  const form = useForm<z.infer<typeof FormSchema>>({
+    resolver: zodResolver(FormSchema),
+    defaultValues: { code: '' },
+    mode: 'onChange',
+  })
+
+  const [factor, setFactor] = useState<{ id: string; type: 'totp'; totp: TOTP } | null>(null)
+
   const { mutate: unenroll } = useMfaUnenrollMutation({ onSuccess: () => onClose() })
   const { mutate: challengeAndVerify, isPending: isVerifying } = useMfaChallengeAndVerifyMutation({
     onError: (error) => {
@@ -137,15 +170,10 @@ const SecondStep = ({
     },
   })
 
-  const [factor, setFactor] = useState<{
-    id: string
-    type: 'totp'
-    totp: {
-      qr_code: string
-      secret: string
-      uri: string
-    }
-  } | null>(null)
+  const onSubmit: SubmitHandler<z.infer<typeof FormSchema>> = async (values) => {
+    if (!factor) return toast.error('Factor required')
+    challengeAndVerify({ factorId: factor.id, code: values.code })
+  }
 
   // this useEffect is to keep the factor until a new one comes. This is a fix to an issue which
   // happens when closing the modal, the outer factor is reset to null too soon and the modal
@@ -153,6 +181,7 @@ const SecondStep = ({
   useEffect(() => {
     if (outerFactor && factor?.id !== outerFactor.id) {
       setFactor(outerFactor)
+      form.reset({ code: '' })
     }
   }, [outerFactor])
 
@@ -160,6 +189,7 @@ const SecondStep = ({
     <ConfirmationModal
       size="medium"
       visible={visible}
+      className="py-5"
       title={`Verify new factor ${factorName}`}
       confirmLabel="Confirm"
       confirmLabelLoading="Confirming"
@@ -170,50 +200,68 @@ const SecondStep = ({
         // unenrolling.
         if (factor) unenroll({ factorId: factor.id })
       }}
-      onConfirm={() => factor && challengeAndVerify({ factorId: factor.id, code })}
+      onConfirm={form.handleSubmit(onSubmit)}
     >
-      <div className="py-4 pb-0 text-sm">
-        <span>
-          Use an authenticator app to scan the following QR code, and provide the code from the app
-          to complete the enrolment.
-        </span>
-      </div>
+      <p className="text-sm">
+        Use an authenticator app to scan the following QR code, and provide the code from the app to
+        complete the enrolment.
+      </p>
+
       {isLoading && (
         <div className="pb-4 px-4">
           <GenericSkeletonLoader />
         </div>
       )}
+
       {factor && (
-        <>
+        <div className="flex flex-col gap-y-4">
           <div className="flex justify-center py-6">
             <div className="h-48 w-48 bg-white rounded">
               <img width={190} height={190} src={factor.totp.qr_code} alt={factor.totp.uri} />
             </div>
           </div>
-          <div>
-            <InformationBox
-              title="Unable to scan?"
-              description={
-                <Input
-                  copy
-                  disabled
-                  id="ref"
-                  size="small"
-                  label="You can also enter this secret key into your authenticator app"
-                  value={factor.totp.secret}
-                />
-              }
-            />
-          </div>
-          <div className="pt-2 pb-4">
-            <Input
-              label="Authentication code"
-              value={code}
-              placeholder="XXXXXX"
-              onChange={(e) => setCode(e.target.value)}
-            />
-          </div>
-        </>
+
+          <InformationBox
+            title="Unable to scan?"
+            description={
+              <Input
+                copy
+                disabled
+                id="ref"
+                size="small"
+                label="You can also enter this secret key into your authenticator app"
+                value={factor.totp.secret}
+              />
+            }
+          />
+
+          <Form_Shadcn_ {...form}>
+            <form
+              id="verify-otp-form"
+              className="flex flex-col gap-4"
+              onSubmit={form.handleSubmit(onSubmit)}
+            >
+              <FormField_Shadcn_
+                key="code"
+                name="code"
+                control={form.control}
+                render={({ field }) => (
+                  <FormItemLayout name="code" label="Authentication code">
+                    <FormControl_Shadcn_>
+                      <Input_Shadcn_
+                        id="code"
+                        autoFocus
+                        {...field}
+                        placeholder="XXXXXX"
+                        className="font-mono"
+                      />
+                    </FormControl_Shadcn_>
+                  </FormItemLayout>
+                )}
+              />
+            </form>
+          </Form_Shadcn_>
+        </div>
       )}
     </ConfirmationModal>
   )

apps/studio/components/interfaces/SignIn/SignInMfaForm.tsx

@@ -11,7 +11,7 @@ import { getReturnToPath } from 'lib/gotrue'
 import { Lock } from 'lucide-react'
 import Link from 'next/link'
 import { useRouter } from 'next/router'
-import { useEffect, useState } from 'react'
+import { useEffect, useRef, useState } from 'react'
 import { SubmitHandler, useForm } from 'react-hook-form'
 import { Button, Form_Shadcn_, FormControl_Shadcn_, FormField_Shadcn_, Input_Shadcn_ } from 'ui'
 import { FormItemLayout } from 'ui-patterns/form/FormItemLayout/FormItemLayout'
@@ -34,12 +34,15 @@ export const SignInMfaForm = ({ context = 'sign-in' }: SignInMfaFormProps) => {
   const router = useRouter()
   const signOut = useSignOut()
   const queryClient = useQueryClient()
+
   const [selectedFactor, setSelectedFactor] = useState<Factor | null>(null)
   const form = useForm<z.infer<typeof schema>>({
     resolver: zodResolver(schema),
     defaultValues: { code: '' },
   })
 
+  const { code } = form.watch()
+
   const {
     data: factors,
     error: factorsError,
@@ -89,6 +92,10 @@ export const SignInMfaForm = ({ context = 'sign-in' }: SignInMfaFormProps) => {
     }
   }, [factors?.totp, isSuccessFactors, router, queryClient])
 
+  useEffect(() => {
+    if (code.length === 6) form.handleSubmit(onSubmit)()
+  }, [code])
+
   const error = useAuthError()
 
   if (error) {
@@ -134,7 +141,7 @@ export const SignInMfaForm = ({ context = 'sign-in' }: SignInMfaFormProps) => {
                       </div>
                       <Input_Shadcn_
                         id="code"
-                        className="pl-10"
+                        className="pl-10 font-mono"
                         {...field}
                         autoFocus
                         autoComplete="off"
@@ -150,7 +157,7 @@ export const SignInMfaForm = ({ context = 'sign-in' }: SignInMfaFormProps) => {
               )}
             />
 
-            <div className="flex items-center justify-between space-x-2">
+            <div className="flex items-center justify-between gap-x-2">
               <Button
                 block
                 type="outline"