fix: wallet state corruption round 4 (#8519)

* wip: dont set wallet context device id until the wallet is fully connected

* fix: reset deviceId when disconnecting wallet

* chore: cleanup

* fix: prompt user to open native wallet on reload

* chore: move early return before isDemoWallet check

* chore: use better checks on receiveAddress react query skipToken

* chore: memoize wallet get from keyring

* chore: dont try respond to events for empty device ID

* chore: remove redundant coalesce

* chore: better typeof of SET_NATIVE_PENDING_DEVICE_ID action

* chore: cleanup logging of wallet ids

src/components/ManageAccountsDrawer/components/ImportAccounts.tsx

@@ -305,6 +305,11 @@ export const ImportAccounts = ({ chainId, onClose }: ImportAccountsProps) => {
   }, [])
 
   const handleDone = useCallback(async () => {
+    if (!walletDeviceId) {
+      console.error('Missing walletDeviceId')
+      return
+    }
+
     if (isDemoWallet) {
       walletDispatch({ type: WalletActions.SET_WALLET_MODAL, payload: true })
       accountManagementPopover.close()

src/components/MultiHopTrade/hooks/useReceiveAddress.tsx

@@ -62,41 +62,42 @@ export const useReceiveAddress = ({
       sellAccountId,
       buyAccountId,
     ],
-    queryFn: isInitializing
-      ? skipToken
-      : async () => {
-          // Already partially covered in isInitializing, but TypeScript lyfe mang.
-          if (!buyAsset || !wallet || !buyAccountId || !buyAccountMetadata) {
-            return undefined
-          }
+    queryFn:
+      !isInitializing && buyAsset && wallet && buyAccountId && buyAccountMetadata && deviceId
+        ? async () => {
+            // Already partially covered in isInitializing, but TypeScript lyfe mang.
+            if (!buyAsset || !wallet || !buyAccountId || !buyAccountMetadata || !deviceId) {
+              return undefined
+            }
 
-          const buyAssetChainId = buyAsset.chainId
-          const buyAssetAccountChainId = buyAccountId
-            ? fromAccountId(buyAccountId).chainId
-            : undefined
+            const buyAssetChainId = buyAsset.chainId
+            const buyAssetAccountChainId = buyAccountId
+              ? fromAccountId(buyAccountId).chainId
+              : undefined
 
-          /**
-           * do NOT remove
-           * super dangerous - don't use the wrong bip44 params to generate receive addresses
-           */
-          if (buyAssetChainId !== buyAssetAccountChainId) {
-            return undefined
-          }
+            /**
+             * do NOT remove
+             * super dangerous - don't use the wrong bip44 params to generate receive addresses
+             */
+            if (buyAssetChainId !== buyAssetAccountChainId) {
+              return undefined
+            }
 
-          if (isUtxoAccountId(buyAccountId) && !buyAccountMetadata?.accountType)
-            throw new Error(`Missing accountType for UTXO account ${buyAccountId}`)
+            if (isUtxoAccountId(buyAccountId) && !buyAccountMetadata?.accountType)
+              throw new Error(`Missing accountType for UTXO account ${buyAccountId}`)
 
-          const shouldFetchUnchainedAddress = Boolean(wallet && isLedger(wallet))
-          const walletReceiveAddress = await getReceiveAddress({
-            asset: buyAsset,
-            wallet,
-            accountMetadata: buyAccountMetadata,
-            deviceId,
-            pubKey: shouldFetchUnchainedAddress ? fromAccountId(buyAccountId).account : undefined,
-          })
+            const shouldFetchUnchainedAddress = Boolean(wallet && isLedger(wallet))
+            const walletReceiveAddress = await getReceiveAddress({
+              asset: buyAsset,
+              wallet,
+              accountMetadata: buyAccountMetadata,
+              deviceId,
+              pubKey: shouldFetchUnchainedAddress ? fromAccountId(buyAccountId).account : undefined,
+            })
 
-          return walletReceiveAddress
-        },
+            return walletReceiveAddress
+          }
+        : skipToken,
     staleTime: Infinity,
   })
 

src/context/WalletProvider/KeepKey/components/Passphrase.tsx

@@ -22,7 +22,7 @@ export const KeepKeyPassphrase = () => {
     state: { deviceId, keyring },
     dispatch,
   } = useWallet()
-  const wallet = keyring.get(deviceId)
+  const wallet = keyring.get(deviceId ?? '')
   const walletId = useAppSelector(selectWalletId)
   const appDispatch = useAppDispatch()
 

src/context/WalletProvider/KeepKey/components/Pin.tsx

@@ -2,7 +2,7 @@ import type { ButtonProps, SimpleGridProps } from '@chakra-ui/react'
 import { Alert, AlertDescription, AlertIcon, Button, Input, SimpleGrid } from '@chakra-ui/react'
 import type { Event } from '@shapeshiftoss/hdwallet-core'
 import type { KeyboardEvent } from 'react'
-import { useCallback, useEffect, useRef, useState } from 'react'
+import { useCallback, useEffect, useMemo, useRef, useState } from 'react'
 import { CircleIcon } from 'components/Icons/Circle'
 import { Text } from 'components/Text'
 import { WalletActions } from 'context/WalletProvider/actions'
@@ -38,7 +38,11 @@ export const KeepKeyPin = ({
     },
     dispatch,
   } = useWallet()
-  const wallet = keyring.get(deviceId)
+
+  const wallet = useMemo(() => {
+    if (!deviceId) return null
+    return keyring.get(deviceId)
+  }, [deviceId, keyring])
 
   const pinFieldRef = useRef<HTMLInputElement | null>(null)
 
@@ -138,10 +142,10 @@ export const KeepKeyPin = ({
       }
     }
 
-    keyring.on(['KeepKey', deviceId, String(MessageType.FAILURE)], handleError)
+    deviceId && keyring.on(['KeepKey', deviceId, String(MessageType.FAILURE)], handleError)
 
     return () => {
-      keyring.off(['KeepKey', deviceId, String(MessageType.FAILURE)], handleError)
+      deviceId && keyring.off(['KeepKey', deviceId, String(MessageType.FAILURE)], handleError)
     }
   }, [deviceId, keyring])
 

src/context/WalletProvider/Ledger/components/Chains.tsx

@@ -52,7 +52,9 @@ export const LedgerChains = () => {
 
   const handleConnectClick = useCallback(
     async (chainId: ChainId) => {
-      if (!walletState?.wallet) {
+      const { wallet, deviceId } = walletState
+
+      if (!wallet || !deviceId) {
         console.error('No wallet found')
         return
       }
@@ -67,7 +69,7 @@ export const LedgerChains = () => {
         ]({
           accountNumber: 0,
           chainIds,
-          wallet: walletState.wallet,
+          wallet,
           isSnapInstalled: false,
         })
 
@@ -102,7 +104,7 @@ export const LedgerChains = () => {
           const accountMetadata = accountMetadataByAccountId[accountId]
           const payload = {
             accountMetadataByAccountId: { [accountId]: accountMetadata },
-            walletId: walletState.deviceId,
+            walletId: deviceId,
           }
 
           dispatch(portfolio.actions.upsertAccountMetadata(payload))
@@ -120,7 +122,7 @@ export const LedgerChains = () => {
         setLoadingChains(prevLoading => ({ ...prevLoading, [chainId]: false }))
       }
     },
-    [dispatch, walletState.deviceId, walletState.wallet],
+    [dispatch, walletState],
   )
 
   const chainsRows = useMemo(

src/context/WalletProvider/NativeWallet/components/EnterPassword.tsx

@@ -43,7 +43,7 @@ export const EnterPassword = () => {
   const translate = useTranslate()
   const { state, dispatch, disconnect } = useWallet()
   const localWallet = useLocalWallet()
-  const { deviceId, keyring } = state
+  const { nativeWalletPendingDeviceId: deviceId, keyring } = state
 
   const [showPw, setShowPw] = useState<boolean>(false)
 
@@ -58,6 +58,7 @@ export const EnterPassword = () => {
   const onSubmit = useCallback(
     async (values: FieldValues) => {
       try {
+        if (!deviceId) return
         const wallet = keyring.get<NativeHDWallet>(deviceId)
         const Vault = await import('@shapeshiftoss/hdwallet-native-vault').then(m => m.Vault)
         const vault = await Vault.open(deviceId, values.password)
@@ -83,6 +84,7 @@ export const EnterPassword = () => {
           type: WalletActions.SET_IS_CONNECTED,
           payload: true,
         })
+        dispatch({ type: WalletActions.RESET_NATIVE_PENDING_DEVICE_ID })
         dispatch({ type: WalletActions.SET_LOCAL_WALLET_LOADING, payload: false })
         dispatch({ type: WalletActions.SET_WALLET_MODAL, payload: false })
       } catch (e) {

src/context/WalletProvider/NativeWallet/components/NativeLoad.tsx

@@ -81,6 +81,13 @@ export const NativeLoad = ({ history }: RouteComponentProps) => {
     if (adapter) {
       const { name, icon } = NativeConfig
       try {
+        // Set a pending device ID so the event handler doesn't redirect the user to password input
+        // for the previous wallet
+        dispatch({
+          type: WalletActions.SET_NATIVE_PENDING_DEVICE_ID,
+          payload: deviceId,
+        })
+
         const wallet = await adapter.pairDevice(deviceId)
         if (!(await wallet?.isInitialized())) {
           // This will trigger the password modal and the modal will set the wallet on state
@@ -103,6 +110,7 @@ export const NativeLoad = ({ history }: RouteComponentProps) => {
             type: WalletActions.SET_IS_CONNECTED,
             payload: true,
           })
+          dispatch({ type: WalletActions.RESET_NATIVE_PENDING_DEVICE_ID })
           // The wallet is already initialized so we can close the modal
           dispatch({ type: WalletActions.SET_WALLET_MODAL, payload: false })
         }

src/context/WalletProvider/NativeWallet/hooks/useNativeEventHandler.ts

@@ -6,20 +6,28 @@ import type { ActionTypes } from 'context/WalletProvider/actions'
 import { WalletActions } from 'context/WalletProvider/actions'
 import type { InitialState } from 'context/WalletProvider/WalletProvider'
 import { isMobile as isMobileApp } from 'lib/globals'
+import { assertUnreachable } from 'lib/utils'
 
 export const useNativeEventHandler = (state: InitialState, dispatch: Dispatch<ActionTypes>) => {
   const { keyring, modal, modalType } = state
 
   useEffect(() => {
     const handleEvent = (e: [deviceId: string, message: Event]) => {
       const deviceId = e[0]
-      switch (e[1].message_type) {
+      const messageType = e[1].message_type as NativeEvents
+      switch (messageType) {
         case NativeEvents.MNEMONIC_REQUIRED:
           if (!deviceId) break
+
+          // Don't show password input for previous wallet when switching
+          if (deviceId !== state.nativeWalletPendingDeviceId) {
+            break
+          }
+
           // If we're on the native mobile app we don't need to handle the MNEMONIC_REQUIRED event as we use the device's native authentication instead
           // Reacting to this event will incorrectly open the native password modal after authentication completes when on the mobile app
           if (isMobileApp) break
-          dispatch({ type: WalletActions.NATIVE_PASSWORD_OPEN, payload: { modal: true, deviceId } })
+          dispatch({ type: WalletActions.NATIVE_PASSWORD_OPEN, payload: { modal: true } })
 
           break
         case NativeEvents.READY:
@@ -28,8 +36,7 @@ export const useNativeEventHandler = (state: InitialState, dispatch: Dispatch<Ac
           }
           break
         default:
-          // If there wasn't an enum value, then we'll check the message type
-          break
+          assertUnreachable(messageType)
       }
     }
 
@@ -48,5 +55,5 @@ export const useNativeEventHandler = (state: InitialState, dispatch: Dispatch<Ac
       keyring.off(['Native', '*', NativeEvents.MNEMONIC_REQUIRED], handleEvent)
       keyring.off(['Native', '*', NativeEvents.READY], handleEvent)
     }
-  }, [modalType, dispatch, keyring, modal])
+  }, [modalType, dispatch, keyring, modal, state.nativeWalletPendingDeviceId, state.deviceId])
 }

src/context/WalletProvider/WalletConnectV2/useWalletConnectV2EventHandler.ts

@@ -21,8 +21,7 @@ export const useWalletConnectV2EventHandler = (
      */
     state.wallet?.disconnect?.()
     dispatch({ type: WalletActions.RESET_STATE })
-    localWallet.clearLocalWallet()
-  }, [dispatch, localWallet, state.wallet])
+  }, [dispatch, state.wallet])
 
   useEffect(() => {
     // This effect should never run for wallets other than WalletConnectV2 since we explicitly tap into @walletconnect/ethereum-provider provider

src/context/WalletProvider/WalletProvider.tsx

@@ -25,7 +25,7 @@ import {
   selectWalletRdns,
   selectWalletType,
 } from 'state/slices/localWalletSlice/selectors'
-import { portfolio } from 'state/slices/portfolioSlice/portfolioSlice'
+import { portfolio as portfolioSlice } from 'state/slices/portfolioSlice/portfolioSlice'
 import { store } from 'state/store'
 
 import type { ActionTypes } from './actions'
@@ -93,11 +93,12 @@ export type InitialState = {
   isLocked: boolean
   modal: boolean
   isLoadingLocalWallet: boolean
-  deviceId: string
+  deviceId: string | null
   showBackButton: boolean
   keepKeyPinRequestType: PinMatrixRequestType | null
   deviceState: DeviceState
   disconnectOnCloseModal: boolean
+  nativeWalletPendingDeviceId: string | null
 } & (
   | {
       modalType: KeyManager | null
@@ -124,11 +125,12 @@ const initialState: InitialState = {
   isLocked: false,
   modal: false,
   isLoadingLocalWallet: false,
-  deviceId: '',
+  deviceId: null,
   showBackButton: true,
   keepKeyPinRequestType: null,
   deviceState: initialDeviceState,
   disconnectOnCloseModal: false,
+  nativeWalletPendingDeviceId: null,
 }
 
 const reducer = (state: InitialState, action: ActionTypes): InitialState => {
@@ -140,14 +142,15 @@ const reducer = (state: InitialState, action: ActionTypes): InitialState => {
       if (currentConnectedType === 'walletconnectv2') {
         state.wallet?.disconnect?.()
         store.dispatch(localWalletSlice.actions.clearLocalWallet())
+        store.dispatch(portfolioSlice.actions.setWalletMeta(undefined))
       }
       const { deviceId, name, wallet, icon, meta, isDemoWallet, connectedType } = action.payload
       // set wallet metadata in redux store
       const walletMeta = {
         walletId: deviceId,
         walletName: name,
       }
-      store.dispatch(portfolio.actions.setWalletMeta(walletMeta))
+      store.dispatch(portfolioSlice.actions.setWalletMeta(walletMeta))
       return {
         ...state,
         deviceId,
@@ -221,7 +224,8 @@ const reducer = (state: InitialState, action: ActionTypes): InitialState => {
         modal: action.payload.modal,
         modalType: KeyManager.Native,
         showBackButton: !state.isLoadingLocalWallet,
-        deviceId: action.payload.deviceId,
+        deviceId: null,
+        walletInfo: null,
         initialRoute: NativeWalletRoutes.EnterPassword,
       }
     case WalletActions.OPEN_KEEPKEY_PIN: {
@@ -290,9 +294,10 @@ const reducer = (state: InitialState, action: ActionTypes): InitialState => {
     case WalletActions.SET_LOCAL_WALLET_LOADING:
       return { ...state, isLoadingLocalWallet: action.payload }
     case WalletActions.RESET_STATE:
-      const resetProperties = omit(initialState, ['keyring', 'adapters', 'modal', 'deviceId'])
+      const resetProperties = omit(initialState, ['keyring', 'adapters', 'modal'])
       // reset wallet meta in redux store
-      store.dispatch(portfolio.actions.setWalletMeta(undefined))
+      store.dispatch(localWalletSlice.actions.clearLocalWallet())
+      store.dispatch(portfolioSlice.actions.setWalletMeta(undefined))
       return { ...state, ...resetProperties }
     // TODO: Remove this once we update SET_DEVICE_STATE to allow explicitly setting falsey values
     case WalletActions.RESET_LAST_DEVICE_INTERACTION_STATE: {
@@ -320,6 +325,21 @@ const reducer = (state: InitialState, action: ActionTypes): InitialState => {
         modalType: KeyManager.KeepKey,
         initialRoute: KeepKeyRoutes.Disconnect,
       }
+    case WalletActions.SET_NATIVE_PENDING_DEVICE_ID:
+      store.dispatch(localWalletSlice.actions.clearLocalWallet())
+      store.dispatch(portfolioSlice.actions.setWalletMeta(undefined))
+      return {
+        ...state,
+        isConnected: false,
+        deviceId: null,
+        walletInfo: null,
+        nativeWalletPendingDeviceId: action.payload,
+      }
+    case WalletActions.RESET_NATIVE_PENDING_DEVICE_ID:
+      return {
+        ...state,
+        nativeWalletPendingDeviceId: null,
+      }
     default:
       return state
   }
@@ -334,6 +354,7 @@ const getInitialState = () => {
      */
     return {
       ...initialState,
+      nativeWalletPendingDeviceId: localWalletDeviceId,
       isLoadingLocalWallet: true,
     }
   }
@@ -413,7 +434,6 @@ export const WalletProvider = ({ children }: { children: React.ReactNode }): JSX
      */
     state.wallet?.disconnect?.()
     dispatch({ type: WalletActions.RESET_STATE })
-    store.dispatch(localWalletSlice.actions.clearLocalWallet())
   }, [state.wallet])
 
   const load = useCallback(() => {