check multiple derivation paths when restoring (#792)

* check multiple derivation paths when restoring * non null assertion
helium · Aug 15, 2024 · bd86027 · bd86027
1 parent 97e4e9c
commit bd86027
Show file tree

Hide file tree

Showing 2 changed files with 69 additions and 27 deletions.
diff --git a/src/features/onboarding/import/AccountImportScreen.tsx b/src/features/onboarding/import/AccountImportScreen.tsx
@@ -11,12 +11,11 @@ import React, { useCallback, useEffect, useRef, useState } from 'react'
 import { useTranslation } from 'react-i18next'
 import { FlatList } from 'react-native'
 import { KeyboardAwareScrollView } from 'react-native-keyboard-aware-scroll-view'
+import { MAIN_DERIVATION_PATHS } from '@hooks/useDerivationAccounts'
+import { Keypair } from '@solana/web3.js'
+import { CSAccount } from '@storage/cloudStorage'
 import { useAccountStorage } from '../../../storage/AccountStorageProvider'
-import {
-  DEFAULT_DERIVATION_PATH,
-  createKeypair,
-  toSecureAccount,
-} from '../../../storage/secureStorage'
+import { createKeypair, toSecureAccount } from '../../../storage/secureStorage'
 import { useOnboarding } from '../OnboardingProvider'
 import { OnboardingNavigationProp } from '../onboardingTypes'
 import PassphraseAutocomplete from './PassphraseAutocomplete'
@@ -149,27 +148,68 @@ const AccountImportScreen = () => {
 
   const handleNext = useCallback(async () => {
     try {
+      let keypair: Keypair | undefined
       const filteredWords: string[] = words.flatMap((w) => (w ? [w] : []))
-      const { keypair } = await createKeypair({
-        givenMnemonic: filteredWords,
-        use24Words: words?.length === 24,
-        derivationPath:
-          Object.values(accounts || {}).find(
-            (a) => a.address === accountAddress,
-          )?.derivationPath || DEFAULT_DERIVATION_PATH,
-      })
+      const foundDerivation = Object.values(accounts || {}).find(
+        (a) => a.address === accountAddress,
+      )?.derivationPath
+
+      if (foundDerivation) {
+        keypair = (
+          await createKeypair({
+            givenMnemonic: filteredWords,
+            use24Words: words?.length === 24,
+            derivationPath: foundDerivation,
+          })
+        ).keypair
+      }
+
       if (restoringAccount) {
+        let restoredAccount: CSAccount | undefined
         if (!accounts || !accountAddress) {
           await showOKAlert({
             title: t('restoreAccount.errorAlert.title'),
             message: t('restoreAccount.errorAlert.message'),
           })
           return
         }
-        const restoredAccount = Object.values(accounts).find(
-          (a) => a.solanaAddress === keypair.publicKey.toBase58(),
-        )
-        if (!restoredAccount || accountAddress !== restoredAccount.address) {
+
+        if (keypair) {
+          restoredAccount = Object.values(accounts).find(
+            // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+            (a) => a.solanaAddress === keypair!.publicKey.toBase58(),
+          )
+        } else {
+          const keypairs = await Promise.all(
+            MAIN_DERIVATION_PATHS.map((dpath) =>
+              createKeypair({
+                givenMnemonic: filteredWords,
+                use24Words: words?.length === 24,
+                derivationPath: dpath,
+              }),
+            ),
+          )
+
+          restoredAccount = Object.values(accounts).find((a) =>
+            keypairs.some(
+              (k) => a.solanaAddress === k.keypair.publicKey.toBase58(),
+            ),
+          )
+
+          if (restoredAccount) {
+            keypair = keypairs.find(
+              (k) =>
+                restoredAccount?.solanaAddress ===
+                k.keypair.publicKey.toBase58(),
+            )?.keypair
+          }
+        }
+
+        if (
+          !keypair ||
+          !restoredAccount ||
+          accountAddress !== restoredAccount.address
+        ) {
           await showOKAlert({
             title: t('restoreAccount.errorAlert.title'),
             message: t('restoreAccount.errorAlert.message'),
@@ -183,10 +223,9 @@ const AccountImportScreen = () => {
         })
         reset()
         navigation.popToTop()
-      } else {
-        setOnboardingData((prev) => ({ ...prev, words: filteredWords }))
-        navigation.navigate('ImportSubAccounts')
       }
+      setOnboardingData((prev) => ({ ...prev, words: filteredWords }))
+      navigation.navigate('ImportSubAccounts')
     } catch (error) {
       await showOKAlert({
         title: t('accountImport.alert.title'),

diff --git a/src/hooks/useDerivationAccounts.ts b/src/hooks/useDerivationAccounts.ts
@@ -65,6 +65,11 @@ export type ResolvedPath = {
 }
 
 export const HELIUM_DERIVATION = 'Helium L1'
+export const MAIN_DERIVATION_PATHS = [
+  HELIUM_DERIVATION,
+  heliumDerivation(-1),
+  solanaDerivation(-1, undefined),
+]
 export const useDerivationAccounts = ({ mnemonic }: { mnemonic?: string }) => {
   const { connection } = useSolana()
   const [resolvedGroups, setResolvedGroups] = useState<ResolvedPath[][]>([])
@@ -75,12 +80,6 @@ export const useDerivationAccounts = ({ mnemonic }: { mnemonic?: string }) => {
     [resolvedGroups],
   )
 
-  const mains = [
-    HELIUM_DERIVATION,
-    heliumDerivation(-1),
-    solanaDerivation(-1, undefined),
-  ]
-
   const solanaWithChange = (start: number, end: number) =>
     new Array(end - start).fill(0).map((_, i) => solanaDerivation(i + start, 0))
 
@@ -90,7 +89,11 @@ export const useDerivationAccounts = ({ mnemonic }: { mnemonic?: string }) => {
       .map((_, i) => solanaDerivation(i + start, undefined))
 
   const [groups, setGroups] = useState([
-    [...mains, ...solanaWithChange(0, 10), ...solanaWithoutChange(0, 10)],
+    [
+      ...MAIN_DERIVATION_PATHS,
+      ...solanaWithChange(0, 10),
+      ...solanaWithoutChange(0, 10),
+    ],
   ])
 
   // When mnemonic changes, reset resolved groups