From f343ccf4568cf94204570e0287dfe7de1b7cb915 Mon Sep 17 00:00:00 2001 From: Monte Lai Date: Thu, 8 May 2025 22:48:04 +0800 Subject: [PATCH 01/12] fix: restore snap and imported srps --- .../SnapKeyring/MultichainWalletSnapClient.ts | 17 +- app/core/Vault.js | 157 +++++++++++++++++- 2 files changed, 154 insertions(+), 20 deletions(-) diff --git a/app/core/SnapKeyring/MultichainWalletSnapClient.ts b/app/core/SnapKeyring/MultichainWalletSnapClient.ts index 7333f7d94c22..01e86cdd4313 100644 --- a/app/core/SnapKeyring/MultichainWalletSnapClient.ts +++ b/app/core/SnapKeyring/MultichainWalletSnapClient.ts @@ -85,18 +85,11 @@ export abstract class MultichainWalletSnapClient { * */ protected async withSnapKeyring( - callback: (keyring: SnapKeyring) => Promise, + callback: (keyring: SnapKeyring) => Promise, ) { - const controllerMessenger = Engine.controllerMessenger; - await Engine.getSnapKeyring(); - - return await controllerMessenger.call( - 'KeyringController:withKeyring', - { type: KeyringTypes.snap }, - async ({ keyring }) => { - await callback(keyring as unknown as SnapKeyring); - }, - ); + const snapKeyring = (await Engine.getSnapKeyring()) as SnapKeyring; + + return await callback(snapKeyring); } /** @@ -117,7 +110,7 @@ export abstract class MultichainWalletSnapClient { ); return await this.withSnapKeyring(async (keyring) => { - (keyring as unknown as SnapKeyring).createAccount( + return await keyring.createAccount( this.snapId, options as unknown as Record, this.snapKeyringOptions, diff --git a/app/core/Vault.js b/app/core/Vault.js index 15195daee46e..a49fd790c7f2 100644 --- a/app/core/Vault.js +++ b/app/core/Vault.js @@ -2,6 +2,11 @@ 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, SolAccountType } from '@metamask/keyring-api'; /** * Restore the given serialized QR keyring. @@ -36,14 +41,72 @@ export const restoreLedgerKeyring = async (serializedLedgerKeyring) => { } }; +export const restoreImportedSrp = async (seedPhrase, numberOfAccounts) => { + const { KeyringController } = Engine.context; + Logger.log('restoreImportedSrp', seedPhrase, numberOfAccounts); + try { + const { id: keyringId } = await KeyringController.addNewKeyring( + KeyringTypes.hd, + { + mnemonic: seedPhrase, + }, + ); + + Logger.log('keyringId', keyringId); + + for (let i = 0; i < numberOfAccounts; i++) { + Logger.log('adding account', i); + await KeyringController.withKeyring( + { id: keyringId }, + async ({ keyring }) => await keyring.addAccounts(1), + ); + } + + return keyringId; + } catch (e) { + Logger.log( + e, + 'error while trying to restore imported srp accounts on recreate vault', + ); + } +}; + +export const restoreSnapAccounts = async (accountType, entropySource) => { + let walletClientType; + switch (accountType) { + case SolAccountType.DataAccount: { + walletClientType = WalletClientType.Solana; + break; + } + case BtcAccountType.P2wpkh: { + walletClientType = WalletClientType.Bitcoin; + break; + } + default: + throw new Error('Unsupported account type'); + } + + try { + const client = MultichainWalletSnapFactory.createClient(walletClientType); + await client.createAccount({ + entropySource, + }); + } 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 +121,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((keyring) => { + try { + return 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 +173,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 +201,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 +218,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 +232,40 @@ export const recreateVaultWithNewPassword = async ( } catch (e) { Logger.error(e, 'error while trying to import accounts on recreate vault'); } + + // recreate import srp accounts + const importedSrpKeyringIds = []; + for (const [index, otherSeedPhrase] of otherSeedPhrases.entries()) { + const importedSrpKeyring = await restoreImportedSrp( + otherSeedPhrase, + otherKeyringAccountCounts[index], + ); + importedSrpKeyringIds.push(importedSrpKeyring); + } + + const newHdKeyringIds = [newPrimaryKeyringId, ...importedSrpKeyringIds]; + Logger.log('newHdKeyringIds', newHdKeyringIds); + // 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) { + Logger.log( + 'snapAccount', + snapAccount, + keyringIdMap.get(snapAccount.options.entropySource), + ); + 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) { From 58deb356695db5844e044c1e6c69fb90cc527e0c Mon Sep 17 00:00:00 2001 From: Monte Lai Date: Thu, 8 May 2025 22:50:43 +0800 Subject: [PATCH 02/12] fix: add fence for multisrp --- app/core/Vault.js | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/app/core/Vault.js b/app/core/Vault.js index a49fd790c7f2..bdaa25830a5f 100644 --- a/app/core/Vault.js +++ b/app/core/Vault.js @@ -235,6 +235,7 @@ export const recreateVaultWithNewPassword = async ( // recreate import srp accounts const importedSrpKeyringIds = []; + ///: BEGIN:ONLY_INCLUDE_IF(multi-srp) for (const [index, otherSeedPhrase] of otherSeedPhrases.entries()) { const importedSrpKeyring = await restoreImportedSrp( otherSeedPhrase, @@ -242,9 +243,8 @@ export const recreateVaultWithNewPassword = async ( ); importedSrpKeyringIds.push(importedSrpKeyring); } - + ///: END:ONLY_INCLUDE_IF(multi-srp) const newHdKeyringIds = [newPrimaryKeyringId, ...importedSrpKeyringIds]; - Logger.log('newHdKeyringIds', newHdKeyringIds); // map old keyring id to new keyring id const keyringIdMap = new Map(); for (const [index, keyring] of hdKeyringsWithMetadata.entries()) { @@ -253,11 +253,6 @@ export const recreateVaultWithNewPassword = async ( // recreate snap accounts for (const snapAccount of firstPartySnapAccounts) { - Logger.log( - 'snapAccount', - snapAccount, - keyringIdMap.get(snapAccount.options.entropySource), - ); await restoreSnapAccounts( snapAccount.type, keyringIdMap.get(snapAccount.options.entropySource), From 07b8dd01602ce2f388c8f053b3e7e1d7a6b4f8eb Mon Sep 17 00:00:00 2001 From: Monte Lai Date: Thu, 8 May 2025 23:13:43 +0800 Subject: [PATCH 03/12] fix: lint --- .../SnapKeyring/MultichainWalletSnapClient.ts | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/app/core/SnapKeyring/MultichainWalletSnapClient.ts b/app/core/SnapKeyring/MultichainWalletSnapClient.ts index 01e86cdd4313..64dcc9a0f02c 100644 --- a/app/core/SnapKeyring/MultichainWalletSnapClient.ts +++ b/app/core/SnapKeyring/MultichainWalletSnapClient.ts @@ -12,7 +12,6 @@ import { SolanaWalletSnapSender, } from './SolanaWalletSnap'; import Engine from '../Engine'; -import { KeyringTypes } from '@metamask/keyring-controller'; import { SnapKeyring } from '@metamask/eth-snap-keyring'; import { MultichainNetwork } from '@metamask/multichain-transactions-controller'; import { store } from '../../store'; @@ -109,13 +108,14 @@ export abstract class MultichainWalletSnapClient { }), ); - return await this.withSnapKeyring(async (keyring) => { - return await keyring.createAccount( - this.snapId, - options as unknown as Record, - this.snapKeyringOptions, - ); - }); + return await this.withSnapKeyring( + async (keyring) => + await keyring.createAccount( + this.snapId, + options as unknown as Record, + this.snapKeyringOptions, + ), + ); } /** From 7abf52afdc72df1c66ee7863d0a16940aa4c62dc Mon Sep 17 00:00:00 2001 From: Monte Lai Date: Fri, 9 May 2025 08:50:51 +0800 Subject: [PATCH 04/12] fix: remove logs --- app/core/Vault.js | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/app/core/Vault.js b/app/core/Vault.js index bdaa25830a5f..d4dfaf11d0d5 100644 --- a/app/core/Vault.js +++ b/app/core/Vault.js @@ -43,7 +43,6 @@ export const restoreLedgerKeyring = async (serializedLedgerKeyring) => { export const restoreImportedSrp = async (seedPhrase, numberOfAccounts) => { const { KeyringController } = Engine.context; - Logger.log('restoreImportedSrp', seedPhrase, numberOfAccounts); try { const { id: keyringId } = await KeyringController.addNewKeyring( KeyringTypes.hd, @@ -52,10 +51,7 @@ export const restoreImportedSrp = async (seedPhrase, numberOfAccounts) => { }, ); - Logger.log('keyringId', keyringId); - for (let i = 0; i < numberOfAccounts; i++) { - Logger.log('adding account', i); await KeyringController.withKeyring( { id: keyringId }, async ({ keyring }) => await keyring.addAccounts(1), @@ -244,6 +240,7 @@ export const recreateVaultWithNewPassword = async ( 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(); From c978fbd663aa6d334f74389df308823dff72ca47 Mon Sep 17 00:00:00 2001 From: Monte Lai Date: Fri, 9 May 2025 15:41:25 +0800 Subject: [PATCH 05/12] feat: add unit tests --- app/core/Vault.js | 2 +- app/core/Vault.test.ts | 385 ++++++++++++++++++- app/util/test/accountsControllerTestUtils.ts | 54 ++- 3 files changed, 415 insertions(+), 26 deletions(-) diff --git a/app/core/Vault.js b/app/core/Vault.js index d4dfaf11d0d5..a16f73111067 100644 --- a/app/core/Vault.js +++ b/app/core/Vault.js @@ -60,7 +60,7 @@ export const restoreImportedSrp = async (seedPhrase, numberOfAccounts) => { return keyringId; } catch (e) { - Logger.log( + Logger.error( e, 'error while trying to restore imported srp accounts on recreate vault', ); diff --git a/app/core/Vault.test.ts b/app/core/Vault.test.ts index 038f73e26683..a47161cc66b6 100644 --- a/app/core/Vault.test.ts +++ b/app/core/Vault.test.ts @@ -1,14 +1,192 @@ +import { EthAccountType, SolAccountType } from '@metamask/keyring-api'; import Logger from '../util/Logger'; import Engine from './Engine'; import { withLedgerKeyring } from './Ledger/Ledger'; -import { restoreLedgerKeyring, restoreQRKeyring } from './Vault'; +import { + restoreLedgerKeyring, + restoreQRKeyring, + restoreSnapAccounts, + restoreImportedSrp, + recreateVaultWithNewPassword, +} from './Vault'; +import { KeyringSelector, KeyringTypes } from '@metamask/keyring-controller'; +import { + createMockInternalAccount, + createMockSnapInternalAccount, +} from '../util/test/accountsControllerTestUtils'; + +const mockAddNewKeyring = jest.fn(); +const mockWithKeyring = jest.fn(); +const mockExportSeedPhrase = jest.fn(); +const mockListMultichainAccounts = jest.fn(); +const mockCreateNewVaultAndRestore = jest.fn(); +const mockAddNewAccount = jest.fn(); +const mockRestoreQRKeyring = jest.fn(); +const mockRestoreLedgerKeyring = jest.fn(); +const mockExportAccount = jest.fn(); +const mockImportAccountWithStrategy = jest.fn(); + +const mockHdAccount1 = createMockInternalAccount( + '0x67B2fAf7959fB61eb9746571041476Bbd0672569', + 'Hd Account 1', + KeyringTypes.hd, + EthAccountType.Eoa, +); +const mockHdAccount2 = createMockInternalAccount( + '0xf57E323fD8C7Bb908A13557b1cE4441c8213824c', + 'Hd Account 2', + KeyringTypes.hd, + EthAccountType.Eoa, +); +const mockHdAccount3 = createMockInternalAccount( + '0x67B2fAf7959fB61eb9746571041476Bbd0672569', + 'Hd Account 3', + KeyringTypes.hd, + EthAccountType.Eoa, +); +const mockHdAccount4 = createMockInternalAccount( + '0xf57E323fD8C7Bb908A13557b1cE4441c8213824c', + 'Hd Account 4', + KeyringTypes.hd, + EthAccountType.Eoa, +); +const mockQrAccount = createMockInternalAccount( + '0xB6ce536fe74d9d5a25356B249424E1D3BDaADC57', + 'Qr Account', + KeyringTypes.qr, + EthAccountType.Eoa, +); +const mockLedgerAccount = createMockInternalAccount( + '0xE3D19DCfE5255C3448CB80299906ac15Dee2cF29', + 'Ledger Account', + KeyringTypes.ledger, + EthAccountType.Eoa, +); +const mockSolanaAccount = createMockSnapInternalAccount( + '', + 'Solana Account 1', + SolAccountType.DataAccount, + 'hd-keyring-id-2', +); +const mockThirdPartySnapAccount = createMockSnapInternalAccount( + '0x571CE5f203f93662301D43A5020aDB895f842cC9', + 'Third Party Snap Account', + EthAccountType.Eoa, +); + +const mockPrivateKeyAccount = createMockInternalAccount( + '0x15fB9d189aEa6233bE7487294702A431a7656810', + 'Private Key Account', + KeyringTypes.simple, + EthAccountType.Eoa, +); + +const mockHdKeyring = { + type: KeyringTypes.hd, + accounts: [mockHdAccount1.address, mockHdAccount2.address], +}; + +const mockHdKeyring2 = { + type: KeyringTypes.hd, + accounts: [mockHdAccount3.address, mockHdAccount4.address], +}; + +const mockQrKeyring = { + type: KeyringTypes.qr, + accounts: [mockQrAccount.address], +}; + +const mockLedgerKeyring = { + type: KeyringTypes.ledger, + accounts: [mockLedgerAccount.address], +}; + +const mockSnapKeyring = { + type: KeyringTypes.snap, + accounts: [mockSolanaAccount.address, mockThirdPartySnapAccount.address], +}; + +const mockSimpleKeyring = { + type: KeyringTypes.simple, + accounts: [mockPrivateKeyAccount.address], +}; + +const mockHdKeyringMetadata = { + id: 'hd-keyring-id', + name: '', +}; + +const mockHdKeyringMetadata2 = { + id: 'hd-keyring-id-2', + name: '', +}; + +const mockQrKeyringMetadata = { + id: 'qr-keyring-id', + name: '', +}; + +const mockLedgerKeyringMetadata = { + id: 'ledger-keyring-id', + name: '', +}; + +const mockSnapKeyringMetadata = { + id: 'snap-keyring-id', + name: '', +}; + +const mockSimpleKeyringMetadata = { + id: 'simple-keyring-id', + name: '', +}; jest.mock('./Engine', () => ({ context: { KeyringController: { - restoreQRKeyring: jest.fn(), - withKeyring: jest.fn(), + withKeyring: (selectedKeyring: KeyringSelector, callback: any) => + mockWithKeyring(selectedKeyring, callback), + addNewKeyring: (type: KeyringTypes, options: { mnemonic?: string }) => + mockAddNewKeyring(type, options), + exportSeedPhrase: (password: string, keyringId: string) => + mockExportSeedPhrase(password, keyringId), + createNewVaultAndRestore: (password: string, seedPhrase: string) => + mockCreateNewVaultAndRestore(password, seedPhrase), + addNewAccount: () => mockAddNewAccount(), + restoreQRKeyring: (serializedKeyring: string) => + mockRestoreQRKeyring(serializedKeyring), + restoreLedgerKeyring: (serializedKeyring: string) => + mockRestoreLedgerKeyring(serializedKeyring), + exportAccount: (password: string, account: string) => + mockExportAccount(password, account), + importAccountWithStrategy: (strategy: string, accounts: string[]) => + mockImportAccountWithStrategy(strategy, accounts), + state: { + get keyrings() { + return [ + mockHdKeyring, + mockHdKeyring2, + mockQrKeyring, + mockLedgerKeyring, + mockSnapKeyring, + mockSimpleKeyring, + ]; + }, + get keyringsMetadata() { + return [ + mockHdKeyringMetadata, + mockHdKeyringMetadata2, + mockQrKeyringMetadata, + mockLedgerKeyringMetadata, + mockSnapKeyringMetadata, + mockSimpleKeyringMetadata, + ]; + }, + }, + }, + AccountsController: { + listMultichainAccounts: () => mockListMultichainAccounts(), }, }, })); @@ -23,6 +201,19 @@ jest.mock('../util/Logger', () => ({ error: jest.fn(), })); +const mockMultichainWalletSnapClient = { + createAccount: jest.fn(), +}; + +jest.mock('./SnapKeyring/MultichainWalletSnapClient', () => ({ + ...jest.requireActual('./SnapKeyring/MultichainWalletSnapClient'), + MultichainWalletSnapFactory: { + createClient: jest + .fn() + .mockImplementation(() => mockMultichainWalletSnapClient), + }, +})); + describe('Vault', () => { beforeEach(() => { jest.clearAllMocks(); @@ -34,19 +225,15 @@ describe('Vault', () => { await restoreQRKeyring(mockSerializedQrKeyring); - expect( - MockEngine.context.KeyringController.restoreQRKeyring, - ).toHaveBeenCalled(); - expect(KeyringController.restoreQRKeyring).toHaveBeenCalledWith( + expect(mockRestoreQRKeyring).toHaveBeenCalled(); + expect(mockRestoreQRKeyring).toHaveBeenCalledWith( mockSerializedQrKeyring, ); }); it('should log error if an exception is thrown', async () => { const error = new Error('Test error'); - MockEngine.context.KeyringController.restoreQRKeyring.mockRejectedValue( - error, - ); + mockRestoreQRKeyring.mockRejectedValue(error); const mockSerializedQrKeyring = 'serialized-keyring'; await restoreQRKeyring(mockSerializedQrKeyring); @@ -60,30 +247,30 @@ describe('Vault', () => { describe('restoreLedgerKeyring', () => { it('should restore ledger keyring if it exists', async () => { - const mockLedgerKeyring = { + const mockLedgerKeyringInstance = { deserialize: jest.fn(), }; mockWithLedgerKeyring.mockImplementation( // @ts-expect-error The Ledger keyring is not compatible with our keyring type yet - (operation) => operation(mockLedgerKeyring), + (operation) => operation(mockLedgerKeyringInstance), ); const mockSerializedLedgerKeyring = 'serialized-keyring'; await restoreLedgerKeyring(mockSerializedLedgerKeyring); - expect(mockLedgerKeyring.deserialize).toHaveBeenCalledWith( + expect(mockLedgerKeyringInstance.deserialize).toHaveBeenCalledWith( mockSerializedLedgerKeyring, ); }); it('should log error if the Ledger keyring throws an error', async () => { const error = new Error('Test error'); - const mockLedgerKeyring = { + const mockLedgerKeyringInstance = { deserialize: jest.fn().mockRejectedValue(error), }; mockWithLedgerKeyring.mockImplementation( // @ts-expect-error The Ledger keyring is not compatible with our keyring type yet - (operation) => operation(mockLedgerKeyring), + (operation) => operation(mockLedgerKeyringInstance), ); const mockSerializedLedgerKeyring = 'serialized-keyring'; @@ -108,4 +295,172 @@ describe('Vault', () => { ); }); }); + + describe('restoreSnapAccounts', () => { + it('creates a snap account if it exists', async () => { + await restoreSnapAccounts(SolAccountType.DataAccount, 'entropy-source'); + + expect(mockMultichainWalletSnapClient.createAccount).toHaveBeenCalledWith( + { + entropySource: 'entropy-source', + }, + ); + }); + + it('logs error if snap account creation fails', async () => { + const error = new Error('Test error'); + mockMultichainWalletSnapClient.createAccount.mockRejectedValue(error); + + await restoreSnapAccounts(SolAccountType.DataAccount, 'entropy-source'); + + expect(Logger.error).toHaveBeenCalledWith( + error, + 'error while trying to restore snap accounts on recreate vault', + ); + }); + }); + describe('restoreImportedSrp', () => { + it('restore imported srp accounts if they exist', async () => { + mockAddNewKeyring.mockResolvedValue({ id: 'keyring-id' }); + mockWithKeyring.mockResolvedValue(null); + await restoreImportedSrp('seed-phrase', 5); + + expect(mockAddNewKeyring).toHaveBeenCalledWith(KeyringTypes.hd, { + mnemonic: 'seed-phrase', + }); + + expect(mockWithKeyring).toHaveBeenCalledTimes(5); + expect(mockWithKeyring).toHaveBeenCalledWith( + { id: 'keyring-id' }, + expect.any(Function), + ); + }); + + it('logs error if srp account creation fails', async () => { + const error = new Error('Test error'); + mockAddNewKeyring.mockRejectedValue(error); + + await restoreImportedSrp('seed-phrase', 5); + + expect(Logger.error).toHaveBeenCalledWith( + error, + 'error while trying to restore imported srp accounts on recreate vault', + ); + }); + }); + + describe('recreateVaultWithNewPassword', () => { + it('should recreate vault with new password', async () => { + const newPassword = 'new-password'; + const primarySeedPhrase = 'seed-phrase'; + const secondarySeedPhrase = 'seed-phrase-2'; + const mockNewKeyringIdForSecondSrp = 'new-keyring-id'; + const mockPrivateKey = 'very-private-key'; + + mockExportSeedPhrase + .mockResolvedValueOnce([primarySeedPhrase]) + .mockResolvedValueOnce([secondarySeedPhrase]); + mockListMultichainAccounts.mockReturnValue([ + mockHdAccount1, + mockHdAccount2, + mockQrAccount, + mockLedgerAccount, + mockSolanaAccount, + mockThirdPartySnapAccount, + ]); + mockCreateNewVaultAndRestore.mockResolvedValue(null); + mockAddNewAccount.mockResolvedValue(''); + mockRestoreQRKeyring.mockResolvedValue(null); + mockRestoreLedgerKeyring.mockResolvedValue(null); + mockAddNewKeyring.mockResolvedValue({ id: mockNewKeyringIdForSecondSrp }); // New srp index + mockWithKeyring + // 1st call is to get serialized ledger keyring + .mockResolvedValueOnce(mockLedgerKeyring) + // 2nd call is to get serialized qr keyring + .mockResolvedValueOnce(mockQrKeyring) + // 3rd call is to add accounts when restoring imported srp + .mockResolvedValueOnce(null); + mockExportAccount.mockResolvedValue(mockPrivateKey); + + await recreateVaultWithNewPassword( + 'password', + newPassword, + mockHdAccount1.address, + ); + + // There are 2 hd keyrings, so 2 seed phrases are exported + expect(mockExportSeedPhrase).toHaveBeenCalledWith( + 'password', + mockHdKeyringMetadata.id, + ); + expect(mockExportSeedPhrase).toHaveBeenCalledWith( + 'password', + mockHdKeyringMetadata2.id, + ); + + expect(mockCreateNewVaultAndRestore).toHaveBeenCalledWith(newPassword, [ + primarySeedPhrase, + ]); + + // Only 2 account is in the primary keyring + expect(mockAddNewAccount).toHaveBeenCalledTimes(1); + expect(mockAddNewAccount).toHaveBeenCalledWith(); + + // Import private key accounts + expect(mockExportAccount).toHaveBeenCalledWith( + 'password', + mockPrivateKeyAccount.address, + ); + expect(mockImportAccountWithStrategy).toHaveBeenCalledWith('privateKey', [ + mockPrivateKey, + ]); + + // Ledger accounts + expect(mockWithKeyring).toHaveBeenNthCalledWith( + 1, + { type: KeyringTypes.ledger }, + expect.any(Function), + ); + expect(withLedgerKeyring).toHaveBeenCalledWith(expect.any(Function)); + + // QR Accounts + expect(mockWithKeyring).toHaveBeenNthCalledWith( + 2, + { type: KeyringTypes.qr }, + expect.any(Function), + ); + expect(mockRestoreQRKeyring).toHaveBeenCalledWith(mockQrKeyring); + + // Imported SRP Accounts + expect(mockAddNewKeyring).toHaveBeenCalledWith(KeyringTypes.hd, { + mnemonic: [secondarySeedPhrase], + }); + // 3rd and 4th call is to add accounts when restoring imported srp + expect(mockWithKeyring).toHaveBeenNthCalledWith( + 3, + { + id: mockNewKeyringIdForSecondSrp, + }, + expect.any(Function), + ); + expect(mockWithKeyring).toHaveBeenNthCalledWith( + 4, + { + id: mockNewKeyringIdForSecondSrp, + }, + expect.any(Function), + ); + + // Snap accounts + // only called once because third party snaps are not restored + expect( + mockMultichainWalletSnapClient.createAccount, + ).toHaveBeenCalledTimes(1); + expect(mockMultichainWalletSnapClient.createAccount).toHaveBeenCalledWith( + { + entropySource: mockNewKeyringIdForSecondSrp, + }, + ); + }); + }); }); diff --git a/app/util/test/accountsControllerTestUtils.ts b/app/util/test/accountsControllerTestUtils.ts index cdbed0ed34b2..579e2dcc9a90 100644 --- a/app/util/test/accountsControllerTestUtils.ts +++ b/app/util/test/accountsControllerTestUtils.ts @@ -109,7 +109,45 @@ export function createMockInternalAccount( export function createMockSnapInternalAccount( address: string, nickname: string, + accountType: KeyringAccountType = EthAccountType.Eoa, + entropySource: string = '', ): InternalAccount { + let methods: string[] = []; + switch (accountType) { + case EthAccountType.Eoa: + methods = [ + EthMethod.PersonalSign, + EthMethod.SignTransaction, + EthMethod.SignTypedDataV1, + EthMethod.SignTypedDataV3, + EthMethod.SignTypedDataV4, + ]; + break; + case BtcAccountType.P2wpkh: + methods = [BtcMethod.SendBitcoin]; + break; + case SolAccountType.DataAccount: + methods = [SolMethod.SendAndConfirmTransaction]; + break; + default: + throw new Error(`Unsupported account type: ${accountType}`); + } + + let type: KeyringAccountType; + switch (accountType) { + case EthAccountType.Eoa: + type = EthAccountType.Eoa; + break; + case BtcAccountType.P2wpkh: + type = BtcAccountType.P2wpkh; + break; + case SolAccountType.DataAccount: + type = SolAccountType.DataAccount; + break; + default: + throw new Error(`Unsupported account type: ${accountType}`); + } + return { address, id: createMockUuidFromAddress(address), @@ -125,16 +163,12 @@ export function createMockSnapInternalAccount( enabled: true, }, }, - options: {}, - methods: [ - EthMethod.PersonalSign, - EthMethod.SignTransaction, - EthMethod.SignTypedDataV1, - EthMethod.SignTypedDataV3, - EthMethod.SignTypedDataV4, - ], - type: EthAccountType.Eoa, - scopes: [EthScope.Eoa], + options: { + entropySource, + }, + methods, + type, + scopes: getAccountTypeScopes(accountType), }; } From b0b005ef640bcb70a6a47d4abda5c541e8ff6bbd Mon Sep 17 00:00:00 2001 From: Monte Lai Date: Fri, 9 May 2025 16:05:29 +0800 Subject: [PATCH 06/12] fix: lint --- app/core/Vault.test.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/app/core/Vault.test.ts b/app/core/Vault.test.ts index a47161cc66b6..0b0c69e55db9 100644 --- a/app/core/Vault.test.ts +++ b/app/core/Vault.test.ts @@ -145,6 +145,7 @@ const mockSimpleKeyringMetadata = { jest.mock('./Engine', () => ({ context: { KeyringController: { + // @typescript-eslint/no-explicit-any - Mock any callback. withKeyring: (selectedKeyring: KeyringSelector, callback: any) => mockWithKeyring(selectedKeyring, callback), addNewKeyring: (type: KeyringTypes, options: { mnemonic?: string }) => @@ -220,7 +221,6 @@ describe('Vault', () => { }); describe('restoreQRKeyring', () => { it('should restore QR keyring if it exists', async () => { - const { KeyringController } = MockEngine.context; const mockSerializedQrKeyring = 'serialized-keyring'; await restoreQRKeyring(mockSerializedQrKeyring); From 1627b9b5622b65e09ca3ec43af12a52c630490ad Mon Sep 17 00:00:00 2001 From: Monte Lai Date: Fri, 9 May 2025 16:53:21 +0800 Subject: [PATCH 07/12] fix:lint --- app/core/Vault.test.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/app/core/Vault.test.ts b/app/core/Vault.test.ts index 0b0c69e55db9..68ceb799e412 100644 --- a/app/core/Vault.test.ts +++ b/app/core/Vault.test.ts @@ -145,7 +145,7 @@ const mockSimpleKeyringMetadata = { jest.mock('./Engine', () => ({ context: { KeyringController: { - // @typescript-eslint/no-explicit-any - Mock any callback. + // eslint-disable-next-line @typescript-eslint/no-explicit-any - Mock any callback. withKeyring: (selectedKeyring: KeyringSelector, callback: any) => mockWithKeyring(selectedKeyring, callback), addNewKeyring: (type: KeyringTypes, options: { mnemonic?: string }) => @@ -191,7 +191,7 @@ jest.mock('./Engine', () => ({ }, }, })); -const MockEngine = jest.mocked(Engine); +jest.mocked(Engine); jest.mock('./Ledger/Ledger', () => ({ withLedgerKeyring: jest.fn(), From 7a14c1b486360508c0b3d4091607f58c63955830 Mon Sep 17 00:00:00 2001 From: Monte Lai Date: Fri, 9 May 2025 17:33:04 +0800 Subject: [PATCH 08/12] fix: tests --- .../MultichainWalletSnapClient.test.ts | 55 ++++++++----------- 1 file changed, 23 insertions(+), 32 deletions(-) diff --git a/app/core/SnapKeyring/MultichainWalletSnapClient.test.ts b/app/core/SnapKeyring/MultichainWalletSnapClient.test.ts index 33a3e46dbacc..adf8f593b974 100644 --- a/app/core/SnapKeyring/MultichainWalletSnapClient.test.ts +++ b/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, From 9926d8b8065a6436c690fb96b156f98e50662752 Mon Sep 17 00:00:00 2001 From: Monte Lai Date: Fri, 9 May 2025 19:14:36 +0800 Subject: [PATCH 09/12] fix: add await --- app/core/SnapKeyring/MultichainWalletSnapClient.ts | 4 ++-- app/core/Vault.js | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/app/core/SnapKeyring/MultichainWalletSnapClient.ts b/app/core/SnapKeyring/MultichainWalletSnapClient.ts index 06cbe08cf158..7419837ae0fa 100644 --- a/app/core/SnapKeyring/MultichainWalletSnapClient.ts +++ b/app/core/SnapKeyring/MultichainWalletSnapClient.ts @@ -120,7 +120,7 @@ export abstract class MultichainWalletSnapClient { ); return await this.withSnapKeyring(async (keyring) => { - keyring.createAccount( + await keyring.createAccount( this.snapId, { ...options, @@ -205,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, diff --git a/app/core/Vault.js b/app/core/Vault.js index a16f73111067..dfd15f613354 100644 --- a/app/core/Vault.js +++ b/app/core/Vault.js @@ -128,9 +128,9 @@ export const recreateVaultWithNewPassword = async ( .filter((keyring) => keyring.type === KeyringTypes.hd); const seedPhrases = await Promise.all( - hdKeyringsWithMetadata.map((keyring) => { + hdKeyringsWithMetadata.map(async (keyring) => { try { - return getSeedPhrase(password, keyring.metadata.id); + return await getSeedPhrase(password, keyring.metadata.id); } catch (e) { Logger.error( e, From 54699b1ba3ea89244ca8c8870aebe160982c1077 Mon Sep 17 00:00:00 2001 From: Monte Lai Date: Fri, 9 May 2025 20:14:25 +0800 Subject: [PATCH 10/12] fix: eslint disable --- app/core/Vault.test.ts | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/app/core/Vault.test.ts b/app/core/Vault.test.ts index 68ceb799e412..5bb158bf4eea 100644 --- a/app/core/Vault.test.ts +++ b/app/core/Vault.test.ts @@ -145,7 +145,8 @@ const mockSimpleKeyringMetadata = { jest.mock('./Engine', () => ({ context: { KeyringController: { - // eslint-disable-next-line @typescript-eslint/no-explicit-any - Mock any callback. + // Using any to mock any callback. + // eslint-disable-next-line @typescript-eslint/no-explicit-any withKeyring: (selectedKeyring: KeyringSelector, callback: any) => mockWithKeyring(selectedKeyring, callback), addNewKeyring: (type: KeyringTypes, options: { mnemonic?: string }) => From b562e8d51b0a732a199b30b5ab49a7f68c763191 Mon Sep 17 00:00:00 2001 From: Monte Lai Date: Fri, 9 May 2025 22:37:34 +0800 Subject: [PATCH 11/12] fix: solana snap name after restore --- app/core/Vault.js | 10 +++++++++- app/core/Vault.test.ts | 7 ++++++- 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/app/core/Vault.js b/app/core/Vault.js index dfd15f613354..243920bdd556 100644 --- a/app/core/Vault.js +++ b/app/core/Vault.js @@ -6,7 +6,11 @@ import { MultichainWalletSnapFactory, WalletClientType, } from './SnapKeyring/MultichainWalletSnapClient'; -import { BtcAccountType, SolAccountType } from '@metamask/keyring-api'; +import { + BtcAccountType, + SolAccountType, + SolScope, +} from '@metamask/keyring-api'; /** * Restore the given serialized QR keyring. @@ -69,13 +73,16 @@ export const restoreImportedSrp = async (seedPhrase, numberOfAccounts) => { 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: @@ -86,6 +93,7 @@ export const restoreSnapAccounts = async (accountType, entropySource) => { const client = MultichainWalletSnapFactory.createClient(walletClientType); await client.createAccount({ entropySource, + scope, }); } catch (e) { Logger.error( diff --git a/app/core/Vault.test.ts b/app/core/Vault.test.ts index 5bb158bf4eea..76f06f9d8311 100644 --- a/app/core/Vault.test.ts +++ b/app/core/Vault.test.ts @@ -1,4 +1,8 @@ -import { EthAccountType, SolAccountType } from '@metamask/keyring-api'; +import { + EthAccountType, + SolAccountType, + SolScope, +} from '@metamask/keyring-api'; import Logger from '../util/Logger'; import Engine from './Engine'; import { withLedgerKeyring } from './Ledger/Ledger'; @@ -460,6 +464,7 @@ describe('Vault', () => { expect(mockMultichainWalletSnapClient.createAccount).toHaveBeenCalledWith( { entropySource: mockNewKeyringIdForSecondSrp, + scope: SolScope.Mainnet, }, ); }); From ef9533627c03aedd87713d94215873feb4d871b4 Mon Sep 17 00:00:00 2001 From: Monte Lai Date: Fri, 9 May 2025 23:37:15 +0800 Subject: [PATCH 12/12] fix: lint and test --- app/core/Vault.js | 1 + app/core/Vault.test.ts | 1 + 2 files changed, 2 insertions(+) diff --git a/app/core/Vault.js b/app/core/Vault.js index 243920bdd556..d3c50acfc2fd 100644 --- a/app/core/Vault.js +++ b/app/core/Vault.js @@ -8,6 +8,7 @@ import { } from './SnapKeyring/MultichainWalletSnapClient'; import { BtcAccountType, + BtcScope, SolAccountType, SolScope, } from '@metamask/keyring-api'; diff --git a/app/core/Vault.test.ts b/app/core/Vault.test.ts index 76f06f9d8311..b545614ba862 100644 --- a/app/core/Vault.test.ts +++ b/app/core/Vault.test.ts @@ -308,6 +308,7 @@ describe('Vault', () => { expect(mockMultichainWalletSnapClient.createAccount).toHaveBeenCalledWith( { entropySource: 'entropy-source', + scope: SolScope.Mainnet, }, ); });