Fix: CP-9487 && CP-9637 accounts and wallets changes (#111)

Co-authored-by: Gergely Lovas <[email protected]>
ava-labs · Jan 8, 2025 · e390a94 · e390a94
1 parent 8cce5d2
commit e390a94
Show file tree

Hide file tree

Showing 15 changed files with 162 additions and 36 deletions.
diff --git a/src/background/services/onboarding/handlers/ledgerOnboardingHandler.test.ts b/src/background/services/onboarding/handlers/ledgerOnboardingHandler.test.ts
@@ -146,14 +146,15 @@ 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([
       {
         pubKeys: ['pubkey1', 'pubkey2', 'pubkey3'],
         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
@@ -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
@@ -14,6 +14,7 @@ describe('src/background/services/wallet/handlers/importLedger', () => {
 
   const accountsService = {
     addPrimaryAccount: jest.fn(),
+    activateAccount: jest.fn(),
   } as unknown as jest.Mocked<AccountsService>;
 
   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
@@ -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
@@ -17,6 +17,7 @@ describe('src/background/services/wallet/handlers/importSeedPhrase', () => {
 
   const accountsService = {
     addPrimaryAccount: jest.fn(),
+    activateAccount: jest.fn(),
   } as unknown as jest.Mocked<AccountsService>;
 
   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
@@ -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
@@ -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
@@ -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<AddressType[]>([]);
   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,19 +172,20 @@ export function LedgerConnector({
         publicKeys: undefined,
         hasPublicKeys: true,
         pathSpec: DerivationPath.BIP44,
+        lastAccountIndexWithBalance: lastAccountIndexWithBalance.current,
       });
     } catch {
       capture('OnboardingLedgerConnectionFailed');
       setPublicKeyState(LedgerStatus.LEDGER_CONNECTION_FAILED);
       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,
     ],
   );