fix: wallet state corruption round 4

src/App.tsx

@@ -11,7 +11,13 @@ import { IconCircle } from 'components/IconCircle'
 import { useBridgeClaimNotification } from 'hooks/useBridgeClaimNotification/useBridgeClaimNotification'
 import { useHasAppUpdated } from 'hooks/useHasAppUpdated/useHasAppUpdated'
 import { useModal } from 'hooks/useModal/useModal'
-import { selectShowConsentBanner, selectShowWelcomeModal } from 'state/slices/selectors'
+import { useWallet } from 'hooks/useWallet/useWallet'
+import {
+  selectShowConsentBanner,
+  selectShowWelcomeModal,
+  selectWalletId,
+} from 'state/slices/selectors'
+import { useAppSelector } from 'state/store'
 
 const flexGap = { base: 2, md: 3 }
 const flexDir: ResponsiveValue<Property.FlexDirection> = { base: 'column', md: 'row' }
@@ -69,6 +75,19 @@ export const App = () => {
     }
   }, [isNativeOnboardOpen, openNativeOnboard, showWelcomeModal])
 
+  const portfolioWalletId = useAppSelector(selectWalletId)
+
+  const {
+    state: { deviceId: contextDeviceId },
+  } = useWallet()
+
+  useEffect(() => {
+    console.log({
+      portfolioWalletId,
+      contextDeviceId,
+    })
+  }, [contextDeviceId, portfolioWalletId])
+
   return (
     <>
       {showConsentBanner && <ConsentBanner />}

src/components/ManageAccountsDrawer/components/ImportAccounts.tsx

@@ -317,6 +317,11 @@ export const ImportAccounts = ({ chainId, onClose }: ImportAccountsProps) => {
       return
     }
 
+    if (!walletDeviceId) {
+      console.error('Missing walletDeviceId')
+      return
+    }
+
     setIsSubmitting(true)
 
     // For every new account that is active, fetch the account and upsert it into the redux state

src/components/MultiHopTrade/hooks/useReceiveAddress.tsx

@@ -66,7 +66,7 @@ export const useReceiveAddress = ({
       ? skipToken
       : async () => {
           // Already partially covered in isInitializing, but TypeScript lyfe mang.
-          if (!buyAsset || !wallet || !buyAccountId || !buyAccountMetadata) {
+          if (!buyAsset || !wallet || !buyAccountId || !buyAccountMetadata || !deviceId) {
             return undefined
           }
 

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

@@ -38,7 +38,7 @@ export const KeepKeyPin = ({
     },
     dispatch,
   } = useWallet()
-  const wallet = keyring.get(deviceId)
+  const wallet = keyring.get(deviceId ?? '')
 
   const pinFieldRef = useRef<HTMLInputElement | null>(null)
 
@@ -138,10 +138,10 @@ export const KeepKeyPin = ({
       }
     }
 
-    keyring.on(['KeepKey', deviceId, String(MessageType.FAILURE)], handleError)
+    keyring.on(['KeepKey', deviceId ?? '', String(MessageType.FAILURE)], handleError)
 
     return () => {
-      keyring.off(['KeepKey', deviceId, String(MessageType.FAILURE)], handleError)
+      keyring.off(['KeepKey', deviceId ?? '', String(MessageType.FAILURE)], handleError)
     }
   }, [deviceId, keyring])
 

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

@@ -52,7 +52,8 @@ 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 +68,7 @@ export const LedgerChains = () => {
         ]({
           accountNumber: 0,
           chainIds,
-          wallet: walletState.wallet,
+          wallet,
           isSnapInstalled: false,
         })
 
@@ -102,7 +103,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 +121,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(() => {

src/context/WalletProvider/WalletViewsSwitch.tsx

@@ -15,8 +15,6 @@ import { Route, Switch, useHistory, useLocation, useRouteMatch } from 'react-rou
 import { SlideTransition } from 'components/SlideTransition'
 import { WalletActions } from 'context/WalletProvider/actions'
 import { useWallet } from 'hooks/useWallet/useWallet'
-import { localWalletSlice } from 'state/slices/localWalletSlice/localWalletSlice'
-import { store } from 'state/store'
 
 import { SUPPORTED_WALLETS } from './config'
 import { KeyManager } from './KeyManager'
@@ -63,14 +61,11 @@ export const WalletViewsSwitch = () => {
     if (disposition === 'initializing' || disposition === 'recovering') {
       await wallet?.cancel()
       disconnect()
-      store.dispatch(localWalletSlice.actions.clearLocalWallet())
       dispatch({ type: WalletActions.OPEN_KEEPKEY_DISCONNECT })
     } else {
       history.replace(INITIAL_WALLET_MODAL_ROUTE)
       if (disconnectOnCloseModal) {
         disconnect()
-        dispatch({ type: WalletActions.RESET_STATE })
-        store.dispatch(localWalletSlice.actions.clearLocalWallet())
       } else {
         dispatch({ type: WalletActions.SET_WALLET_MODAL, payload: false })
       }

src/context/WalletProvider/actions.ts

@@ -28,6 +28,8 @@ export enum WalletActions {
   OPEN_KEEPKEY_RECOVERY = 'OPEN_KEEPKEY_RECOVERY',
   OPEN_KEEPKEY_CHARACTER_REQUEST = 'OPEN_KEEPKEY_CHARACTER_REQUEST',
   DOWNLOAD_UPDATER = 'DOWNLOAD_UPDATER',
+  SET_NATIVE_PENDING_DEVICE_ID = 'SET_NATIVE_PENDING_DEVICE_ID',
+  RESET_NATIVE_PENDING_DEVICE_ID = 'RESET_NATIVE_PENDING_DEVICE_ID',
 }
 
 export type ActionTypes =
@@ -62,7 +64,6 @@ export type ActionTypes =
       type: WalletActions.NATIVE_PASSWORD_OPEN
       payload: {
         modal: boolean
-        deviceId: string
       }
     }
   | {
@@ -107,3 +108,8 @@ export type ActionTypes =
       }
     }
   | { type: WalletActions.OPEN_KEEPKEY_DISCONNECT }
+  | {
+      type: WalletActions.SET_NATIVE_PENDING_DEVICE_ID
+      payload: string | null
+    }
+  | { type: WalletActions.RESET_NATIVE_PENDING_DEVICE_ID }