diff --git a/src/background/services/onboarding/handlers/ledgerOnboardingHandler.test.ts b/src/background/services/onboarding/handlers/ledgerOnboardingHandler.test.ts index df1368e1..4e3a1692 100644 --- a/src/background/services/onboarding/handlers/ledgerOnboardingHandler.test.ts +++ b/src/background/services/onboarding/handlers/ledgerOnboardingHandler.test.ts @@ -146,7 +146,7 @@ describe('src/background/services/onboarding/handlers/ledgerOnboardingHandler.ts ).not.toHaveBeenCalled(); }); - it('sets up a ledger wallet with pubkeys correctly', async () => { + it('sets up a ledger wallet with pubkeys correctly with right amount of accounts', async () => { const handler = getHandler(); const request = getRequest([ { @@ -154,6 +154,7 @@ describe('src/background/services/onboarding/handlers/ledgerOnboardingHandler.ts password: 'password', walletName: 'wallet-name', analyticsConsent: false, + numberOfAccountsToCreate: 2, }, ]); @@ -181,9 +182,12 @@ describe('src/background/services/onboarding/handlers/ledgerOnboardingHandler.ts expect(accountsServiceMock.addPrimaryAccount).toHaveBeenNthCalledWith(2, { walletId: WALLET_ID, }); - expect(accountsServiceMock.addPrimaryAccount).toHaveBeenNthCalledWith(3, { - walletId: WALLET_ID, - }); + expect(accountsServiceMock.addPrimaryAccount).not.toHaveBeenNthCalledWith( + 3, + { + walletId: WALLET_ID, + }, + ); expect(settingsServiceMock.setAnalyticsConsent).toHaveBeenCalledWith(false); diff --git a/src/background/services/onboarding/handlers/ledgerOnboardingHandler.ts b/src/background/services/onboarding/handlers/ledgerOnboardingHandler.ts index 04dcbfee..f8b7f563 100644 --- a/src/background/services/onboarding/handlers/ledgerOnboardingHandler.ts +++ b/src/background/services/onboarding/handlers/ledgerOnboardingHandler.ts @@ -26,6 +26,7 @@ type HandlerType = ExtensionRequestHandler< password: string; analyticsConsent: boolean; walletName?: string; + numberOfAccountsToCreate?: number; }, ] >; @@ -46,8 +47,15 @@ export class LedgerOnboardingHandler implements HandlerType { ) {} handle: HandlerType['handle'] = async ({ request }) => { - const { xpub, xpubXP, pubKeys, password, analyticsConsent, walletName } = - (request.params ?? [])[0] ?? {}; + const { + xpub, + xpubXP, + pubKeys, + password, + analyticsConsent, + walletName, + numberOfAccountsToCreate, + } = (request.params ?? [])[0] ?? {}; if ((xpub || xpubXP) && pubKeys?.length) { return { @@ -92,20 +100,15 @@ export class LedgerOnboardingHandler implements HandlerType { }; } - if (xpub) { + for (let i = 0; i < (numberOfAccountsToCreate || 1); i++) { + if (pubKeys && pubKeys.length < i) { + break; + } await this.accountsService.addPrimaryAccount({ walletId, }); } - if (pubKeys?.length) { - for (let i = 0; i < pubKeys.length; i++) { - await this.accountsService.addPrimaryAccount({ - walletId, - }); - } - } - await finalizeOnboarding({ walletId, networkService: this.networkService, diff --git a/src/background/services/wallet/handlers/importLedger.test.ts b/src/background/services/wallet/handlers/importLedger.test.ts index 2ad3cd83..51b1ae92 100644 --- a/src/background/services/wallet/handlers/importLedger.test.ts +++ b/src/background/services/wallet/handlers/importLedger.test.ts @@ -14,6 +14,7 @@ describe('src/background/services/wallet/handlers/importLedger', () => { const accountsService = { addPrimaryAccount: jest.fn(), + activateAccount: jest.fn(), } as unknown as jest.Mocked; const secretsService = { @@ -105,6 +106,7 @@ describe('src/background/services/wallet/handlers/importLedger', () => { xpub: xpubValue, xpubXP: xpubXPValue, name: nameValue, + numberOfAccountsToCreate: 2, }); expect(walletService.addPrimaryWallet).toHaveBeenCalledWith({ @@ -115,13 +117,18 @@ describe('src/background/services/wallet/handlers/importLedger', () => { name: nameValue, }); + expect(accountsService.addPrimaryAccount).toHaveBeenCalledTimes(2); + + expect(accountsService.activateAccount).toHaveBeenCalledTimes(1); + expect(result).toEqual({ type: SecretType.Ledger, name: nameValue, id: walletId, }); }); - it('returns an ImportWalletResult if LedgerLive import is successful', async () => { + + it('only imports accounts with pubkeys', async () => { const walletId = crypto.randomUUID(); const pubKeysValue = [ { @@ -143,6 +150,7 @@ describe('src/background/services/wallet/handlers/importLedger', () => { secretType: SecretType.LedgerLive, pubKeys: pubKeysValue, name: nameValue, + numberOfAccountsToCreate: 2, }); expect(walletService.addPrimaryWallet).toHaveBeenCalledWith({ @@ -152,6 +160,59 @@ describe('src/background/services/wallet/handlers/importLedger', () => { name: nameValue, }); + expect(accountsService.addPrimaryAccount).toHaveBeenCalledTimes(1); + expect(accountsService.activateAccount).toHaveBeenCalledTimes(1); + + expect(result).toEqual({ + type: SecretType.LedgerLive, + name: nameValue, + id: walletId, + }); + }); + + it('imports max 3 accounts', async () => { + const walletId = crypto.randomUUID(); + const pubKeysValue = [ + { + evm: 'pubKeyEvm1', + }, + { + evm: 'pubKeyEvm2', + }, + { + evm: 'pubKeyEvm3', + }, + { + evm: 'pubKeyEvm4', + }, + ]; + const nameValue = 'walletName'; + secretsService.isKnownSecret.mockResolvedValueOnce(false); + walletService.addPrimaryWallet.mockResolvedValue(walletId); + secretsService.getWalletAccountsSecretsById.mockResolvedValue({ + secretType: SecretType.LedgerLive, + pubKeys: pubKeysValue, + derivationPath: DerivationPath.LedgerLive, + id: walletId, + name: nameValue, + }); + + const { result } = await handle({ + secretType: SecretType.LedgerLive, + pubKeys: pubKeysValue, + name: nameValue, + numberOfAccountsToCreate: 4, + }); + + expect(walletService.addPrimaryWallet).toHaveBeenCalledWith({ + secretType: SecretType.LedgerLive, + pubKeys: pubKeysValue, + derivationPath: DerivationPath.LedgerLive, + name: nameValue, + }); + + expect(accountsService.addPrimaryAccount).toHaveBeenCalledTimes(3); + expect(result).toEqual({ type: SecretType.LedgerLive, name: nameValue, diff --git a/src/background/services/wallet/handlers/importLedger.ts b/src/background/services/wallet/handlers/importLedger.ts index 3aebaee4..b39baa27 100644 --- a/src/background/services/wallet/handlers/importLedger.ts +++ b/src/background/services/wallet/handlers/importLedger.ts @@ -32,15 +32,27 @@ export class ImportLedgerHandler implements HandlerType { // (i.e. address for 0-index account derived N times). for (let i = 0; i < numberOfAccounts; i++) { - await this.accountsService.addPrimaryAccount({ + const accountId = await this.accountsService.addPrimaryAccount({ walletId, }); + if (i === 0) { + await this.accountsService.activateAccount(accountId); + } } } handle: HandlerType['handle'] = async ({ request }) => { - const [{ xpub, xpubXP, pubKeys, secretType, name, dryRun }] = - request.params; + const [ + { + xpub, + xpubXP, + pubKeys, + secretType, + name, + dryRun, + numberOfAccountsToCreate, + }, + ] = request.params; if ( secretType !== SecretType.Ledger && @@ -105,10 +117,15 @@ export class ImportLedgerHandler implements HandlerType { }); } - const numberOfAccountsToCreate = - secretType === SecretType.LedgerLive ? pubKeys?.length : undefined; + const accountsToBeCreated = numberOfAccountsToCreate || 3; + const accountsToCreate = Math.min( + 3, + pubKeys + ? Math.min(pubKeys.length, accountsToBeCreated) + : accountsToBeCreated, + ); + await this.#addAccounts(id, accountsToCreate); - await this.#addAccounts(id, numberOfAccountsToCreate); const addedWallet = await this.secretsService.getWalletAccountsSecretsById(id); return { diff --git a/src/background/services/wallet/handlers/importSeedPhrase.test.ts b/src/background/services/wallet/handlers/importSeedPhrase.test.ts index 783fc031..212a93a6 100644 --- a/src/background/services/wallet/handlers/importSeedPhrase.test.ts +++ b/src/background/services/wallet/handlers/importSeedPhrase.test.ts @@ -17,6 +17,7 @@ describe('src/background/services/wallet/handlers/importSeedPhrase', () => { const accountsService = { addPrimaryAccount: jest.fn(), + activateAccount: jest.fn(), } as unknown as jest.Mocked; const secretsService = { @@ -135,6 +136,8 @@ describe('src/background/services/wallet/handlers/importSeedPhrase', () => { walletId, }); + expect(accountsService.activateAccount).toHaveBeenCalledTimes(1); + expect(result).toEqual({ type: SecretType.Mnemonic, name, diff --git a/src/background/services/wallet/handlers/importSeedPhrase.ts b/src/background/services/wallet/handlers/importSeedPhrase.ts index 9f687df7..ad34416f 100644 --- a/src/background/services/wallet/handlers/importSeedPhrase.ts +++ b/src/background/services/wallet/handlers/importSeedPhrase.ts @@ -39,7 +39,7 @@ export class ImportSeedPhraseHandler implements HandlerType { async #addAccounts(walletId: string) { // We need to await each of these calls, otherwise there may be race conditions // (i.e. address for 0-index account derived N times). - await this.accountsService.addPrimaryAccount({ + const activeAccount = await this.accountsService.addPrimaryAccount({ walletId, }); await this.accountsService.addPrimaryAccount({ @@ -48,6 +48,8 @@ export class ImportSeedPhraseHandler implements HandlerType { await this.accountsService.addPrimaryAccount({ walletId, }); + + await this.accountsService.activateAccount(activeAccount); } handle: HandlerType['handle'] = async ({ request }) => { diff --git a/src/background/services/wallet/handlers/models.ts b/src/background/services/wallet/handlers/models.ts index 843a3f3b..e6c9610c 100644 --- a/src/background/services/wallet/handlers/models.ts +++ b/src/background/services/wallet/handlers/models.ts @@ -24,6 +24,7 @@ export type ImportLedgerWalletParams = { secretType: SecretType.Ledger | SecretType.LedgerLive; name?: string; dryRun?: boolean; + numberOfAccountsToCreate?: number; }; export type ImportWalletResult = { diff --git a/src/components/ledger/LedgerConnector.tsx b/src/components/ledger/LedgerConnector.tsx index 4dc24ab3..3a426d1a 100644 --- a/src/components/ledger/LedgerConnector.tsx +++ b/src/components/ledger/LedgerConnector.tsx @@ -1,4 +1,4 @@ -import { useCallback, useEffect, useState } from 'react'; +import { useCallback, useEffect, useRef, useState } from 'react'; import { useLedgerContext } from '@src/contexts/LedgerProvider'; import { useAnalyticsContext } from '@src/contexts/AnalyticsProvider'; import { @@ -45,6 +45,7 @@ export interface LedgerConnectorData { publicKeys: PubKeyType[] | undefined; hasPublicKeys: boolean; pathSpec: DerivationPath; + lastAccountIndexWithBalance: number; } interface LedgerConnectorProps { @@ -81,6 +82,8 @@ export function LedgerConnector({ const [addresses, setAddresses] = useState([]); const [hasPublicKeys, setHasPublicKeys] = useState(false); const [dropdownDisabled, setDropdownDisabled] = useState(true); + const lastAccountIndexWithBalance = useRef(0); + const { t } = useTranslation(); const { importLedger } = useImportLedger(); @@ -98,12 +101,18 @@ export function LedgerConnector({ addressList: AddressType[] = [], ) => { const address = getAddressFromXPub(xpubValue, accountIndex); + const { balance } = await getAvaxBalance(address); + const newAddresses = [ ...addressList, { address, balance: balance.balanceDisplayValue || '0' }, ]; setAddresses(newAddresses); + lastAccountIndexWithBalance.current = Math.max( + 0, + newAddresses.findLastIndex((addr) => addr.balance !== '0'), + ); if (accountIndex < 2) { await getAddressFromXpubKey(xpubValue, accountIndex + 1, newAddresses); } @@ -163,6 +172,7 @@ export function LedgerConnector({ publicKeys: undefined, hasPublicKeys: true, pathSpec: DerivationPath.BIP44, + lastAccountIndexWithBalance: lastAccountIndexWithBalance.current, }); } catch { capture('OnboardingLedgerConnectionFailed'); @@ -170,12 +180,12 @@ export function LedgerConnector({ popDeviceSelection(); } }, [ - capture, - isLedgerWalletExist, + getExtendedPublicKey, checkIfWalletExists, + capture, getAddressFromXpubKey, - getExtendedPublicKey, onSuccess, + isLedgerWalletExist, popDeviceSelection, ]); @@ -214,6 +224,10 @@ export function LedgerConnector({ { address, balance: balance.balanceDisplayValue || '0' }, ]; setAddresses(newAddresses); + lastAccountIndexWithBalance.current = Math.max( + 0, + newAddresses.findLastIndex((addr) => addr.balance !== '0'), + ); if (accountIndex < 2) { await getPubKeys(derivationPathSpec, accountIndex + 1, newAddresses, [ ...pubKeys, @@ -242,6 +256,7 @@ export function LedgerConnector({ publicKeys: publicKeyValue, hasPublicKeys: true, pathSpec: DerivationPath.LedgerLive, + lastAccountIndexWithBalance: lastAccountIndexWithBalance.current, }); } } catch { @@ -251,12 +266,12 @@ export function LedgerConnector({ } }, [ + getPublicKey, + getAvaxBalance, capture, - isLedgerWalletExist, checkIfWalletExists, - getAvaxBalance, - getPublicKey, onSuccess, + isLedgerWalletExist, popDeviceSelection, ], ); diff --git a/src/contexts/OnboardingProvider.tsx b/src/contexts/OnboardingProvider.tsx index 42700b36..d057478b 100644 --- a/src/contexts/OnboardingProvider.tsx +++ b/src/contexts/OnboardingProvider.tsx @@ -82,6 +82,7 @@ const OnboardingContext = createContext<{ isNewsletterEnabled: boolean; setIsNewsletterEnabled: Dispatch>; onboardingWalletType: WalletType | undefined; + setNumberOfAccountsToCreate: Dispatch>; }>({} as any); export function OnboardingContextProvider({ children }: { children: any }) { @@ -131,6 +132,7 @@ export function OnboardingContextProvider({ children }: { children: any }) { const [walletType, setWalletType] = useState(); const [isSeedlessMfaRequired, setIsSeedlessMfaRequired] = useState(false); + const [numberOfAccountsToCreate, setNumberOfAccountsToCreate] = useState(0); const [onboardingWalletType, setOnboardingWalletType] = useState< WalletType | undefined @@ -155,6 +157,7 @@ export function OnboardingContextProvider({ children }: { children: any }) { setOnboardingWalletType(undefined); setIsNewsletterEnabled(false); setNewsletterEmail(''); + setNumberOfAccountsToCreate(0); }, []); useEffect(() => { @@ -272,11 +275,13 @@ export function OnboardingContextProvider({ children }: { children: any }) { password, analyticsConsent: !!analyticsConsent, walletName: walletName, + numberOfAccountsToCreate, }, ], }); }, [ analyticsConsent, + numberOfAccountsToCreate, password, publicKeys, request, @@ -436,6 +441,7 @@ export function OnboardingContextProvider({ children }: { children: any }) { setIsSeedlessMfaRequired, setOnboardingWalletType, onboardingWalletType, + setNumberOfAccountsToCreate, }} > {/* diff --git a/src/pages/Accounts/AddWalletWithLedger.tsx b/src/pages/Accounts/AddWalletWithLedger.tsx index a5c7d07c..72b84d61 100644 --- a/src/pages/Accounts/AddWalletWithLedger.tsx +++ b/src/pages/Accounts/AddWalletWithLedger.tsx @@ -11,7 +11,7 @@ import { } from '@avalabs/core-k2-components'; import { PageTitle } from '@src/components/common/PageTitle'; import { useTranslation } from 'react-i18next'; -import { useCallback, useEffect, useState } from 'react'; +import { useCallback, useEffect, useRef, useState } from 'react'; import { LedgerWrongVersionOverlay } from '../Ledger/LedgerWrongVersionOverlay'; import { PubKeyType } from '@src/background/services/wallet/models'; @@ -59,6 +59,7 @@ export function AddWalletWithLedger() { const [pathSpec, setPathSpec] = useState( DerivationPath.BIP44, ); + const lastAccountIndexWithBalance = useRef(0); const { popDeviceSelection } = useLedgerContext(); @@ -74,6 +75,7 @@ export function AddWalletWithLedger() { setPublicKeys(data.publicKeys); setHasPublicKeys(data.hasPublicKeys); setPathSpec(data.pathSpec); + lastAccountIndexWithBalance.current = data.lastAccountIndexWithBalance; } const LedgerLiveSupportButton = () => ( @@ -114,6 +116,7 @@ export function AddWalletWithLedger() { pathSpec === DerivationPath.BIP44 ? SecretType.Ledger : SecretType.LedgerLive, + numberOfAccountsToCreate: lastAccountIndexWithBalance.current + 1, }); capture('SeedphraseImportSuccess'); @@ -132,6 +135,7 @@ export function AddWalletWithLedger() { capture, getErrorMessage, importLedger, + lastAccountIndexWithBalance, pathSpec, publicKeys, xpub, diff --git a/src/pages/Accounts/hooks/useKeystoreFileImport.test.ts b/src/pages/Accounts/hooks/useKeystoreFileImport.test.ts index 89a35eca..91f4986c 100644 --- a/src/pages/Accounts/hooks/useKeystoreFileImport.test.ts +++ b/src/pages/Accounts/hooks/useKeystoreFileImport.test.ts @@ -12,10 +12,12 @@ import { usePrivateKeyImport } from './usePrivateKeyImport'; import { useKeystoreFileImport } from './useKeystoreFileImport'; import { SeedphraseImportError } from '@src/background/services/wallet/handlers/models'; import { utils } from '@avalabs/avalanchejs'; +import { useAccountsContext } from '@src/contexts/AccountsProvider'; jest.mock('@src/contexts/AnalyticsProvider'); jest.mock('./useImportSeedphrase'); jest.mock('./usePrivateKeyImport'); +jest.mock('@src/contexts/AccountsProvider'); const getFile = (data) => { const encoder = new TextEncoder(); @@ -42,6 +44,10 @@ describe('src/pages/Accounts/hooks/useKeystoreFileImport', () => { isImporting: false, importSeedphrase: jest.fn(), }); + + jest.mocked(useAccountsContext).mockReturnValue({ + selectAccount: jest.fn(), + } as any); }); describe('isValidKeystoreFile()', () => { diff --git a/src/pages/Accounts/hooks/useKeystoreFileImport.ts b/src/pages/Accounts/hooks/useKeystoreFileImport.ts index 23ecd44e..3b81bc2a 100644 --- a/src/pages/Accounts/hooks/useKeystoreFileImport.ts +++ b/src/pages/Accounts/hooks/useKeystoreFileImport.ts @@ -17,6 +17,7 @@ import { useAnalyticsContext } from '@src/contexts/AnalyticsProvider'; import { useImportSeedphrase } from './useImportSeedphrase'; import { usePrivateKeyImport } from './usePrivateKeyImport'; import { useJsonFileReader } from './useJsonFileReader'; +import { useAccountsContext } from '@src/contexts/AccountsProvider'; export const useKeystoreFileImport = () => { const { capture } = useAnalyticsContext(); @@ -26,6 +27,7 @@ export const useKeystoreFileImport = () => { useImportSeedphrase(); const { isImporting: isImportingPrivateKey, importPrivateKey } = usePrivateKeyImport(); + const { selectAccount } = useAccountsContext(); const extractKeys = useCallback( async (file: File, password: string) => { @@ -62,7 +64,8 @@ export const useKeystoreFileImport = () => { utils.base58check.decode(key.replace('PrivateKey-', '')), ).toString('hex'); - await importPrivateKey(privateKey); + const accountId = await importPrivateKey(privateKey); + await selectAccount(accountId); } else if (type === 'mnemonic') { try { await importSeedphrase({ @@ -82,7 +85,7 @@ export const useKeystoreFileImport = () => { } } }, - [extractKeys, importPrivateKey, importSeedphrase], + [extractKeys, importPrivateKey, importSeedphrase, selectAccount], ); const getKeyCounts = useCallback( diff --git a/src/pages/ImportPrivateKey/ImportPrivateKey.tsx b/src/pages/ImportPrivateKey/ImportPrivateKey.tsx index 3b554f54..993daeb5 100644 --- a/src/pages/ImportPrivateKey/ImportPrivateKey.tsx +++ b/src/pages/ImportPrivateKey/ImportPrivateKey.tsx @@ -51,9 +51,8 @@ export function ImportPrivateKey() { const balance = useBalanceTotalInCurrency(derivedAddresses as Account); const { isImporting: isImportLoading, importPrivateKey } = usePrivateKeyImport(); - const { selectAccount } = useAccountsContext(); const history = useHistory(); - const { allAccounts } = useAccountsContext(); + const { allAccounts, selectAccount } = useAccountsContext(); const [isKnownAccount, setIsKnownAccount] = useState(false); const [isDuplicatedAccountDialogOpen, setIsDuplicatedAccountDialogOpen] = useState(false); diff --git a/src/pages/Onboarding/pages/Ledger/LedgerConnect.tsx b/src/pages/Onboarding/pages/Ledger/LedgerConnect.tsx index eb547d95..47633f3b 100644 --- a/src/pages/Onboarding/pages/Ledger/LedgerConnect.tsx +++ b/src/pages/Onboarding/pages/Ledger/LedgerConnect.tsx @@ -52,6 +52,7 @@ export function LedgerConnect() { setPublicKeys, setOnboardingPhase, setOnboardingWalletType, + setNumberOfAccountsToCreate, } = useOnboardingContext(); const [hasPublicKeys, setHasPublicKeys] = useState(false); @@ -69,6 +70,7 @@ export function LedgerConnect() { setXpubXP(data.xpubXP); setPublicKeys(data.publicKeys); setHasPublicKeys(data.hasPublicKeys); + setNumberOfAccountsToCreate(data.lastAccountIndexWithBalance + 1); } const Content = ( diff --git a/tsconfig.json b/tsconfig.json index d1476f19..02e91865 100644 --- a/tsconfig.json +++ b/tsconfig.json @@ -13,7 +13,7 @@ "sourceMap": true, "jsx": "react-jsx", "esModuleInterop": true, - "lib": ["es2015", "es2021", "dom"], + "lib": ["es2015", "es2021", "es2023", "dom"], "experimentalDecorators": true, "emitDecoratorMetadata": true, "allowSyntheticDefaultImports": true,