fix: restore snap and imported srps when changing password

app/core/SnapKeyring/MultichainWalletSnapClient.test.ts

@@ -6,17 +6,21 @@ import {
   WalletClientType,
 } from './MultichainWalletSnapClient';
 import { CaipChainId, SnapId } from '@metamask/snaps-sdk';
-import { KeyringTypes } from '@metamask/keyring-controller';
 import Engine from '../Engine';
 import { Sender } from '@metamask/keyring-snap-client';
 import { SnapKeyring } from '@metamask/eth-snap-keyring';
 import { BtcScope, SolScope } from '@metamask/keyring-api';
 
+const mockSnapKeyring = {
+  createAccount: jest.fn(),
+  discoverAccounts: jest.fn(),
+};
+
 jest.mock('../Engine', () => ({
   controllerMessenger: {
     call: jest.fn(),
   },
-  getSnapKeyring: jest.fn(),
+  getSnapKeyring: () => mockSnapKeyring,
   context: {
     AccountsController: {
       getNextAvailableAccountName: jest.fn().mockReturnValue('Snap Account 1'),
@@ -94,32 +98,29 @@ describe('MultichainWalletSnapClient', () => {
 
   describe('withSnapKeyring', () => {
     it('calls the callback with the keyring', async () => {
-      const mockCallback = jest.fn();
-      const mockKeyring = { test: 'keyring' };
-
-      (Engine.controllerMessenger.call as jest.Mock).mockImplementationOnce(
-        async (_, __, callback) => {
-          await callback({ keyring: mockKeyring });
-        },
-      );
+      const mockOptions = {};
+      const mockCallback = async (keyring: SnapKeyring) => {
+        await keyring.createAccount(
+          mockSnapId,
+          mockOptions,
+          mockSnapKeyringOptions,
+        );
+      };
 
       await client.testWithSnapKeyring(mockCallback);
 
-      expect(Engine.controllerMessenger.call).toHaveBeenCalledWith(
-        'KeyringController:withKeyring',
-        { type: KeyringTypes.snap },
-        expect.any(Function),
+      expect(mockSnapKeyring.createAccount).toHaveBeenCalledWith(
+        mockSnapId,
+        mockOptions,
+        mockSnapKeyringOptions,
       );
-      expect(mockCallback).toHaveBeenCalledWith(mockKeyring);
     });
 
     it('handles errors from the controller messenger', async () => {
       const mockError = new Error('Test error');
-      (Engine.controllerMessenger.call as jest.Mock).mockRejectedValueOnce(
-        mockError,
-      );
+      const mockCallback = jest.fn().mockRejectedValueOnce(mockError);
 
-      await expect(client.testWithSnapKeyring(jest.fn())).rejects.toThrow(
+      await expect(client.testWithSnapKeyring(mockCallback)).rejects.toThrow(
         'Test error',
       );
     });
@@ -133,19 +134,9 @@ describe('MultichainWalletSnapClient', () => {
         entropySource: 'test-entropy',
       };
 
-      const mockKeyring = {
-        createAccount: jest.fn(),
-      };
-
-      (Engine.controllerMessenger.call as jest.Mock).mockImplementationOnce(
-        async (_, __, callback) => {
-          await callback({ keyring: mockKeyring });
-        },
-      );
-
       await client.createAccount(mockOptions);
 
-      expect(mockKeyring.createAccount).toHaveBeenCalledWith(
+      expect(mockSnapKeyring.createAccount).toHaveBeenCalledWith(
         mockSnapId,
         mockOptions,
         mockSnapKeyringOptions,
@@ -212,8 +203,8 @@ describe('MultichainWalletSnapClient', () => {
 
       await client.addDiscoveredAccounts(mockEntropySource);
 
-      expect(mockKeyring.createAccount).toHaveBeenCalledTimes(1);
-      expect(mockKeyring.createAccount).toHaveBeenCalledWith(
+      expect(mockSnapKeyring.createAccount).toHaveBeenCalledTimes(1);
+      expect(mockSnapKeyring.createAccount).toHaveBeenCalledWith(
         mockSnapId,
         {
           derivationPath: mockDiscoveredAccounts[0].derivationPath,

app/core/SnapKeyring/MultichainWalletSnapClient.ts

@@ -13,7 +13,6 @@ import {
   SolanaWalletSnapSender,
 } from './SolanaWalletSnap';
 import Engine from '../Engine';
-import { KeyringTypes } from '@metamask/keyring-controller';
 import { SnapKeyring } from '@metamask/eth-snap-keyring';
 import { store } from '../../store';
 import { startPerformanceTrace } from '../redux/slices/performance';
@@ -87,16 +86,11 @@ export abstract class MultichainWalletSnapClient {
    *
    */
   protected async withSnapKeyring(
-    callback: (keyring: SnapKeyring) => Promise<void>,
+    callback: (keyring: SnapKeyring) => Promise<unknown>,
   ) {
-    const controllerMessenger = Engine.controllerMessenger;
-    await Engine.getSnapKeyring();
+    const snapKeyring = (await Engine.getSnapKeyring()) as SnapKeyring;
 
-    return await controllerMessenger.call(
-      'KeyringController:withKeyring',
-      { type: KeyringTypes.snap },
-      async ({ keyring }) => await callback(keyring as unknown as SnapKeyring),
-    );
+    return await callback(snapKeyring);
   }
 
   /**
@@ -126,7 +120,7 @@ export abstract class MultichainWalletSnapClient {
     );
 
     return await this.withSnapKeyring(async (keyring) => {
-      (keyring as unknown as SnapKeyring).createAccount(
+      await keyring.createAccount(
         this.snapId,
         {
           ...options,
@@ -211,7 +205,7 @@ export abstract class MultichainWalletSnapClient {
       await this.withSnapKeyring(async (keyring) => {
         const results = await Promise.allSettled(
           discoveredAccounts.map(async (account) => {
-            keyring.createAccount(
+            await keyring.createAccount(
               this.snapId,
               {
                 derivationPath: account.derivationPath,

app/core/Vault.js

@@ -2,6 +2,16 @@ import Engine from './Engine';
 import Logger from '../util/Logger';
 import { KeyringTypes } from '@metamask/keyring-controller';
 import { withLedgerKeyring } from './Ledger/Ledger';
+import {
+  MultichainWalletSnapFactory,
+  WalletClientType,
+} from './SnapKeyring/MultichainWalletSnapClient';
+import {
+  BtcAccountType,
+  BtcScope,
+  SolAccountType,
+  SolScope,
+} from '@metamask/keyring-api';
 
 /**
  * Restore the given serialized QR keyring.
@@ -36,14 +46,72 @@ export const restoreLedgerKeyring = async (serializedLedgerKeyring) => {
   }
 };
 
+export const restoreImportedSrp = async (seedPhrase, numberOfAccounts) => {
+  const { KeyringController } = Engine.context;
+  try {
+    const { id: keyringId } = await KeyringController.addNewKeyring(
+      KeyringTypes.hd,
+      {
+        mnemonic: seedPhrase,
+      },
+    );
+
+    for (let i = 0; i < numberOfAccounts; i++) {
+      await KeyringController.withKeyring(
+        { id: keyringId },
+        async ({ keyring }) => await keyring.addAccounts(1),
+      );
+    }
+
+    return keyringId;
+  } catch (e) {
+    Logger.error(
+      e,
+      'error while trying to restore imported srp accounts on recreate vault',
+    );
+  }
+};
+
+export const restoreSnapAccounts = async (accountType, entropySource) => {
+  let walletClientType;
+  let scope;
+  switch (accountType) {
+    case SolAccountType.DataAccount: {
+      walletClientType = WalletClientType.Solana;
+      scope = SolScope.Mainnet;
+      break;
+    }
+    case BtcAccountType.P2wpkh: {
+      walletClientType = WalletClientType.Bitcoin;
+      scope = BtcScope.Mainnet;
+      break;
+    }
+    default:
+      throw new Error('Unsupported account type');
+  }
+
+  try {
+    const client = MultichainWalletSnapFactory.createClient(walletClientType);
+    await client.createAccount({
+      entropySource,
+      scope,
+    });
+  } catch (e) {
+    Logger.error(
+      e,
+      'error while trying to restore snap accounts on recreate vault',
+    );
+  }
+};
+
 /**
  * Returns current vault seed phrase
  * It does it using an empty password or a password set by the user
  * depending on the state the app is currently in
  */
-export const getSeedPhrase = async (password = '') => {
+export const getSeedPhrase = async (password = '', keyringId) => {
   const { KeyringController } = Engine.context;
-  return await KeyringController.exportSeedPhrase(password);
+  return await KeyringController.exportSeedPhrase(password, keyringId);
 };
 
 /**
@@ -58,8 +126,35 @@ export const recreateVaultWithNewPassword = async (
   newPassword,
   selectedAddress,
 ) => {
-  const { KeyringController } = Engine.context;
-  const seedPhrase = await getSeedPhrase(password);
+  const { KeyringController, AccountsController } = Engine.context;
+  const hdKeyringsWithMetadata = KeyringController.state.keyrings
+    .map((keyring, _index) => ({
+      ...keyring,
+      ///: BEGIN:ONLY_INCLUDE_IF(multi-srp)
+      metadata: KeyringController.state.keyringsMetadata?.[_index] || {},
+      ///: END:ONLY_INCLUDE_IF
+    }))
+    .filter((keyring) => keyring.type === KeyringTypes.hd);
+
+  const seedPhrases = await Promise.all(
+    hdKeyringsWithMetadata.map(async (keyring) => {
+      try {
+        return await getSeedPhrase(password, keyring.metadata.id);
+      } catch (e) {
+        Logger.error(
+          e,
+          'error while trying to get seed phrase on recreate vault',
+        );
+        return null;
+      }
+    }),
+  );
+  const [primaryKeyringSeedPhrase, ...otherSeedPhrases] = seedPhrases;
+  if (!primaryKeyringSeedPhrase) {
+    throw new Error('error while trying to get seed phrase on recreate vault');
+  }
+
+  // START: Getting accounts to be reimported
 
   let importedAccounts = [];
   try {
@@ -83,9 +178,19 @@ export const recreateVaultWithNewPassword = async (
     );
   }
 
+  const firstPartySnapAccounts =
+    AccountsController.listMultichainAccounts().filter(
+      (account) => account.options?.entropySource,
+    );
+
   // Get props to restore vault
-  const hdKeyring = KeyringController.state.keyrings[0];
-  const existingAccountCount = hdKeyring.accounts.length;
+  const hdKeyringsAccountCount = hdKeyringsWithMetadata.map(
+    (keyring) => keyring.accounts.length,
+  );
+  const [primaryKeyringAccountCount, ...otherKeyringAccountCounts] =
+    hdKeyringsAccountCount;
+
+  // END: Getting accounts to be reimported
 
   const serializedLedgerKeyring = hasKeyringType(
     KeyringController.state,
@@ -101,7 +206,14 @@ export const recreateVaultWithNewPassword = async (
     : undefined;
 
   // Recreate keyring with password given to this method
-  await KeyringController.createNewVaultAndRestore(newPassword, seedPhrase);
+  await KeyringController.createNewVaultAndRestore(
+    newPassword,
+    primaryKeyringSeedPhrase,
+  );
+  const [newPrimaryKeyringMetadata] = KeyringController.state.keyringsMetadata;
+  const newPrimaryKeyringId = newPrimaryKeyringMetadata.id;
+
+  // START: Restoring keyrings
 
   if (serializedQrKeyring !== undefined) {
     await restoreQRKeyring(serializedQrKeyring);
@@ -111,7 +223,7 @@ export const recreateVaultWithNewPassword = async (
   }
 
   // Create previous accounts again
-  for (let i = 0; i < existingAccountCount - 1; i++) {
+  for (let i = 0; i < primaryKeyringAccountCount - 1; i++) {
     await KeyringController.addNewAccount();
   }
 
@@ -125,6 +237,36 @@ export const recreateVaultWithNewPassword = async (
   } catch (e) {
     Logger.error(e, 'error while trying to import accounts on recreate vault');
   }
+
+  // recreate import srp accounts
+  const importedSrpKeyringIds = [];
+  ///: BEGIN:ONLY_INCLUDE_IF(multi-srp)
+  for (const [index, otherSeedPhrase] of otherSeedPhrases.entries()) {
+    const importedSrpKeyring = await restoreImportedSrp(
+      otherSeedPhrase,
+      otherKeyringAccountCounts[index],
+    );
+    importedSrpKeyringIds.push(importedSrpKeyring);
+  }
+  ///: END:ONLY_INCLUDE_IF(multi-srp)
+
+  const newHdKeyringIds = [newPrimaryKeyringId, ...importedSrpKeyringIds];
+  // map old keyring id to new keyring id
+  const keyringIdMap = new Map();
+  for (const [index, keyring] of hdKeyringsWithMetadata.entries()) {
+    keyringIdMap.set(keyring.metadata.id, newHdKeyringIds[index]);
+  }
+
+  // recreate snap accounts
+  for (const snapAccount of firstPartySnapAccounts) {
+    await restoreSnapAccounts(
+      snapAccount.type,
+      keyringIdMap.get(snapAccount.options.entropySource),
+    );
+  }
+
+  // END: Restoring keyrings
+
   const recreatedKeyrings = KeyringController.state.keyrings;
   // Reselect previous selected account if still available
   for (const keyring of recreatedKeyrings) {