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 { Lock } from 'lucide-react'
import Link from 'next/link'
import { useRouter } from 'next/router'
import { useEffect, useState } from 'react'
import { SubmitHandler, useForm } from 'react-hook-form'
import z from 'zod'

import AlertError from 'components/ui/AlertError'
import { GenericSkeletonLoader } from 'components/ui/ShimmeringLoader'
import { useMfaChallengeAndVerifyMutation } from 'data/profile/mfa-challenge-and-verify-mutation'
import { useMfaListFactorsQuery } from 'data/profile/mfa-list-factors-query'
import { useSignOut } from 'lib/auth'
import { getReturnToPath } from 'lib/gotrue'
import { Button, Form_Shadcn_, FormControl_Shadcn_, FormField_Shadcn_, Input_Shadcn_ } from 'ui'
import { FormItemLayout } from 'ui-patterns/form/FormItemLayout/FormItemLayout'

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

  const {
    data: factors,
    error: factorsError,
    isError: isErrorFactors,
    isSuccess: isSuccessFactors,
    isLoading: isLoadingFactors,
  } = useMfaListFactorsQuery()
  const {
    mutate: mfaChallengeAndVerify,
    isLoading: isVerifying,
    isSuccess,
  } = useMfaChallengeAndVerifyMutation({
    onSuccess: async () => {
      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])

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

      {isErrorFactors && <AlertError error={factorsError} subject="Failed to retrieve factors" />}

      {isSuccessFactors && (
        <Form_Shadcn_ {...form}>
          <form id={formId} 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={
                    selectedFactor && factors?.totp.length === 2
                      ? `Code generated by ${selectedFactor.friendly_name}`
                      : null
                  }
                >
                  <FormControl_Shadcn_>
                    <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_Shadcn_
                        id="code"
                        className="pl-10"
                        {...field}
                        autoFocus
                        autoComplete="off"
                        autoCorrect="off"
                        autoCapitalize="none"
                        spellCheck="false"
                        placeholder="XXXXXX"
                        disabled={isVerifying}
                      />
                    </div>
                  </FormControl_Shadcn_>
                </FormItemLayout>
              )}
            />

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

      <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>
              <a
                className="text-sm text-foreground-light hover:text-foreground cursor-pointer"
                onClick={() =>
                  setSelectedFactor(factors.totp.find((f) => f.id !== selectedFactor?.id)!)
                }
              >{`Authenticate using ${
                factors.totp.find((f) => f.id !== selectedFactor?.id)?.friendly_name
              }?`}</a>
            </li>
          )}
          <li>
            <Link
              href="/logout"
              className="text-sm transition text-foreground-light hover:text-foreground"
            >
              Force sign out and clear cookies
            </Link>
          </li>
          <li>
            <Link
              target="_blank"
              rel="noreferrer"
              href="/support/new?subject=Unable+to+sign+in+via+MFA&category=Login_issues"
              className="text-sm transition text-foreground-light hover:text-foreground"
            >
              Reach out to us via support
            </Link>
          </li>
        </ul>
      </div>
    </>
  )
}