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

import { zodResolver } from '@hookform/resolvers/zod'
import type { Factor } from '@supabase/supabase-js'
import { useQueryClient } from '@tanstack/react-query'
import { useAuthError, useFlag, useParams } from 'common'
import { Lock } from 'lucide-react'
import Link from 'next/link'
import { useSearchParams } from 'next/navigation'
import { useRouter } from 'next/router'
import { useEffect, useState } from 'react'
import { SubmitHandler, useForm, useWatch } from 'react-hook-form'
import { Button, Form, FormControl, FormField, Input } from 'ui'
import { FormItemLayout } from 'ui-patterns/form/FormItemLayout/FormItemLayout'
import { GenericSkeletonLoader } from 'ui-patterns/ShimmeringLoader'
import z from 'zod'

import {
  getFactorDisplayName,
  getFactorFriendlyName,
  MFA_SUPPORT_EMAIL_HREF,
} from './SignInMfa.utils'
import { AlertError } from '@/components/ui/AlertError'
import { useAddLoginEvent } from '@/data/misc/audit-login-mutation'
import { useMfaChallengeAndVerifyMutation } from '@/data/profile/mfa-challenge-and-verify-mutation'
import { useMfaListFactorsQuery } from '@/data/profile/mfa-list-factors-query'
import { useRecoveryCodesStatusQuery } from '@/data/recovery-codes/recovery-codes-status-query'
import { useSignOut } from '@/lib/auth'
import { getReturnToPath } from '@/lib/gotrue'
import { useTrack } from '@/lib/telemetry/track'

const schema = z.object({
  code: z.string().min(1, 'MFA Code is required'),
})

const formId = 'sign-in-mfa-form'

interface SignInMfaFormProps {
  context?: 'forgot-password' | 'sign-in'
}

export const SignInMfaForm = ({ context = 'sign-in' }: SignInMfaFormProps) => {
  const router = useRouter()
  const signOut = useSignOut()
  const queryClient = useQueryClient()
  const searchParams = useSearchParams()
  const { method: signInMethod = 'unknown' } = useParams()
  const enableAuthRecoveryCodes = useFlag('enableAuthRecoveryCodes')
  const { data: recoveryCodesStatus } = useRecoveryCodesStatusQuery({
    enabled: enableAuthRecoveryCodes,
  })

  const track = useTrack()
  const { mutate: addLoginEvent } = useAddLoginEvent()

  const [selectedFactor, setSelectedFactor] = useState<Factor | null>(null)
  const form = useForm<z.infer<typeof schema>>({
    resolver: zodResolver(schema),
    defaultValues: { code: '' },
  })

  const code = useWatch({ control: form.control, name: 'code' })

  const {
    data: factors,
    error: factorsError,
    isError: isErrorFactors,
    isSuccess: isSuccessFactors,
    isPending: isLoadingFactors,
  } = useMfaListFactorsQuery()
  const {
    mutate: mfaChallengeAndVerify,
    isPending: isVerifying,
    isSuccess,
  } = useMfaChallengeAndVerifyMutation({
    onSuccess: async () => {
      if (context === 'sign-in') {
        track('sign_in', { category: 'account', method: signInMethod })
        addLoginEvent({})
      }

      await queryClient.resetQueries()

      if (context === 'forgot-password') {
        router.push({
          pathname: '/reset-password',
          query: router.query,
        })
      } else {
        router.push(getReturnToPath())
      }
    },
  })

  const onClickLogout = async () => {
    await signOut()
    await router.replace('/sign-in')
  }

  const onSubmit: SubmitHandler<z.infer<typeof schema>> = async ({ code }) => {
    if (selectedFactor) {
      mfaChallengeAndVerify({ factorId: selectedFactor.id, code, refreshFactors: false })
    }
  }

  useEffect(() => {
    if (isSuccessFactors) {
      // if the user wanders into this page and he has no MFA setup, send the user to the next screen
      if (factors.totp.length === 0) {
        queryClient.resetQueries().then(() => router.push(getReturnToPath()))
      }
      if (factors.totp.length > 0) {
        setSelectedFactor(factors.totp[0])
      }
    }
  }, [factors?.totp, isSuccessFactors, router, queryClient])

  useEffect(() => {
    if (code.length === 6) form.handleSubmit(onSubmit)()
  }, [code])

  const error = useAuthError()
  const selectedFactorName = getFactorFriendlyName(selectedFactor)

  if (error) {
    return (
      <AlertError
        error={error}
        subject="Error while signing in"
        hideContactSupport
        additionalActions={
          <>
            <Button asChild>
              <Link href="/sign-in">Back to sign in</Link>
            </Button>
            <Button asChild>
              <a href={MFA_SUPPORT_EMAIL_HREF}>Email support</a>
            </Button>
          </>
        }
      />
    )
  }

  return (
    <>
      {isLoadingFactors && <GenericSkeletonLoader />}

      {isErrorFactors && (
        <AlertError
          error={factorsError}
          subject="Failed to retrieve factors"
          description="Try refreshing your browser. If the issue persists, email support@supabase.com."
          hideContactSupport
        />
      )}

      {isSuccessFactors && (
        <Form {...form}>
          <form
            id={formId}
            method="POST"
            className="flex flex-col gap-4"
            onSubmit={form.handleSubmit(onSubmit)}
          >
            <FormField
              key="code"
              name="code"
              control={form.control}
              render={({ field }) => (
                <FormItemLayout
                  name="code"
                  label="Authentication code"
                  description={
                    selectedFactorName ? (
                      <>
                        From the authenticator app registered as{' '}
                        <strong className="text-foreground">{selectedFactorName}</strong>
                      </>
                    ) : (
                      'From the authenticator app registered to this account'
                    )
                  }
                >
                  <FormControl>
                    <div className="relative">
                      <div className="absolute inset-y-0 left-0 pl-3 flex items-center pointer-events-none text-foreground-light [&_svg]:stroke-[1.5] [&_svg]:h-[20px] [&_svg]:w-[20px]">
                        <Lock />
                      </div>
                      <Input
                        id="code"
                        className="pl-10 font-mono"
                        {...field}
                        autoFocus
                        autoComplete="off"
                        autoCorrect="off"
                        autoCapitalize="none"
                        spellCheck="false"
                        placeholder="XXXXXX"
                        disabled={isVerifying}
                      />
                    </div>
                  </FormControl>
                </FormItemLayout>
              )}
            />

            <div className="flex items-center justify-between gap-x-2">
              <Button
                block
                variant="outline"
                size="large"
                disabled={isVerifying || isSuccess}
                onClick={onClickLogout}
                className="opacity-80 hover:opacity-100 transition"
              >
                Cancel
              </Button>
              <Button
                variant="primary"
                block
                form={formId}
                type="submit"
                size="large"
                disabled={isVerifying || isSuccess}
                loading={isVerifying || isSuccess}
              >
                {isVerifying ? 'Verifying' : isSuccess ? 'Signing in' : 'Verify'}
              </Button>
            </div>
          </form>
        </Form>
      )}

      <div className="my-8">
        <div className="text-sm">
          <span className="text-foreground-light">Unable to sign in?</span>{' '}
        </div>
        <ul className="list-disc pl-6">
          {factors?.totp.length === 2 && (
            <li>
              <button
                tabIndex={0}
                className="text-sm text-foreground-light hover:text-foreground cursor-pointer"
                onClick={() =>
                  setSelectedFactor(factors.totp.find((f) => f.id !== selectedFactor?.id)!)
                }
              >
                Authenticate using{' '}
                <strong className="text-foreground">
                  {getFactorDisplayName(factors.totp.find((f) => f.id !== selectedFactor?.id))}
                </strong>
                ?
              </button>
            </li>
          )}
          {enableAuthRecoveryCodes && recoveryCodesStatus?.status === 'available' && (
            <li>
              <Link
                href={`/sign-in-recovery-code?${searchParams}`}
                className="text-sm transition text-foreground-light hover:text-foreground"
              >
                Authenticate using a recovery code
              </Link>
            </li>
          )}
          <li>
            <Link
              href="/logout"
              className="text-sm transition text-foreground-light hover:text-foreground"
            >
              Force sign out and clear cookies
            </Link>
          </li>
          <li>
            <a
              href={MFA_SUPPORT_EMAIL_HREF}
              className="text-sm transition text-foreground-light hover:text-foreground"
            >
              Email support to regain access
            </a>
          </li>
        </ul>
      </div>
    </>
  )
}