diff --git a/packages/account-tree-controller/CHANGELOG.md b/packages/account-tree-controller/CHANGELOG.md index b09828e5624..02ccc548bad 100644 --- a/packages/account-tree-controller/CHANGELOG.md +++ b/packages/account-tree-controller/CHANGELOG.md @@ -7,6 +7,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Changed + +- **BREAKING:** Properly removes "outdated" accounts from the `MultichainAccountService` ([#7038](https://github.com/MetaMask/core/pull/7038)) + - This change will ensure that "disabled" accounts are not properly removed from the account tree too. + - The `MultichainAccountService:multichainAccountGroupCreated` event is now required. + - The `MultichainAccountService:multichainAccountGroupUpdated` event is now required. + - The `MultichainAccountService:getMultichainAccountWallets` action is now required. + ## [2.0.0] ### Changed diff --git a/packages/account-tree-controller/src/AccountTreeController.test.ts b/packages/account-tree-controller/src/AccountTreeController.test.ts index d98b819b97d..fe7b055e078 100644 --- a/packages/account-tree-controller/src/AccountTreeController.test.ts +++ b/packages/account-tree-controller/src/AccountTreeController.test.ts @@ -1,7 +1,12 @@ -import type { AccountWalletId, Bip44Account } from '@metamask/account-api'; +import type { + AccountWalletId, + Bip44Account, + MultichainAccountGroup, +} from '@metamask/account-api'; import { AccountGroupType, AccountWalletType, + isBip44Account, toAccountGroupId, toAccountWalletId, toMultichainAccountGroupId, @@ -10,6 +15,7 @@ import { } from '@metamask/account-api'; import type { AccountId } from '@metamask/accounts-controller'; import { deriveStateFromMetadata } from '@metamask/base-controller'; +import type { EntropySourceId, KeyringAccount } from '@metamask/keyring-api'; import { BtcAccountType, EthAccountType, @@ -232,6 +238,48 @@ const MOCK_HARDWARE_ACCOUNT_1: InternalAccount = { const mockGetSelectedMultichainAccountActionHandler = jest.fn(); +/** + * Get multichain account wallets for the given list of accounts. + * + * @param accounts Accounts to build wallets for. + * @returns Multichain account wallets. + */ +function getMultichainAccountWalletsFromAccounts(accounts: KeyringAccount[]) { + const bip44Accounts = accounts.filter(isBip44Account); + + const wallets: { + [entropySource: EntropySourceId]: { + [groupIndex: number]: Set>; + }; + } = {}; + for (const bip44Account of bip44Accounts) { + const { id: entropySource, groupIndex } = bip44Account.options.entropy; + + wallets[entropySource] ??= {}; + wallets[entropySource][groupIndex] ??= new Set(); + wallets[entropySource][groupIndex].add(bip44Account); + } + + return Object.entries(wallets).map(([entropySource, groups]) => { + const walletId = toMultichainAccountWalletId(entropySource); + return { + id: walletId, + getMultichainAccountGroups: jest.fn().mockImplementation(() => { + return Object.entries(groups).map(([groupIndex, groupAccounts]) => { + const groupId = toMultichainAccountGroupId( + walletId, + Number(groupIndex), + ); + return { + id: groupId, + getAccounts: jest.fn().mockReturnValue(Array.from(groupAccounts)), + }; + }); + }), + }; + }); +} + /** * Sets up the AccountTreeController for testing. * @@ -304,6 +352,9 @@ function setup({ getSelectedMultichainAccount: jest.Mock; getAccount: jest.Mock; }; + MultichainAccountService: { + getMultichainAccountWallets: jest.Mock; + }; UserStorageController: { performGetStorage: jest.Mock; performGetStorageAllFeatureEntries: jest.Mock; @@ -327,6 +378,9 @@ function setup({ getAccount: jest.fn(), getSelectedMultichainAccount: jest.fn(), }, + MultichainAccountService: { + getMultichainAccountWallets: jest.fn(), + }, UserStorageController: { getState: jest.fn(), performGetStorage: jest.fn(), @@ -410,6 +464,17 @@ function setup({ 'UserStorageController:performBatchSetStorage', mocks.UserStorageController.performBatchSetStorage, ); + + mocks.MultichainAccountService.getMultichainAccountWallets.mockImplementation( + () => + getMultichainAccountWalletsFromAccounts( + mocks.AccountsController.accounts, + ), + ); + messenger.registerActionHandler( + 'MultichainAccountService:getMultichainAccountWallets', + mocks.MultichainAccountService.getMultichainAccountWallets, + ); } if (keyrings) { @@ -943,12 +1008,18 @@ describe('AccountTreeController', () => { }); const now = Date.now(); - mocks.AccountsController.listMultichainAccounts.mockReturnValue([ + const mockAccounts = [ // Faking accounts to be out of order: mockAccountWith(1, now + 1000), mockAccountWith(2, now + 2000), mockAccountWith(0, now), - ]); + ]; + mocks.AccountsController.listMultichainAccounts.mockReturnValue( + mockAccounts, + ); + mocks.MultichainAccountService.getMultichainAccountWallets.mockImplementation( + () => getMultichainAccountWalletsFromAccounts(mockAccounts), + ); controller.init(); @@ -964,6 +1035,30 @@ describe('AccountTreeController', () => { expect(groupIds[1]).toBe(toMultichainAccountGroupId(walletId, 1)); expect(groupIds[2]).toBe(toMultichainAccountGroupId(walletId, 2)); }); + + it('skips disabled bip-44 accounts', () => { + const { controller, mocks } = setup({ + accounts: [MOCK_HD_ACCOUNT_1, MOCK_TRX_ACCOUNT_1], + keyrings: [MOCK_HD_KEYRING_1], + }); + + // We simulate TRX account to be "disabled" here, so the `AccountsController` has 2 accounts (EVM + TRX), but + // our wallets will only returns the EVM account here, meaning TRX will be skipped. + mocks.MultichainAccountService.getMultichainAccountWallets.mockImplementation( + () => getMultichainAccountWalletsFromAccounts([MOCK_HD_ACCOUNT_1]), + ); + + controller.init(); + + const group = controller.getAccountGroupObject( + toMultichainAccountGroupId( + toMultichainAccountWalletId(MOCK_HD_KEYRING_1.metadata.id), + MOCK_HD_ACCOUNT_1.options.entropy.groupIndex, + ), + ); + expect(group).toBeDefined(); + expect(group?.accounts).toHaveLength(1); // Just EVM, no TRX. + }); }); describe('getAccountGroupObject', () => { @@ -1340,6 +1435,82 @@ describe('AccountTreeController', () => { }); }); + describe('on MultichainAccountService:multichainAccountGroupUpdated', () => { + it('removes account from the tree if their no longer part of the multichain account group', () => { + const { controller, messenger } = setup({ + accounts: [MOCK_HD_ACCOUNT_1, MOCK_TRX_ACCOUNT_1], + keyrings: [MOCK_HD_KEYRING_1], + }); + + messenger.registerActionHandler( + 'SnapController:get', + () => + // TODO: Update this to avoid the unknown cast if possible. + MOCK_SNAP_3 as unknown as ReturnType< + SnapControllerGetSnap['handler'] + >, + ); + + controller.init(); + + const walletId = toMultichainAccountWalletId( + MOCK_HD_KEYRING_1.metadata.id, + ); + const groupId = toMultichainAccountGroupId(walletId, 0); + + let group = controller.getAccountGroupObject(groupId); + expect(group).toBeDefined(); + expect(group?.accounts).toHaveLength(2); + + messenger.publish( + 'MultichainAccountService:multichainAccountGroupUpdated', + // Partial implementation, so we need to cast here. + { + id: groupId, + // Simulate a group with just 1 account now, meaning that `MOCK_TRX_ACCOUNT_1` got + // "removed/disabled" from this multichain account group. + getAccounts: jest.fn().mockReturnValue([MOCK_HD_ACCOUNT_1]), + } as unknown as MultichainAccountGroup>, + ); + + group = controller.getAccountGroupObject(groupId); + expect(group).toBeDefined(); + expect(group?.accounts).toHaveLength(1); + }); + }); + + describe('on MultichainAccountService:multichainAccountGroupCreated', () => { + it('does nothing if an multichain account group got created before inserting any accounts', () => { + const { controller, messenger } = setup({ + accounts: [], + keyrings: [MOCK_HD_KEYRING_1], + }); + + controller.init(); + + const walletId = toMultichainAccountWalletId( + MOCK_HD_KEYRING_1.metadata.id, + ); + const groupId = toMultichainAccountGroupId(walletId, 0); + + let group = controller.getAccountGroupObject(groupId); + expect(group).toBeUndefined(); + + messenger.publish( + 'MultichainAccountService:multichainAccountGroupCreated', + // Partial implementation, so we need to cast here. + { + id: groupId, + getAccounts: jest.fn().mockReturnValue([MOCK_HD_ACCOUNT_1]), + } as unknown as MultichainAccountGroup>, + ); + + // Still undefined, since we're only reacting to `:accountAdded` for now. + group = controller.getAccountGroupObject(groupId); + expect(group).toBeUndefined(); + }); + }); + describe('account ordering by type', () => { it('orders accounts in group according to ACCOUNT_TYPE_TO_SORT_ORDER regardless of insertion order', () => { const evmAccount = MOCK_HD_ACCOUNT_1; diff --git a/packages/account-tree-controller/src/AccountTreeController.ts b/packages/account-tree-controller/src/AccountTreeController.ts index 2180076e233..dc2bdeca87a 100644 --- a/packages/account-tree-controller/src/AccountTreeController.ts +++ b/packages/account-tree-controller/src/AccountTreeController.ts @@ -1,16 +1,23 @@ -import { AccountWalletType, select } from '@metamask/account-api'; +import { + AccountWalletType, + isBip44Account, + select, +} from '@metamask/account-api'; import type { AccountGroupId, AccountWalletId, AccountSelector, MultichainAccountWalletId, AccountGroupType, + Bip44Account, } from '@metamask/account-api'; import type { MultichainAccountWalletStatus } from '@metamask/account-api'; +import type { MultichainAccountGroup } from '@metamask/account-api'; import { type AccountId } from '@metamask/accounts-controller'; import type { StateMetadata } from '@metamask/base-controller'; import { BaseController } from '@metamask/base-controller'; import type { TraceCallback } from '@metamask/controller-utils'; +import type { KeyringAccount } from '@metamask/keyring-api'; import { isEvmAccountType } from '@metamask/keyring-api'; import type { InternalAccount } from '@metamask/keyring-internal-api'; import { assert } from '@metamask/utils'; @@ -247,6 +254,20 @@ export class AccountTreeController extends BaseController< }, ); + this.messenger.subscribe( + 'MultichainAccountService:multichainAccountGroupCreated', + (group) => { + this.#handleMultichainAccountWalletGroupCreatedOrUpdated(group); + }, + ); + + this.messenger.subscribe( + 'MultichainAccountService:multichainAccountGroupUpdated', + (group) => { + this.#handleMultichainAccountWalletGroupCreatedOrUpdated(group); + }, + ); + this.#registerMessageHandlers(); } @@ -292,8 +313,22 @@ export class AccountTreeController extends BaseController< (a, b) => a.metadata.importTime - b.metadata.importTime, ); + // HACK: Since we're not consuming accounts directly from the service, we have to + // cross-check with its internal state before inserting initial BIP-44 accounts. + // Those accounts could be marked as "disabled" (using the + // `AccountProviderWrapper`) and if they are, we don't want to insert them in the tree. + // The real fix for this would be consume accounts from the service instead of the + // `AccountsController`, but this requires a bit more testing, for now we keep this + // simple. + const bip44Accounts = this.#getBip44AccountIds(); + // For now, we always re-compute all wallets, we do not re-use the existing state. for (const account of accounts) { + if (isBip44Account(account) && !bip44Accounts.has(account.id)) { + // We skip those, so they won't be part of the initial tree. + continue; + } + this.#insert(wallets, account); } @@ -846,43 +881,63 @@ export class AccountTreeController extends BaseController< return; } - const context = this.#accountIdToContext.get(accountId); + this.#removeAccounts([accountId]); + } - if (context) { - const { walletId, groupId } = context; + /** + * Removes an account from the tree if it exists. If not, this method does nothing. + * + * @param accountIds - Account IDs to remove from the tree. + */ + #removeAccounts(accountIds: AccountId[]) { + const removedAccounts: AccountId[] = []; - const previousSelectedAccountGroup = - this.state.accountTree.selectedAccountGroup; - let selectedAccountGroupChanged = false; + const previousSelectedAccountGroup = + this.state.accountTree.selectedAccountGroup; + let selectedAccountGroupChanged = false; - this.update((state) => { - const accounts = - state.accountTree.wallets[walletId]?.groups[groupId]?.accounts; - - if (accounts) { - const index = accounts.indexOf(accountId); - if (index !== -1) { - accounts.splice(index, 1); - - // Check if we need to update selectedAccountGroup after removal - if ( - state.accountTree.selectedAccountGroup === groupId && - accounts.length === 0 - ) { - // The currently selected group is now empty, find a new group to select - const newSelectedAccountGroup = this.#getDefaultAccountGroupId( - state.accountTree.wallets, - ); - state.accountTree.selectedAccountGroup = newSelectedAccountGroup; - selectedAccountGroupChanged = - newSelectedAccountGroup !== previousSelectedAccountGroup; + this.update((state) => { + for (const accountId of accountIds) { + const context = this.#accountIdToContext.get(accountId); + + if (context) { + const { walletId, groupId } = context; + + const accounts = + state.accountTree.wallets[walletId]?.groups[groupId]?.accounts; + + if (accounts) { + const index = accounts.indexOf(accountId); + if (index !== -1) { + accounts.splice(index, 1); + + // Now we know this account got removed. + removedAccounts.push(accountId); + + // Check if we need to update selectedAccountGroup after removal + if ( + state.accountTree.selectedAccountGroup === groupId && + accounts.length === 0 + ) { + // The currently selected group is now empty, find a new group to select + const newSelectedAccountGroup = this.#getDefaultAccountGroupId( + state.accountTree.wallets, + ); + state.accountTree.selectedAccountGroup = + newSelectedAccountGroup; + selectedAccountGroupChanged = + newSelectedAccountGroup !== previousSelectedAccountGroup; + } + } + if (accounts.length === 0) { + this.#pruneEmptyGroupAndWallet(state, walletId, groupId); } - } - if (accounts.length === 0) { - this.#pruneEmptyGroupAndWallet(state, walletId, groupId); } } - }); + } + }); + + if (removedAccounts.length) { this.messenger.publish( `${controllerName}:accountTreeChange`, this.state.accountTree, @@ -898,7 +953,9 @@ export class AccountTreeController extends BaseController< } // Clear reverse-mapping for that account. - this.#accountIdToContext.delete(accountId); + removedAccounts.forEach((accountId) => + this.#accountIdToContext.delete(accountId), + ); } } @@ -1219,6 +1276,72 @@ export class AccountTreeController extends BaseController< }); } + /** + * Get account IDs from a multichain account group. + * + * @param group - Multichain account group to extract account IDs from. + * @returns Set of BIP-44 account IDs for the given group. + */ + #getBip44AccountIdsFromMultichainAccountGroup( + group: MultichainAccountGroup>, + ): Set['id']> { + // FIXME: We should introduce a way to just get the account IDs since getting + // account object might be a bit more costly (at least with the current + // architecture). + return new Set(group.getAccounts().map((account) => account.id)); + } + + /** + * Get account IDs owned by the multichain account service. + * + * @returns Set of BIP-44 account IDs. + */ + #getBip44AccountIds() { + const accounts = new Set(); + + // Since we're building the tree out of of the account list, we don't know if those + // accounts got disabled at the service level. To make sure both states are "in-sync" + // we check each multichain account groups and remove extra accounts if there's any! + for (const wallet of this.messenger.call( + 'MultichainAccountService:getMultichainAccountWallets', + )) { + for (const group of wallet.getMultichainAccountGroups()) { + for (const account of this.#getBip44AccountIdsFromMultichainAccountGroup( + group, + )) { + accounts.add(account); + } + } + } + + return accounts; + } + + /** + * Handles multichain account group updates. + * + * @param multichainGroup - Multichain account group that got updated. + */ + #handleMultichainAccountWalletGroupCreatedOrUpdated( + multichainGroup: MultichainAccountGroup>, + ): void { + const bip44GroupAccounts = + this.#getBip44AccountIdsFromMultichainAccountGroup(multichainGroup); + + const group = this.#getAccountGroup(multichainGroup.id); + if (group) { + this.#removeAccounts( + group.accounts.filter( + // We only check if some accounts are no longer part of the multichain + // account group. This is required mainly with the `AccountProviderWrapper`s that + // can be disabled when "Basic functionality" is turned OFF or when the wrappers + // are controlled by remote feature flags. + (account) => !bip44GroupAccounts.has(account), + ), + ); + } + } + /** * Gets account group object. * diff --git a/packages/account-tree-controller/src/types.ts b/packages/account-tree-controller/src/types.ts index da288a1b565..89e38afc911 100644 --- a/packages/account-tree-controller/src/types.ts +++ b/packages/account-tree-controller/src/types.ts @@ -39,7 +39,12 @@ import type { AccountWalletObject, AccountTreeWalletPersistedMetadata, } from './wallet'; -import type { MultichainAccountServiceWalletStatusChangeEvent } from '../../multichain-account-service/src/types'; +import type { + MultichainAccountServiceGetMultichainAccountWalletsAction, + MultichainAccountServiceMultichainAccountGroupCreatedEvent, + MultichainAccountServiceMultichainAccountGroupUpdatedEvent, + MultichainAccountServiceWalletStatusChangeEvent, +} from '../../multichain-account-service/src/types'; // Backward compatibility aliases using indexed access types /** @@ -127,7 +132,8 @@ export type AllowedActions = | UserStorageController.UserStorageControllerPerformSetStorage | UserStorageController.UserStorageControllerPerformBatchSetStorage | AuthenticationController.AuthenticationControllerGetSessionProfile - | MultichainAccountServiceCreateMultichainAccountGroupAction; + | MultichainAccountServiceCreateMultichainAccountGroupAction + | MultichainAccountServiceGetMultichainAccountWalletsAction; export type AccountTreeControllerActions = | AccountTreeControllerGetStateAction @@ -167,7 +173,9 @@ export type AllowedEvents = | AccountsControllerAccountRemovedEvent | AccountsControllerSelectedAccountChangeEvent | UserStorageController.UserStorageControllerStateChangeEvent - | MultichainAccountServiceWalletStatusChangeEvent; + | MultichainAccountServiceWalletStatusChangeEvent + | MultichainAccountServiceMultichainAccountGroupCreatedEvent + | MultichainAccountServiceMultichainAccountGroupUpdatedEvent; export type AccountTreeControllerEvents = | AccountTreeControllerStateChangeEvent diff --git a/packages/account-tree-controller/tests/mockMessenger.ts b/packages/account-tree-controller/tests/mockMessenger.ts index 10fa770a675..4dfa8081d3f 100644 --- a/packages/account-tree-controller/tests/mockMessenger.ts +++ b/packages/account-tree-controller/tests/mockMessenger.ts @@ -50,6 +50,8 @@ export function getAccountTreeControllerMessenger( 'AccountsController:selectedAccountChange', 'UserStorageController:stateChange', 'MultichainAccountService:walletStatusChange', + 'MultichainAccountService:multichainAccountGroupCreated', + 'MultichainAccountService:multichainAccountGroupUpdated', ], actions: [ 'AccountsController:listMultichainAccounts', @@ -63,6 +65,7 @@ export function getAccountTreeControllerMessenger( 'UserStorageController:performBatchSetStorage', 'AuthenticationController:getSessionProfile', 'MultichainAccountService:createMultichainAccountGroup', + 'MultichainAccountService:getMultichainAccountWallets', 'KeyringController:getState', 'SnapController:get', ],