diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 89b5e450cdd..7694f4e6f06 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -9,6 +9,7 @@ ## Accounts Team /packages/accounts-controller @MetaMask/accounts-engineers /packages/multichain-transactions-controller @MetaMask/accounts-engineers +/packages/account-tree-controller @MetaMask/accounts-engineers ## Assets Team /packages/assets-controllers @MetaMask/metamask-assets @@ -88,6 +89,8 @@ /packages/foundryup @MetaMask/mobile-platform @MetaMask/extension-platform ## Package Release related +/packages/account-tree-controller/package.json @MetaMask/accounts-engineers @MetaMask/wallet-framework-engineers +/packages/account-tree-controller/CHANGELOG.md @MetaMask/accounts-engineers @MetaMask/wallet-framework-engineers /packages/accounts-controller/package.json @MetaMask/accounts-engineers @MetaMask/wallet-framework-engineers /packages/accounts-controller/CHANGELOG.md @MetaMask/accounts-engineers @MetaMask/wallet-framework-engineers /packages/address-book-controller/package.json @MetaMask/confirmations @MetaMask/wallet-framework-engineers diff --git a/README.md b/README.md index bd069d45383..e6289249a33 100644 --- a/README.md +++ b/README.md @@ -20,6 +20,7 @@ Each package in this repository has its own README where you can find installati +- [`@metamask/account-tree-controller`](packages/account-tree-controller) - [`@metamask/accounts-controller`](packages/accounts-controller) - [`@metamask/address-book-controller`](packages/address-book-controller) - [`@metamask/announcement-controller`](packages/announcement-controller) @@ -78,6 +79,7 @@ Each package in this repository has its own README where you can find installati %%{ init: { 'flowchart': { 'curve': 'bumpX' } } }%% graph LR; linkStyle default opacity:0.5 + account_wallet_controller(["@metamask/account-tree-controller"]); accounts_controller(["@metamask/accounts-controller"]); address_book_controller(["@metamask/address-book-controller"]); announcement_controller(["@metamask/announcement-controller"]); @@ -127,6 +129,9 @@ linkStyle default opacity:0.5 token_search_discovery_controller(["@metamask/token-search-discovery-controller"]); transaction_controller(["@metamask/transaction-controller"]); user_operation_controller(["@metamask/user-operation-controller"]); + account_wallet_controller --> base_controller; + account_wallet_controller --> accounts_controller; + account_wallet_controller --> keyring_controller; accounts_controller --> base_controller; accounts_controller --> keyring_controller; accounts_controller --> network_controller; diff --git a/packages/account-tree-controller/CHANGELOG.md b/packages/account-tree-controller/CHANGELOG.md new file mode 100644 index 00000000000..b518709c7b8 --- /dev/null +++ b/packages/account-tree-controller/CHANGELOG.md @@ -0,0 +1,10 @@ +# Changelog + +All notable changes to this project will be documented in this file. + +The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), +and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). + +## [Unreleased] + +[Unreleased]: https://github.com/MetaMask/core/ diff --git a/packages/account-tree-controller/LICENSE b/packages/account-tree-controller/LICENSE new file mode 100644 index 00000000000..7d002dced3a --- /dev/null +++ b/packages/account-tree-controller/LICENSE @@ -0,0 +1,20 @@ +MIT License + +Copyright (c) 2025 MetaMask + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE diff --git a/packages/account-tree-controller/README.md b/packages/account-tree-controller/README.md new file mode 100644 index 00000000000..929fa1e14de --- /dev/null +++ b/packages/account-tree-controller/README.md @@ -0,0 +1,15 @@ +# `@metamask/account-tree-controller` + +Manages account wallets according to pre-defined grouping rules (entropy source, Snap IDs, keyring types) and organize wallets/groups of accounts in a tree structure. + +## Installation + +`yarn add @metamask/account-tree-controller` + +or + +`npm install @metamask/account-tree-controller` + +## Contributing + +This package is part of a monorepo. Instructions for contributing can be found in the [monorepo README](https://github.com/MetaMask/core#readme). diff --git a/packages/account-tree-controller/jest.config.js b/packages/account-tree-controller/jest.config.js new file mode 100644 index 00000000000..ca084133399 --- /dev/null +++ b/packages/account-tree-controller/jest.config.js @@ -0,0 +1,26 @@ +/* + * For a detailed explanation regarding each configuration property and type check, visit: + * https://jestjs.io/docs/configuration + */ + +const merge = require('deepmerge'); +const path = require('path'); + +const baseConfig = require('../../jest.config.packages'); + +const displayName = path.basename(__dirname); + +module.exports = merge(baseConfig, { + // The display name when running multiple projects + displayName, + + // An object that configures minimum threshold enforcement for coverage results + coverageThreshold: { + global: { + branches: 100, + functions: 100, + lines: 100, + statements: 100, + }, + }, +}); diff --git a/packages/account-tree-controller/package.json b/packages/account-tree-controller/package.json new file mode 100644 index 00000000000..efa20fe0986 --- /dev/null +++ b/packages/account-tree-controller/package.json @@ -0,0 +1,85 @@ +{ + "name": "@metamask/account-tree-controller", + "version": "0.0.0", + "description": "Controller to group account together based on some pre-defined rules", + "keywords": [ + "MetaMask", + "Ethereum" + ], + "homepage": "https://github.com/MetaMask/core/tree/main/packages/account-tree-controller#readme", + "bugs": { + "url": "https://github.com/MetaMask/core/issues" + }, + "repository": { + "type": "git", + "url": "https://github.com/MetaMask/core.git" + }, + "license": "MIT", + "sideEffects": false, + "exports": { + ".": { + "import": { + "types": "./dist/index.d.mts", + "default": "./dist/index.mjs" + }, + "require": { + "types": "./dist/index.d.cts", + "default": "./dist/index.cjs" + } + }, + "./package.json": "./package.json" + }, + "main": "./dist/index.cjs", + "types": "./dist/index.d.cts", + "files": [ + "dist/" + ], + "scripts": { + "build": "ts-bridge --project tsconfig.build.json --verbose --clean --no-references", + "build:docs": "typedoc", + "changelog:update": "../../scripts/update-changelog.sh @metamask/account-tree-controller", + "changelog:validate": "../../scripts/validate-changelog.sh @metamask/account-tree-controller", + "publish:preview": "yarn npm publish --tag preview", + "since-latest-release": "../../scripts/since-latest-release.sh", + "test": "NODE_OPTIONS=--experimental-vm-modules jest --reporters=jest-silent-reporter", + "test:clean": "NODE_OPTIONS=--experimental-vm-modules jest --clearCache", + "test:verbose": "NODE_OPTIONS=--experimental-vm-modules jest --verbose", + "test:watch": "NODE_OPTIONS=--experimental-vm-modules jest --watch" + }, + "dependencies": { + "@metamask/base-controller": "^8.0.1", + "@metamask/snaps-sdk": "^7.1.0", + "@metamask/snaps-utils": "^9.4.0", + "lodash": "^4.17.21" + }, + "devDependencies": { + "@metamask/accounts-controller": "^30.0.0", + "@metamask/auto-changelog": "^3.4.4", + "@metamask/keyring-api": "^18.0.0", + "@metamask/keyring-controller": "^22.0.1", + "@metamask/providers": "^22.1.0", + "@metamask/snaps-controllers": "^12.3.1", + "@types/jest": "^27.4.1", + "deepmerge": "^4.2.2", + "jest": "^27.5.1", + "ts-jest": "^27.1.4", + "typedoc": "^0.24.8", + "typedoc-plugin-missing-exports": "^2.0.0", + "typescript": "~5.2.2", + "webextension-polyfill": "^0.12.0" + }, + "peerDependencies": { + "@metamask/accounts-controller": "^30.0.0", + "@metamask/keyring-controller": "^22.0.0", + "@metamask/providers": "^22.0.0", + "@metamask/snaps-controllers": "^12.0.0", + "webextension-polyfill": "^0.10.0 || ^0.11.0 || ^0.12.0" + }, + "engines": { + "node": "^18.18 || >=20" + }, + "publishConfig": { + "access": "public", + "registry": "https://registry.npmjs.org/" + } +} diff --git a/packages/account-tree-controller/src/AccountTreeController.test.ts b/packages/account-tree-controller/src/AccountTreeController.test.ts new file mode 100644 index 00000000000..71e48b07c67 --- /dev/null +++ b/packages/account-tree-controller/src/AccountTreeController.test.ts @@ -0,0 +1,576 @@ +import { Messenger } from '@metamask/base-controller'; +import { + EthAccountType, + EthMethod, + EthScope, + SolAccountType, + SolScope, +} from '@metamask/keyring-api'; +import { KeyringTypes } from '@metamask/keyring-controller'; +import type { InternalAccount } from '@metamask/keyring-internal-api'; +import type { GetSnap as SnapControllerGetSnap } from '@metamask/snaps-controllers'; + +import { + AccountTreeController, + AccountWalletCategory, + type AccountTreeControllerMessenger, + type AccountTreeControllerActions, + type AccountTreeControllerEvents, + type AccountTreeControllerState, + type AllowedActions, + type AllowedEvents, + type AccountGroupMetadata, + toDefaultAccountGroupId, + DEFAULT_ACCOUNT_GROUP_NAME, + toAccountWalletId, +} from './AccountTreeController'; +import { getAccountWalletNameFromKeyringType } from './names'; + +const ETH_EOA_METHODS = [ + EthMethod.PersonalSign, + EthMethod.Sign, + EthMethod.SignTransaction, + EthMethod.SignTypedDataV1, + EthMethod.SignTypedDataV3, + EthMethod.SignTypedDataV4, +] as const; + +const MOCK_SNAP_1 = { + id: 'local:mock-snap-id-1', + name: 'Mock Snap 1', + enabled: true, + manifest: { + proposedName: 'Mock Snap 1', + }, +}; + +const MOCK_SNAP_2 = { + id: 'local:mock-snap-id-2', + name: 'Mock Snap 2', + enabled: true, + manifest: { + proposedName: 'Mock Snap 2', + }, +}; + +const MOCK_HD_KEYRING_1 = { + type: KeyringTypes.hd, + metadata: { id: 'mock-keyring-id-1', name: 'HD Keyring 1' }, + accounts: ['0x123'], +}; + +const MOCK_HD_KEYRING_2 = { + type: KeyringTypes.hd, + metadata: { id: 'mock-keyring-id-2', name: 'HD Keyring 1' }, + accounts: ['0x456'], +}; + +const MOCK_HD_ACCOUNT_1: InternalAccount = { + id: 'mock-id-1', + address: '0x123', + options: { entropySource: MOCK_HD_KEYRING_1.metadata.id }, + methods: [...ETH_EOA_METHODS], + type: EthAccountType.Eoa, + scopes: [EthScope.Eoa], + metadata: { + name: 'Account 1', + keyring: { type: KeyringTypes.hd }, + importTime: 0, + lastSelected: 0, + nameLastUpdatedAt: 0, + }, +}; + +const MOCK_HD_ACCOUNT_2: InternalAccount = { + id: 'mock-id-2', + address: '0x456', + options: { entropySource: MOCK_HD_KEYRING_2.metadata.id }, + methods: [...ETH_EOA_METHODS], + type: EthAccountType.Eoa, + scopes: [EthScope.Eoa], + metadata: { + name: 'Account 2', + keyring: { type: KeyringTypes.hd }, + importTime: 0, + lastSelected: 0, + nameLastUpdatedAt: 0, + }, +}; + +const MOCK_SNAP_ACCOUNT_1: InternalAccount = { + id: 'mock-snap-id-1', + address: 'aabbccdd', + options: { entropySource: MOCK_HD_KEYRING_2.metadata.id }, // Note: shares entropy with MOCK_HD_ACCOUNT_2 + methods: [...ETH_EOA_METHODS], + type: SolAccountType.DataAccount, + scopes: [SolScope.Mainnet], + metadata: { + name: 'Snap Acc 1', + keyring: { type: KeyringTypes.snap }, + snap: MOCK_SNAP_1, + importTime: 0, + lastSelected: 0, + }, +}; + +const MOCK_SNAP_ACCOUNT_2: InternalAccount = { + id: 'mock-snap-id-2', + address: '0x789', + options: {}, + methods: [...ETH_EOA_METHODS], + type: EthAccountType.Eoa, + scopes: [EthScope.Eoa], + metadata: { + name: 'Snap Acc 2', + keyring: { type: KeyringTypes.snap }, + snap: MOCK_SNAP_2, + importTime: 0, + lastSelected: 0, + }, +}; + +const MOCK_HARDWARE_ACCOUNT_1: InternalAccount = { + id: 'mock-hardware-id-1', + address: '0xABC', + options: {}, + methods: [...ETH_EOA_METHODS], + type: EthAccountType.Eoa, + scopes: [EthScope.Eoa], + metadata: { + name: 'Hardware Acc 1', + keyring: { type: KeyringTypes.ledger }, + importTime: 0, + lastSelected: 0, + }, +}; + +/** + * Creates a new root messenger instance for testing. + * + * @returns A new Messenger instance. + */ +function getRootMessenger() { + return new Messenger< + AccountTreeControllerActions | AllowedActions, + AccountTreeControllerEvents | AllowedEvents + >(); +} + +/** + * Retrieves a restricted messenger for the AccountTreeController. + * + * @param messenger - The root messenger instance. Defaults to a new Messenger created by getRootMessenger(). + * @returns The restricted messenger for the AccountTreeController. + */ +function getAccountTreeControllerMessenger( + messenger = getRootMessenger(), +): AccountTreeControllerMessenger { + return messenger.getRestricted({ + name: 'AccountTreeController', + allowedEvents: [ + 'AccountsController:accountAdded', + 'AccountsController:accountRemoved', + ], + allowedActions: [ + 'AccountsController:listMultichainAccounts', + 'KeyringController:getState', + 'SnapController:get', + ], + }); +} + +/** + * Sets up the AccountTreeController for testing. + * + * @param options - Configuration options for setup. + * @param options.state - Partial initial state for the controller. Defaults to empty object. + * @param options.messenger - An optional messenger instance to use. Defaults to a new Messenger. + * @returns An object containing the controller instance and the messenger. + */ +function setup({ + state = {}, + messenger = getRootMessenger(), +}: { + state?: Partial; + messenger?: Messenger< + AccountTreeControllerActions | AllowedActions, + AccountTreeControllerEvents | AllowedEvents + >; +} = {}): { + controller: AccountTreeController; + messenger: Messenger< + AccountTreeControllerActions | AllowedActions, + AccountTreeControllerEvents | AllowedEvents + >; +} { + const controller = new AccountTreeController({ + messenger: getAccountTreeControllerMessenger(messenger), + state, + }); + return { controller, messenger }; +} + +describe('AccountTreeController', () => { + describe('init', () => { + it('groups accounts by entropy source, then snapId, then wallet type', () => { + const { controller, messenger } = setup(); + messenger.registerActionHandler( + 'AccountsController:listMultichainAccounts', + () => [ + MOCK_HD_ACCOUNT_1, + MOCK_HD_ACCOUNT_2, + MOCK_SNAP_ACCOUNT_1, // Belongs to MOCK_HD_ACCOUNT_2's wallet due to shared entropySource + MOCK_SNAP_ACCOUNT_2, // Has its own Snap wallet + MOCK_HARDWARE_ACCOUNT_1, // Has its own Keyring wallet + ], + ); + messenger.registerActionHandler('KeyringController:getState', () => ({ + isUnlocked: true, + keyrings: [MOCK_HD_KEYRING_1, MOCK_HD_KEYRING_2], + })); + messenger.registerActionHandler( + 'SnapController:get', + () => + // TODO: Update this to avoid the unknown cast if possible. + MOCK_SNAP_1 as unknown as ReturnType< + SnapControllerGetSnap['handler'] + >, + ); + + controller.init(); + + const expectedWalletId1 = toAccountWalletId( + AccountWalletCategory.Entropy, + MOCK_HD_KEYRING_1.metadata.id, + ); + const expectedWalletId1Group = toDefaultAccountGroupId(expectedWalletId1); + const expectedWalletId2 = toAccountWalletId( + AccountWalletCategory.Entropy, + MOCK_HD_KEYRING_2.metadata.id, + ); + const expectedWalletId2Group = toDefaultAccountGroupId(expectedWalletId2); + const expectedSnapWalletId = toAccountWalletId( + AccountWalletCategory.Snap, + MOCK_SNAP_2.id, + ); + const expectedSnapWalletIdGroup = + toDefaultAccountGroupId(expectedSnapWalletId); + const expectedKeyringWalletId = `${AccountWalletCategory.Keyring}:${KeyringTypes.ledger}`; + const expectedKeyringWalletIdGroup = toDefaultAccountGroupId( + expectedKeyringWalletId, + ); + + const mockDefaultGroupMetadata: AccountGroupMetadata = { + name: DEFAULT_ACCOUNT_GROUP_NAME, + }; + + expect(controller.state).toStrictEqual({ + accountTree: { + wallets: { + [expectedWalletId1]: { + id: expectedWalletId1, + groups: { + [expectedWalletId1Group]: { + id: expectedWalletId1Group, + accounts: [MOCK_HD_ACCOUNT_1.id], + metadata: mockDefaultGroupMetadata, + }, + }, + metadata: { name: 'Wallet 1' }, + }, + [expectedWalletId2]: { + id: expectedWalletId2, + groups: { + [expectedWalletId2Group]: { + id: expectedWalletId2Group, + accounts: [MOCK_HD_ACCOUNT_2.id, MOCK_SNAP_ACCOUNT_1.id], + metadata: mockDefaultGroupMetadata, + }, + }, + metadata: { name: 'Wallet 2' }, + }, + [expectedSnapWalletId]: { + id: expectedSnapWalletId, + groups: { + [expectedSnapWalletIdGroup]: { + id: expectedSnapWalletIdGroup, + accounts: [MOCK_SNAP_ACCOUNT_2.id], + metadata: mockDefaultGroupMetadata, + }, + }, + metadata: { name: `Snap: ${MOCK_SNAP_1.manifest.proposedName}` }, + }, + [expectedKeyringWalletId]: { + id: expectedKeyringWalletId, + groups: { + [expectedKeyringWalletIdGroup]: { + id: expectedKeyringWalletIdGroup, + accounts: [MOCK_HARDWARE_ACCOUNT_1.id], + metadata: mockDefaultGroupMetadata, + }, + }, + metadata: { + name: getAccountWalletNameFromKeyringType( + MOCK_HARDWARE_ACCOUNT_1.metadata.keyring.type as KeyringTypes, + ), + }, + }, + }, + }, + } as AccountTreeControllerState); + }); + + it('warns and fall back to wallet type grouping if an HD account is missing entropySource', () => { + const consoleWarnSpy = jest + .spyOn(console, 'warn') + .mockImplementation(() => undefined); + const { controller, messenger } = setup(); + const mockHdAccountWithoutEntropy: InternalAccount = { + ...MOCK_HD_ACCOUNT_1, + id: 'mock-no-entropy-id', + options: {}, + }; + messenger.registerActionHandler( + 'AccountsController:listMultichainAccounts', + () => [mockHdAccountWithoutEntropy], + ); + messenger.registerActionHandler('KeyringController:getState', () => ({ + isUnlocked: true, + keyrings: [], + })); + + controller.init(); + + expect(consoleWarnSpy).toHaveBeenCalledWith( + "! Found an HD account with no entropy source: account won't be associated to its wallet", + ); + const expectedKeyringWalletId = toAccountWalletId( + AccountWalletCategory.Keyring, + KeyringTypes.hd, + ); + const expectedGroupId = toDefaultAccountGroupId(expectedKeyringWalletId); + expect( + controller.state.accountTree.wallets[expectedKeyringWalletId]?.groups[ + expectedGroupId + ]?.accounts, + ).toContain(mockHdAccountWithoutEntropy.id); + consoleWarnSpy.mockRestore(); + }); + + it('handles Snap accounts with entropy source', () => { + const { controller, messenger } = setup(); + const mockSnapAccountWithEntropy: InternalAccount = { + ...MOCK_SNAP_ACCOUNT_2, + options: { entropySource: MOCK_HD_KEYRING_2.metadata.id }, + metadata: { + ...MOCK_SNAP_ACCOUNT_2.metadata, + snap: MOCK_SNAP_2, + }, + }; + messenger.registerActionHandler( + 'AccountsController:listMultichainAccounts', + () => [mockSnapAccountWithEntropy], + ); + messenger.registerActionHandler('KeyringController:getState', () => ({ + isUnlocked: true, + keyrings: [MOCK_HD_KEYRING_2], + })); + + controller.init(); + + const expectedWalletId = toAccountWalletId( + AccountWalletCategory.Entropy, + MOCK_HD_KEYRING_2.metadata.id, + ); + const expectedGroupId = toDefaultAccountGroupId(expectedWalletId); + expect( + controller.state.accountTree.wallets[expectedWalletId]?.groups[ + expectedGroupId + ]?.accounts, + ).toContain(mockSnapAccountWithEntropy.id); + }); + + it('fallback to Snap ID if Snap cannot be found', () => { + const { controller, messenger } = setup(); + messenger.registerActionHandler( + 'AccountsController:listMultichainAccounts', + () => [MOCK_SNAP_ACCOUNT_1], + ); + messenger.registerActionHandler('KeyringController:getState', () => ({ + isUnlocked: true, + keyrings: [], + })); + messenger.registerActionHandler('SnapController:get', () => undefined); // Snap won't be found. + + controller.init(); + + // Since no entropy sources will be found, it will be categorized as a + // "Keyring" wallet + const wallet1Id = toAccountWalletId( + AccountWalletCategory.Snap, + MOCK_SNAP_1.id, + ); + + // FIXME: Do we really want this behavior? + expect( + controller.state.accountTree.wallets[wallet1Id]?.metadata.name, + ).toBe('Snap: mock-snap-id-1'); + }); + + it('fallback to HD keyring category if entropy sources cannot be found', () => { + const { controller, messenger } = setup(); + // Create entropy wallets that will both get "Wallet" as base name, then get numbered + const mockHdAccount1: InternalAccount = { + ...MOCK_HD_ACCOUNT_1, + options: { entropySource: MOCK_HD_KEYRING_1.metadata.id }, + }; + const mockHdAccount2: InternalAccount = { + ...MOCK_HD_ACCOUNT_2, + options: { entropySource: MOCK_HD_KEYRING_2.metadata.id }, + }; + messenger.registerActionHandler( + 'AccountsController:listMultichainAccounts', + () => [mockHdAccount1, mockHdAccount2], + ); + messenger.registerActionHandler('KeyringController:getState', () => ({ + isUnlocked: true, + keyrings: [], // Entropy sources won't be found. + })); + + controller.init(); + + // Since no entropy sources will be found, it will be categorized as a + // "Keyring" wallet + const wallet1Id = toAccountWalletId( + AccountWalletCategory.Keyring, + mockHdAccount1.metadata.keyring.type, + ); + const wallet2Id = toAccountWalletId( + AccountWalletCategory.Keyring, + mockHdAccount1.metadata.keyring.type, + ); + + // FIXME: Do we really want this behavior? + expect( + controller.state.accountTree.wallets[wallet1Id]?.metadata.name, + ).toBe('HD Wallet'); + expect( + controller.state.accountTree.wallets[wallet2Id]?.metadata.name, + ).toBe('HD Wallet'); + }); + }); + + describe('on AccountsController:accountRemoved', () => { + it('removes an account from the tree', () => { + const { controller, messenger } = setup(); + // + // 2 accounts that share the same entropy source (thus, same wallet). + const mockHdAccount1 = { + ...MOCK_HD_ACCOUNT_1, + options: { + entropySource: MOCK_HD_KEYRING_1.metadata.id, + }, + }; + const mockHdAccount2 = { + ...MOCK_HD_ACCOUNT_2, + options: { + entropySource: MOCK_HD_KEYRING_1.metadata.id, + }, + }; + + // Create entropy wallets that will both get "Wallet" as base name, then get numbered + messenger.registerActionHandler( + 'AccountsController:listMultichainAccounts', + () => [mockHdAccount1, mockHdAccount2], + ); + messenger.registerActionHandler('KeyringController:getState', () => ({ + isUnlocked: true, + keyrings: [MOCK_HD_KEYRING_1], + })); + + controller.init(); + + messenger.publish('AccountsController:accountRemoved', mockHdAccount1.id); + + const walletId1 = toAccountWalletId( + AccountWalletCategory.Entropy, + MOCK_HD_KEYRING_1.metadata.id, + ); + const walletId1Group = toDefaultAccountGroupId(walletId1); + expect(controller.state).toStrictEqual({ + accountTree: { + wallets: { + [walletId1]: { + id: walletId1, + groups: { + [walletId1Group]: { + id: walletId1Group, + metadata: { name: DEFAULT_ACCOUNT_GROUP_NAME }, + accounts: [mockHdAccount2.id], // HD account 1 got removed. + }, + }, + metadata: { name: 'Wallet 1' }, + }, + }, + }, + } as AccountTreeControllerState); + }); + }); + + describe('on AccountsController:accountAdded', () => { + it('adds an account from the tree', () => { + const { controller, messenger } = setup(); + // + // 2 accounts that share the same entropy source (thus, same wallet). + const mockHdAccount1 = { + ...MOCK_HD_ACCOUNT_1, + options: { + entropySource: MOCK_HD_KEYRING_1.metadata.id, + }, + }; + const mockHdAccount2 = { + ...MOCK_HD_ACCOUNT_2, + options: { + entropySource: MOCK_HD_KEYRING_1.metadata.id, + }, + }; + + // Create entropy wallets that will both get "Wallet" as base name, then get numbered + messenger.registerActionHandler( + 'AccountsController:listMultichainAccounts', + () => [mockHdAccount1], + ); + messenger.registerActionHandler('KeyringController:getState', () => ({ + isUnlocked: true, + keyrings: [MOCK_HD_KEYRING_1], + })); + + controller.init(); + + messenger.publish('AccountsController:accountAdded', mockHdAccount2); + + const walletId1 = toAccountWalletId( + AccountWalletCategory.Entropy, + MOCK_HD_KEYRING_1.metadata.id, + ); + const walletId1Group = toDefaultAccountGroupId(walletId1); + expect(controller.state).toStrictEqual({ + accountTree: { + wallets: { + [walletId1]: { + id: walletId1, + groups: { + [walletId1Group]: { + id: walletId1Group, + metadata: { name: DEFAULT_ACCOUNT_GROUP_NAME }, + accounts: [mockHdAccount1.id, mockHdAccount2.id], // HD account 2 got added. + }, + }, + metadata: { name: 'Wallet 1' }, + }, + }, + }, + } as AccountTreeControllerState); + }); + }); +}); diff --git a/packages/account-tree-controller/src/AccountTreeController.ts b/packages/account-tree-controller/src/AccountTreeController.ts new file mode 100644 index 00000000000..6ab241a4c79 --- /dev/null +++ b/packages/account-tree-controller/src/AccountTreeController.ts @@ -0,0 +1,438 @@ +import type { + AccountId, + AccountsControllerAccountAddedEvent, + AccountsControllerAccountRemovedEvent, + AccountsControllerListMultichainAccountsAction, +} from '@metamask/accounts-controller'; +import type { StateMetadata } from '@metamask/base-controller'; +import { + type ControllerGetStateAction, + type ControllerStateChangeEvent, + type RestrictedMessenger, + BaseController, +} from '@metamask/base-controller'; +import type { KeyringControllerGetStateAction } from '@metamask/keyring-controller'; +import { KeyringTypes } from '@metamask/keyring-controller'; +import type { InternalAccount } from '@metamask/keyring-internal-api'; +import type { GetSnap as SnapControllerGetSnap } from '@metamask/snaps-controllers'; +import type { SnapId } from '@metamask/snaps-sdk'; +import { stripSnapPrefix } from '@metamask/snaps-utils'; + +import { getAccountWalletNameFromKeyringType } from './names'; + +const controllerName = 'AccountTreeController'; + +export enum AccountWalletCategory { + Entropy = 'entropy', + Keyring = 'keyring', + Snap = 'snap', +} + +type AccountTreeRuleMatch = { + category: AccountWalletCategory; + id: AccountWalletId; + name: string; +}; + +type AccountTreeRuleFunction = ( + account: InternalAccount, +) => AccountTreeRuleMatch | undefined; + +type AccountReverseMapping = { + walletId: AccountWalletId; + groupId: AccountGroupId; +}; + +export type AccountWalletId = `${AccountWalletCategory}:${string}`; +export type AccountGroupId = `${AccountWalletId}:${string}`; + +// Do not export this one, we just use it to have a common type interface between group and wallet metadata. +type Metadata = { + name: string; +}; + +export type AccountWalletMetadata = Metadata; + +export type AccountGroupMetadata = Metadata; + +export type AccountGroup = { + id: AccountGroupId; + // Blockchain Accounts: + accounts: AccountId[]; + metadata: AccountGroupMetadata; +}; + +export type AccountWallet = { + id: AccountWalletId; + // Account groups OR Multichain accounts (once available). + groups: { + [groupId: AccountGroupId]: AccountGroup; + }; + metadata: AccountGroupMetadata; // Assuming Metadata is a defined type +}; + +export type AccountTreeControllerState = { + accountTree: { + wallets: { + // Wallets: + [walletId: AccountWalletId]: AccountWallet; + }; + }; +}; + +export type AccountTreeControllerGetStateAction = ControllerGetStateAction< + typeof controllerName, + AccountTreeControllerState +>; + +export type AllowedActions = + | AccountsControllerListMultichainAccountsAction + | KeyringControllerGetStateAction + | SnapControllerGetSnap; + +export type AccountTreeControllerActions = never; + +export type AccountTreeControllerStateChangeEvent = ControllerStateChangeEvent< + typeof controllerName, + AccountTreeControllerState +>; + +export type AllowedEvents = + | AccountsControllerAccountAddedEvent + | AccountsControllerAccountRemovedEvent; + +export type AccountTreeControllerEvents = never; + +export type AccountTreeControllerMessenger = RestrictedMessenger< + typeof controllerName, + AccountTreeControllerActions | AllowedActions, + AccountTreeControllerEvents | AllowedEvents, + AllowedActions['type'], + AllowedEvents['type'] +>; + +const accountTreeControllerMetadata: StateMetadata = + { + accountTree: { + persist: false, // We do re-recompute this state everytime. + anonymous: false, + }, + }; + +/** + * Gets default state of the `AccountTreeController`. + * + * @returns The default state of the `AccountTreeController`. + */ +export function getDefaultAccountTreeControllerState(): AccountTreeControllerState { + return { + accountTree: { + wallets: {}, + }, + }; +} + +// TODO: For now we use this for the 2nd-level of the tree until we implements proper multichain accounts. +export const DEFAULT_ACCOUNT_GROUP_UNIQUE_ID: string = 'default'; // This might need to be re-evaluated based on new structure +export const DEFAULT_ACCOUNT_GROUP_NAME: string = 'Default'; + +/** + * Convert a unique ID to a wallet ID for a given category. + * + * @param category - A wallet category. + * @param id - A unique ID. + * @returns A wallet ID. + */ +export function toAccountWalletId( + category: AccountWalletCategory, + id: string, +): AccountWalletId { + return `${category}:${id}`; +} + +/** + * Convert a wallet ID and a unique ID, to a group ID. + * + * @param walletId - A wallet ID. + * @param id - A unique ID. + * @returns A group ID. + */ +export function toAccountGroupId( + walletId: AccountWalletId, + id: string, +): AccountGroupId { + return `${walletId}:${id}`; +} + +/** + * Convert a wallet ID to the default group ID. + * + * @param walletId - A wallet ID. + * @returns The default group ID. + */ +export function toDefaultAccountGroupId( + walletId: AccountWalletId, +): AccountGroupId { + return toAccountGroupId(walletId, DEFAULT_ACCOUNT_GROUP_UNIQUE_ID); +} + +export class AccountTreeController extends BaseController< + typeof controllerName, + AccountTreeControllerState, + AccountTreeControllerMessenger +> { + readonly #reverse: Map; + + readonly #rules: AccountTreeRuleFunction[]; + + /** + * Constructor for AccountTreeController. + * + * @param options - The controller options. + * @param options.messenger - The messenger object. + * @param options.state - Initial state to set on this controller + */ + constructor({ + messenger, + state, + }: { + messenger: AccountTreeControllerMessenger; + state?: Partial; + }) { + super({ + messenger, + name: controllerName, + metadata: accountTreeControllerMetadata, + state: { + ...getDefaultAccountTreeControllerState(), + ...state, + }, + }); + + // Reverse map to allow fast node access from an account ID. + this.#reverse = new Map(); + + // Rules to apply to construct the wallets tree. + this.#rules = [ + // 1. We group by entropy-source + (account: InternalAccount) => this.#matchGroupByEntropySource(account), + // 2. We group by Snap ID + (account: InternalAccount) => this.#matchGroupBySnapId(account), + // 3. We group by wallet type (this rule cannot fail and will group all non-matching accounts) + (account: InternalAccount) => this.#matchGroupByKeyringType(account), + ]; + + this.messagingSystem.subscribe( + 'AccountsController:accountAdded', + (account) => { + this.#handleAccountAdded(account); + }, + ); + + this.messagingSystem.subscribe( + 'AccountsController:accountRemoved', + (accountId) => { + this.#handleAccountRemoved(accountId); + }, + ); + } + + init() { + const wallets = {}; + + // For now, we always re-compute all wallets, we do not re-use the existing state. + for (const account of this.#listAccounts()) { + this.#insert(wallets, account); + } + + this.update((state) => { + state.accountTree.wallets = wallets; + }); + } + + #handleAccountAdded(account: InternalAccount) { + this.update((state) => { + this.#insert(state.accountTree.wallets, account); + }); + } + + #handleAccountRemoved(accountId: AccountId) { + const found = this.#reverse.get(accountId); + + if (found) { + const { walletId, groupId } = found; + this.update((state) => { + const { accounts } = + state.accountTree.wallets[walletId].groups[groupId]; + + const index = accounts.indexOf(accountId); + if (index !== -1) { + accounts.splice(index, 1); + } + }); + } + } + + #hasKeyringType(account: InternalAccount, type: KeyringTypes): boolean { + return account.metadata.keyring.type === (type as string); + } + + #matchGroupByEntropySource( + account: InternalAccount, + ): AccountTreeRuleMatch | undefined { + let entropySource: string | undefined; + + if (this.#hasKeyringType(account, KeyringTypes.hd)) { + // TODO: Maybe use superstruct to validate the structure of HD account since they are not strongly-typed for now? + if (!account.options.entropySource) { + console.warn( + "! Found an HD account with no entropy source: account won't be associated to its wallet", + ); + return undefined; + } + + entropySource = account.options.entropySource as string; + } + + // TODO: For now, we're not checking if the Snap is a preinstalled one, and we probably should... + if ( + this.#hasKeyringType(account, KeyringTypes.snap) && + account.metadata.snap?.enabled + ) { + // Not all Snaps have an entropy-source and options are not typed yet, so we have to check manually here. + if (account.options.entropySource) { + // We blindly trust the `entropySource` for now, but it could be wrong since it comes from a Snap. + entropySource = account.options.entropySource as string; + } + } + + if (!entropySource) { + return undefined; + } + + // We check if we can get the name for that entropy source, if not this means this entropy does not match + // any HD keyrings, thus, is invalid (this account will be grouped by another rule). + const entropySourceName = this.#getEntropySourceName(entropySource); + if (!entropySourceName) { + console.warn( + '! Tried to name a wallet using an unknown entropy, this should not be possible.', + ); + return undefined; + } + + return { + category: AccountWalletCategory.Entropy, + id: toAccountWalletId(AccountWalletCategory.Entropy, entropySource), + name: entropySourceName, + }; + } + + #matchGroupBySnapId( + account: InternalAccount, + ): AccountTreeRuleMatch | undefined { + if ( + this.#hasKeyringType(account, KeyringTypes.snap) && + account.metadata.snap && + account.metadata.snap.enabled + ) { + const { id } = account.metadata.snap; + + return { + category: AccountWalletCategory.Snap, + id: toAccountWalletId(AccountWalletCategory.Snap, id), + name: this.#getSnapName(id as SnapId), + }; + } + + return undefined; + } + + #matchGroupByKeyringType( + account: InternalAccount, + ): AccountTreeRuleMatch | undefined { + const { type } = account.metadata.keyring; + + return { + category: AccountWalletCategory.Keyring, + id: toAccountWalletId(AccountWalletCategory.Keyring, type), + name: getAccountWalletNameFromKeyringType(type as KeyringTypes), + }; + } + + #getSnapName(snapId: SnapId): string { + const snap = this.messagingSystem.call('SnapController:get', snapId); + const snapName = snap + ? // TODO: Handle localization here, but that's a "client thing", so we don't have a `core` controller + // to refer too. + snap.manifest.proposedName + : stripSnapPrefix(snapId); + + return `Snap: ${snapName}`; + } + + #getEntropySourceName(entropySource: string): string | undefined { + const { keyrings } = this.messagingSystem.call( + 'KeyringController:getState', + ); + + const index = keyrings + .filter((keyring) => keyring.type === (KeyringTypes.hd as string)) + .findIndex((keyring) => keyring.metadata.id === entropySource); + + if (index === -1) { + return undefined; + } + + return `Wallet ${index + 1}`; // Use human indexing. + } + + #insert( + wallets: { [walletId: AccountWalletId]: AccountWallet }, + account: InternalAccount, + ) { + for (const rule of this.#rules) { + const match = rule(account); + + if (!match) { + // No match for that rule, we go to the next one. + continue; + } + + const walletId = match.id; + const walletName = match.name; + const groupId = toDefaultAccountGroupId(walletId); // Use a single-group for now until multichain accounts is supported. + const groupName = DEFAULT_ACCOUNT_GROUP_NAME; + + if (!wallets[walletId]) { + wallets[walletId] = { + id: walletId, + groups: { + [groupId]: { + id: groupId, + accounts: [], + metadata: { name: groupName }, + }, + }, + metadata: { + name: walletName, + }, + }; + } + wallets[walletId].groups[groupId].accounts.push(account.id); + + // Update the reverse mapping for this account. + this.#reverse.set(account.id, { + walletId, + groupId, + }); + + return; + } + } + + #listAccounts(): InternalAccount[] { + return this.messagingSystem.call( + 'AccountsController:listMultichainAccounts', + ); + } +} diff --git a/packages/account-tree-controller/src/index.ts b/packages/account-tree-controller/src/index.ts new file mode 100644 index 00000000000..9e144ad0594 --- /dev/null +++ b/packages/account-tree-controller/src/index.ts @@ -0,0 +1,19 @@ +export type { + AccountTreeControllerState, + AccountTreeControllerGetStateAction, + AccountTreeControllerActions, + AccountTreeControllerStateChangeEvent, + AccountTreeControllerEvents, + AccountTreeControllerMessenger, + AccountWallet, + AccountWalletId, + AccountWalletMetadata, + AccountWalletCategory, + AccountGroup, + AccountGroupId, + AccountGroupMetadata, +} from './AccountTreeController'; +export { + AccountTreeController, + getDefaultAccountTreeControllerState, +} from './AccountTreeController'; diff --git a/packages/account-tree-controller/src/names.test.ts b/packages/account-tree-controller/src/names.test.ts new file mode 100644 index 00000000000..d354a71bc30 --- /dev/null +++ b/packages/account-tree-controller/src/names.test.ts @@ -0,0 +1,26 @@ +import { KeyringTypes } from '@metamask/keyring-controller'; + +import { getAccountWalletNameFromKeyringType } from './names'; + +describe('names', () => { + describe('getWalletNameFromKeyringType', () => { + it.each(Object.values(KeyringTypes))( + 'computes wallet name from: %s', + (type) => { + const name = getAccountWalletNameFromKeyringType(type as KeyringTypes); + + expect(name).toBeDefined(); + expect(name.length).toBeGreaterThan(0); + }, + ); + + it('defaults to "Unknown" if keyring type is not known', () => { + const name = getAccountWalletNameFromKeyringType( + 'Not A Keyring Type' as KeyringTypes, + ); + + expect(name).toBe('Unknown'); + expect(name.length).toBeGreaterThan(0); + }); + }); +}); diff --git a/packages/account-tree-controller/src/names.ts b/packages/account-tree-controller/src/names.ts new file mode 100644 index 00000000000..2a56d70046c --- /dev/null +++ b/packages/account-tree-controller/src/names.ts @@ -0,0 +1,39 @@ +import { KeyringTypes } from '@metamask/keyring-controller'; + +/** + * Get wallet name from a keyring type. + * + * @param type - Keyring's type. + * @returns Wallet name. + */ +export function getAccountWalletNameFromKeyringType(type: KeyringTypes) { + switch (type) { + case KeyringTypes.simple: { + return 'Private Keys'; + } + case KeyringTypes.hd: { + return 'HD Wallet'; + } + case KeyringTypes.trezor: { + return 'Trezor'; + } + case KeyringTypes.oneKey: { + return 'OneKey'; + } + case KeyringTypes.ledger: { + return 'Ledger'; + } + case KeyringTypes.lattice: { + return 'Lattice'; + } + case KeyringTypes.qr: { + return 'QR'; + } + case KeyringTypes.snap: { + return 'Snap Wallet'; + } + default: { + return 'Unknown'; + } + } +} diff --git a/packages/account-tree-controller/tsconfig.build.json b/packages/account-tree-controller/tsconfig.build.json new file mode 100644 index 00000000000..5e3f6b10dd6 --- /dev/null +++ b/packages/account-tree-controller/tsconfig.build.json @@ -0,0 +1,14 @@ +{ + "extends": "../../tsconfig.packages.build.json", + "compilerOptions": { + "baseUrl": "./", + "outDir": "./dist", + "rootDir": "./src" + }, + "references": [ + { "path": "../accounts-controller/tsconfig.build.json" }, + { "path": "../base-controller/tsconfig.build.json" }, + { "path": "../keyring-controller/tsconfig.build.json" } + ], + "include": ["../../types", "./src"] +} diff --git a/packages/account-tree-controller/tsconfig.json b/packages/account-tree-controller/tsconfig.json new file mode 100644 index 00000000000..8b6228af6b8 --- /dev/null +++ b/packages/account-tree-controller/tsconfig.json @@ -0,0 +1,18 @@ +{ + "extends": "../../tsconfig.packages.json", + "compilerOptions": { + "baseUrl": "./" + }, + "references": [ + { + "path": "../base-controller" + }, + { + "path": "../keyring-controller" + }, + { + "path": "../accounts-controller" + } + ], + "include": ["../../types", "./src"] +} diff --git a/packages/account-tree-controller/typedoc.json b/packages/account-tree-controller/typedoc.json new file mode 100644 index 00000000000..c9da015dbf8 --- /dev/null +++ b/packages/account-tree-controller/typedoc.json @@ -0,0 +1,7 @@ +{ + "entryPoints": ["./src/index.ts"], + "excludePrivate": true, + "hideGenerator": true, + "out": "docs", + "tsconfig": "./tsconfig.build.json" +} diff --git a/teams.json b/teams.json index 35be7f59fd5..68c0d879fe2 100644 --- a/teams.json +++ b/teams.json @@ -1,5 +1,6 @@ { "metamask/accounts-controller": "team-accounts", + "metamask/account-tree-controller": "team-accounts", "metamask/address-book-controller": "team-confirmations", "metamask/announcement-controller": "team-wallet-ux", "metamask/app-metadata-controller": "team-mobile-platform", diff --git a/tsconfig.build.json b/tsconfig.build.json index a418d6b71fc..312bbf40103 100644 --- a/tsconfig.build.json +++ b/tsconfig.build.json @@ -1,5 +1,6 @@ { "references": [ + { "path": "./packages/account-tree-controller/tsconfig.build.json" }, { "path": "./packages/accounts-controller/tsconfig.build.json" }, { "path": "./packages/address-book-controller/tsconfig.build.json" }, { "path": "./packages/announcement-controller/tsconfig.build.json" }, diff --git a/tsconfig.json b/tsconfig.json index 5e6b6271b7d..b4268e7d965 100644 --- a/tsconfig.json +++ b/tsconfig.json @@ -6,6 +6,7 @@ "noEmit": true }, "references": [ + { "path": "./packages/account-tree-controller" }, { "path": "./packages/accounts-controller" }, { "path": "./packages/address-book-controller" }, { "path": "./packages/announcement-controller" }, diff --git a/yarn.lock b/yarn.lock index 0747cd79fb4..8f91f81a7a2 100644 --- a/yarn.lock +++ b/yarn.lock @@ -2436,6 +2436,37 @@ __metadata: languageName: node linkType: hard +"@metamask/account-tree-controller@workspace:packages/account-tree-controller": + version: 0.0.0-use.local + resolution: "@metamask/account-tree-controller@workspace:packages/account-tree-controller" + dependencies: + "@metamask/accounts-controller": "npm:^30.0.0" + "@metamask/auto-changelog": "npm:^3.4.4" + "@metamask/base-controller": "npm:^8.0.1" + "@metamask/keyring-api": "npm:^18.0.0" + "@metamask/keyring-controller": "npm:^22.0.1" + "@metamask/providers": "npm:^22.1.0" + "@metamask/snaps-controllers": "npm:^12.3.1" + "@metamask/snaps-sdk": "npm:^7.1.0" + "@metamask/snaps-utils": "npm:^9.4.0" + "@types/jest": "npm:^27.4.1" + deepmerge: "npm:^4.2.2" + jest: "npm:^27.5.1" + lodash: "npm:^4.17.21" + ts-jest: "npm:^27.1.4" + typedoc: "npm:^0.24.8" + typedoc-plugin-missing-exports: "npm:^2.0.0" + typescript: "npm:~5.2.2" + webextension-polyfill: "npm:^0.12.0" + peerDependencies: + "@metamask/accounts-controller": ^30.0.0 + "@metamask/keyring-controller": ^22.0.0 + "@metamask/providers": ^22.0.0 + "@metamask/snaps-controllers": ^12.0.0 + webextension-polyfill: ^0.10.0 || ^0.11.0 || ^0.12.0 + languageName: unknown + linkType: soft + "@metamask/accounts-controller@npm:^30.0.0, @metamask/accounts-controller@workspace:packages/accounts-controller": version: 0.0.0-use.local resolution: "@metamask/accounts-controller@workspace:packages/accounts-controller"