feat: signer account page to enable MFA, deviceFactor recovery on login

src/components/common/ConnectWallet/MPCWallet.tsx

@@ -1,9 +1,14 @@
+import { MPCWalletState } from '@/hooks/wallets/mpc/useMPCWallet'
 import { Box, Button, CircularProgress } from '@mui/material'
 import { useContext } from 'react'
 import { MpcWalletContext } from './MPCWalletProvider'
+import { PasswordRecovery } from './PasswordRecovery'
 
 export const MPCWallet = () => {
-  const { loginPending, triggerLogin, resetAccount, userInfo } = useContext(MpcWalletContext)
+  const { loginPending, triggerLogin, resetAccount, userInfo, walletState, recoverFactorWithPassword } =
+    useContext(MpcWalletContext)
+
+  console.log(walletState)
 
   return (
     <>
@@ -28,6 +33,10 @@ export const MPCWallet = () => {
           )}
         </Button>
       )}
+
+      {walletState === MPCWalletState.MANUAL_RECOVERY && (
+        <PasswordRecovery recoverFactorWithPassword={recoverFactorWithPassword} />
+      )}
     </>
   )
 }

src/components/common/ConnectWallet/MPCWalletProvider.tsx

@@ -6,7 +6,7 @@ type MPCWalletContext = {
   triggerLogin: () => Promise<void>
   resetAccount: () => Promise<void>
   upsertPasswordBackup: (password: string) => Promise<void>
-  recoverFactorWithPassword: (password: string) => Promise<void>
+  recoverFactorWithPassword: (password: string, storeDeviceFactor: boolean) => Promise<void>
   walletState: MPCWalletState
   userInfo: {
     email: string | undefined

src/components/common/ConnectWallet/PasswordRecovery.tsx

@@ -0,0 +1,64 @@
+import { VisibilityOff, Visibility } from '@mui/icons-material'
+import {
+  DialogContent,
+  Typography,
+  TextField,
+  IconButton,
+  FormControlLabel,
+  Checkbox,
+  Button,
+  Box,
+} from '@mui/material'
+import { useState } from 'react'
+import ModalDialog from '../ModalDialog'
+
+export const PasswordRecovery = ({
+  recoverFactorWithPassword,
+}: {
+  recoverFactorWithPassword: (password: string, storeDeviceFactor: boolean) => Promise<void>
+}) => {
+  const [showPassword, setShowPassword] = useState(false)
+  const [recoveryPassword, setRecoveryPassword] = useState<string>('')
+  const [storeDeviceFactor, setStoreDeviceFactor] = useState(false)
+  return (
+    <ModalDialog open dialogTitle="Enter your recovery password" hideChainIndicator>
+      <DialogContent>
+        <Box>
+          <Typography>
+            This browser is not registered with your Account yet. Please enter your recovery password to restore access
+            to this account.
+          </Typography>
+          <Box mt={2} display="flex" flexDirection="column" alignItems="baseline" gap={2}>
+            <TextField
+              label="Recovery password"
+              type={showPassword ? 'text' : 'password'}
+              value={recoveryPassword}
+              onChange={(event) => {
+                setRecoveryPassword(event.target.value)
+              }}
+              InputProps={{
+                endAdornment: (
+                  <IconButton
+                    aria-label="toggle password visibility"
+                    onClick={() => setShowPassword((prev) => !prev)}
+                    edge="end"
+                  >
+                    {showPassword ? <VisibilityOff /> : <Visibility />}
+                  </IconButton>
+                ),
+              }}
+            />
+            <FormControlLabel
+              control={<Checkbox checked={storeDeviceFactor} onClick={() => setStoreDeviceFactor((prev) => !prev)} />}
+              label="Do not ask again on this device"
+            />
+
+            <Button variant="contained" onClick={() => recoverFactorWithPassword(recoveryPassword, storeDeviceFactor)}>
+              Submit
+            </Button>
+          </Box>
+        </Box>
+      </DialogContent>
+    </ModalDialog>
+  )
+}

src/components/settings/SignerAccountMFA/index.tsx

@@ -1,31 +1,79 @@
 import useMPC from '@/hooks/wallets/mpc/useMPC'
-import { Box, Button, Typography } from '@mui/material'
+import { Box, Button, Checkbox, FormControlLabel, TextField, Typography } from '@mui/material'
 import { COREKIT_STATUS } from '@web3auth/mpc-core-kit'
 import { getPubKeyPoint } from '@tkey-mpc/common-types'
-import useMFASettings from './useMFASettings'
 import { BN } from 'bn.js'
-import { useState } from 'react'
+import { useEffect, useMemo, useState } from 'react'
+import { useSecurityQuestions } from '@/hooks/wallets/mpc/recovery/useSecurityQuestions'
+import useMFASettings from './useMFASettings'
+import { useForm } from 'react-hook-form'
+import { useDeviceShare } from '@/hooks/wallets/mpc/recovery/useDeviceShare'
+
+type SignerAccountFormData = {
+  oldPassword: string | undefined
+  newPassword: string
+  confirmPassword: string
+  storeDeviceShare: boolean
+}
 
 const SignerAccountMFA = () => {
   const mpcCoreKit = useMPC()
-  const mfaSettings = useMFASettings()
+  const mfaSettings = useMFASettings(mpcCoreKit)
+  const securityQuestions = useSecurityQuestions(mpcCoreKit)
+  const deviceShareModule = useDeviceShare(mpcCoreKit)
+
+  const formMethods = useForm<SignerAccountFormData>({
+    mode: 'all',
+  })
+
+  const { register, formState, watch, setValue, handleSubmit } = formMethods
 
   const [enablingMFA, setEnablingMFA] = useState(false)
 
+  const isPasswordSet = useMemo(() => securityQuestions.isEnabled(), [securityQuestions])
+
+  console.log(mpcCoreKit)
+
+  useEffect(() => {
+    deviceShareModule.isEnabled().then((value) => setValue('storeDeviceShare', value))
+  }, [deviceShareModule, setValue])
+
   const enableMFA = async () => {
     if (!mpcCoreKit) {
       return
     }
 
     setEnablingMFA(true)
     try {
-      // First enable MFA in mpcCoreKit
-      const recoveryFactor = await mpcCoreKit.enableMFA({})
+      const { newPassword, oldPassword, storeDeviceShare } = formMethods.getValues()
+      // 1. setup device factor with password recovery
+      await securityQuestions.upsertPassword(newPassword, oldPassword)
+      const securityQuestionFactor = await securityQuestions.recoverWithPassword(newPassword)
+      if (!securityQuestionFactor) {
+        throw Error('Could not recover using the new password recovery')
+      }
 
-      // Then remove the recovery factor the mpcCoreKit creates
-      const recoverKey = new BN(recoveryFactor, 'hex')
-      const recoverPubKey = getPubKeyPoint(recoverKey)
-      await mpcCoreKit.deleteFactor(recoverPubKey, recoverKey)
+      if (!mfaSettings?.mfaEnabled) {
+        // 2. enable MFA in mpcCoreKit
+        const recoveryFactor = await mpcCoreKit.enableMFA({})
+
+        // 3. remove the recovery factor the mpcCoreKit creates
+        const recoverKey = new BN(recoveryFactor, 'hex')
+        const recoverPubKey = getPubKeyPoint(recoverKey)
+        await mpcCoreKit.deleteFactor(recoverPubKey, recoverKey)
+      }
+
+      const hasDeviceShare = await deviceShareModule.isEnabled()
+      if (hasDeviceShare !== storeDeviceShare) {
+        if (storeDeviceShare) {
+          await deviceShareModule.createAndStoreDeviceFactor()
+        } else {
+          // Switch to password recovery factor such that we can delete the device factor
+          await mpcCoreKit.inputFactorKey(new BN(securityQuestionFactor, 'hex'))
+
+          await deviceShareModule.removeDeviceFactor()
+        }
+      }
     } catch (error) {
       console.error(error)
     } finally {
@@ -41,16 +89,75 @@ const SignerAccountMFA = () => {
     )
   }
 
+  const onSubmit = async () => {
+    console.log('submitting')
+    await enableMFA()
+  }
+
   return (
-    <Box>
-      {mfaSettings?.mfaEnabled ? (
-        <Typography>MFA is enabled!</Typography>
-      ) : (
-        <Button disabled={enablingMFA} onClick={enableMFA}>
-          Enable MFA
+    <form onSubmit={handleSubmit(onSubmit)}>
+      <Box display="flex" flexDirection="column" gap={3} alignItems="baseline">
+        {
+          /* TODO: Memoize this*/ securityQuestions.isEnabled() ? (
+            <Typography>You already have a recovery password setup.</Typography>
+          ) : (
+            <Typography>You have no password setup. Secure your account now!</Typography>
+          )
+        }
+        {isPasswordSet && (
+          <TextField
+            placeholder="Old password"
+            label="Old password"
+            type="password"
+            error={!!formState.errors['oldPassword']}
+            helperText={formState.errors['oldPassword']?.message}
+            {...register('oldPassword', {
+              required: securityQuestions.isEnabled(),
+            })}
+          />
+        )}
+        <TextField
+          placeholder="New password"
+          label="New password"
+          type="password"
+          error={!!formState.errors['newPassword']}
+          helperText={formState.errors['newPassword']?.message}
+          {...register('newPassword', {
+            required: true,
+            minLength: 6,
+          })}
+        />
+        <TextField
+          placeholder="Confirm new password"
+          label="Confirm new password"
+          type="password"
+          error={!!formState.errors['confirmPassword']}
+          helperText={formState.errors['confirmPassword']?.message}
+          {...register('confirmPassword', {
+            required: true,
+            validate: (value: string) => {
+              const currentNewPW = watch('newPassword')
+              if (value !== currentNewPW) {
+                return 'Passwords do not match'
+              }
+            },
+          })}
+        />
+
+        <FormControlLabel
+          control={<Checkbox {...register('storeDeviceShare')} />}
+          label="Do not ask for second factor on this device"
+        />
+
+        <Button
+          sx={{ justifySelf: 'flex-start' }}
+          disabled={!formMethods.formState.isValid || enablingMFA}
+          type="submit"
+        >
+          Change
         </Button>
-      )}
-    </Box>
+      </Box>
+    </form>
   )
 }
 

src/components/settings/SignerAccountMFA/useMFASettings.ts

@@ -1,13 +1,10 @@
-import useMPC from '@/hooks/wallets/mpc/useMPC'
-import { COREKIT_STATUS } from '@web3auth/mpc-core-kit'
+import { COREKIT_STATUS, type Web3AuthMPCCoreKit } from '@web3auth/mpc-core-kit'
 
 export type MFASettings = {
   mfaEnabled: boolean
 } | null
 
-const useMFASettings = () => {
-  const mpcCoreKit = useMPC()
-
+const useMFASettings = (mpcCoreKit: Web3AuthMPCCoreKit | undefined) => {
   if (mpcCoreKit?.status !== COREKIT_STATUS.LOGGED_IN) {
     return null
   }

src/hooks/wallets/mpc/__tests__/useMPC.test.ts

@@ -1,6 +1,6 @@
 import * as useOnboard from '@/hooks/wallets/useOnboard'
 import { renderHook, waitFor } from '@/tests/test-utils'
-import { getMPCCoreKitInstance, setMPCCoreKitInstance, useInitMPC } from '../useMPC'
+import { _getMPCCoreKitInstance, setMPCCoreKitInstance, useInitMPC } from '../useMPC'
 import * as useChains from '@/hooks/useChains'
 import { type ChainInfo, RPC_AUTHENTICATION } from '@safe-global/safe-gateway-typescript-sdk'
 import { hexZeroPad } from 'ethers/lib/utils'
@@ -104,7 +104,7 @@ describe('useInitMPC', () => {
     renderHook(() => useInitMPC())
 
     await waitFor(() => {
-      expect(getMPCCoreKitInstance()).toBeDefined()
+      expect(_getMPCCoreKitInstance()).toBeDefined()
       expect(connectWalletSpy).not.toBeCalled()
     })
   })
@@ -151,7 +151,7 @@ describe('useInitMPC', () => {
 
     await waitFor(() => {
       expect(connectWalletSpy).toBeCalled()
-      expect(getMPCCoreKitInstance()).toBeDefined()
+      expect(_getMPCCoreKitInstance()).toBeDefined()
     })
   })
 
@@ -215,7 +215,7 @@ describe('useInitMPC', () => {
 
     await waitFor(() => {
       expect(mockChainChangedListener).toHaveBeenCalledWith('0x5')
-      expect(getMPCCoreKitInstance()).toBeDefined()
+      expect(_getMPCCoreKitInstance()).toBeDefined()
       expect(connectWalletSpy).not.toBeCalled()
     })
   })

src/hooks/wallets/mpc/recovery/useDeviceShare.ts

@@ -0,0 +1,51 @@
+import {
+  BrowserStorage,
+  getWebBrowserFactor,
+  storeWebBrowserFactor,
+  TssShareType,
+  type Web3AuthMPCCoreKit,
+} from '@web3auth/mpc-core-kit'
+import BN from 'bn.js'
+import { getPubKeyPoint } from '@tkey-mpc/common-types'
+
+export const useDeviceShare = (mpcCoreKit: Web3AuthMPCCoreKit | undefined) => {
+  const isEnabled = async () => {
+    if (!mpcCoreKit || !mpcCoreKit.tKey.metadata) {
+      return false
+    }
+    return !!(await getWebBrowserFactor(mpcCoreKit))
+  }
+
+  const createAndStoreDeviceFactor = async () => {
+    if (!mpcCoreKit) {
+      return
+    }
+    const userAgent = navigator.userAgent
+
+    const deviceFactorKey = new BN(
+      await mpcCoreKit.createFactor({ shareType: TssShareType.DEVICE, additionalMetadata: { userAgent } }),
+      'hex',
+    )
+    await storeWebBrowserFactor(deviceFactorKey, mpcCoreKit)
+  }
+
+  const removeDeviceFactor = async () => {
+    if (!mpcCoreKit) {
+      return
+    }
+    const deviceFactor = await getWebBrowserFactor(mpcCoreKit)
+    const key = new BN(deviceFactor, 'hex')
+    const pubKey = getPubKeyPoint(key)
+    const pubKeyX = pubKey.x.toString('hex', 64)
+    await mpcCoreKit.deleteFactor(pubKey)
+    const currentStorage = BrowserStorage.getInstance('mpc_corekit_store')
+    debugger
+    currentStorage.set(pubKeyX, undefined)
+  }
+
+  return {
+    isEnabled,
+    createAndStoreDeviceFactor,
+    removeDeviceFactor,
+  }
+}