diff --git a/app/scripts/lib/AccountIdentitiesPetnamesBridge.test.ts b/app/scripts/lib/AccountIdentitiesPetnamesBridge.test.ts index 7e1a23a5902e..c17ce27d68aa 100644 --- a/app/scripts/lib/AccountIdentitiesPetnamesBridge.test.ts +++ b/app/scripts/lib/AccountIdentitiesPetnamesBridge.test.ts @@ -27,7 +27,6 @@ const MOCK_INTERNAL_ACCOUNT = createMockInternalAccount({ address: ADDRESS_MOCK, name: NAME_MOCK, keyringType: KeyringTypes.hd, - is4337: false, snapOptions: undefined, }); diff --git a/app/scripts/lib/accounts/BalancesController.test.ts b/app/scripts/lib/accounts/BalancesController.test.ts new file mode 100644 index 000000000000..01ce1f88c608 --- /dev/null +++ b/app/scripts/lib/accounts/BalancesController.test.ts @@ -0,0 +1,122 @@ +import { ControllerMessenger } from '@metamask/base-controller'; +import { + Balance, + BtcAccountType, + CaipAssetType, + InternalAccount, +} from '@metamask/keyring-api'; +import { createMockInternalAccount } from '../../../../test/jest/mocks'; +import { + BalancesController, + AllowedActions, + BalancesControllerEvents, + BalancesControllerState, + defaultState, +} from './BalancesController'; +import { Poller } from './Poller'; + +const mockBtcAccount = createMockInternalAccount({ + address: '', + name: 'Btc Account', + // @ts-expect-error - account type may be btc or eth, mock file is not typed + type: BtcAccountType.P2wpkh, + snapOptions: { + id: 'mock-btc-snap', + name: 'mock-btc-snap', + enabled: true, + }, +}); + +const mockBalanceResult = { + 'bip122:000000000933ea01ad0ee984209779ba/slip44:0': { + amount: '0.00000000', + unit: 'BTC', + }, +}; + +const setupController = ({ + state = defaultState, + mocks, +}: { + state?: BalancesControllerState; + mocks?: { + listMultichainAccounts?: InternalAccount[]; + handleRequestReturnValue?: Record; + }; +} = {}) => { + const controllerMessenger = new ControllerMessenger< + AllowedActions, + BalancesControllerEvents + >(); + + const balancesControllerMessenger = controllerMessenger.getRestricted({ + name: 'BalancesController', + allowedActions: ['SnapController:handleRequest'], + allowedEvents: [], + }); + + const mockSnapHandleRequest = jest.fn(); + controllerMessenger.registerActionHandler( + 'SnapController:handleRequest', + mockSnapHandleRequest.mockReturnValue( + mocks?.handleRequestReturnValue ?? mockBalanceResult, + ), + ); + + // TODO: remove when listMultichainAccounts action is available + const mockListMultichainAccounts = jest + .fn() + .mockReturnValue(mocks?.listMultichainAccounts ?? [mockBtcAccount]); + + const controller = new BalancesController({ + messenger: balancesControllerMessenger, + state, + // TODO: remove when listMultichainAccounts action is available + listMultichainAccounts: mockListMultichainAccounts, + }); + + return { + controller, + mockSnapHandleRequest, + mockListMultichainAccounts, + }; +}; + +describe('BalancesController', () => { + it('initialize with default state', () => { + const { controller } = setupController({}); + expect(controller.state).toEqual({ balances: {} }); + }); + + it('starts polling when calling start', async () => { + const spyPoller = jest.spyOn(Poller.prototype, 'start'); + const { controller } = setupController(); + await controller.start(); + expect(spyPoller).toHaveBeenCalledTimes(1); + }); + + it('stops polling when calling stop', async () => { + const spyPoller = jest.spyOn(Poller.prototype, 'stop'); + const { controller } = setupController(); + await controller.start(); + await controller.stop(); + expect(spyPoller).toHaveBeenCalledTimes(1); + }); + + it('update balances when calling updateBalances', async () => { + const { controller } = setupController(); + + await controller.updateBalances(); + + expect(controller.state).toEqual({ + balances: { + [mockBtcAccount.id]: { + 'bip122:000000000933ea01ad0ee984209779ba/slip44:0': { + amount: '0.00000000', + unit: 'BTC', + }, + }, + }, + }); + }); +}); diff --git a/app/scripts/lib/accounts/BalancesController.ts b/app/scripts/lib/accounts/BalancesController.ts new file mode 100644 index 000000000000..eee4ac11889a --- /dev/null +++ b/app/scripts/lib/accounts/BalancesController.ts @@ -0,0 +1,255 @@ +import type { Json, JsonRpcRequest } from '@metamask/utils'; + +import { + BaseController, + type ControllerGetStateAction, + type ControllerStateChangeEvent, + type RestrictedControllerMessenger, +} from '@metamask/base-controller'; +import { + BtcAccountType, + KeyringClient, + type Balance, + type CaipAssetType, + type InternalAccount, +} from '@metamask/keyring-api'; +import type { HandleSnapRequest } from '@metamask/snaps-controllers'; +import type { SnapId } from '@metamask/snaps-sdk'; +import { HandlerType } from '@metamask/snaps-utils'; +import type { Draft } from 'immer'; +import { Poller } from './Poller'; + +const controllerName = 'BalancesController'; + +/** + * State used by the {@link BalancesController} to cache account balances. + */ +export type BalancesControllerState = { + balances: { + [account: string]: { + [asset: string]: { + amount: string; + unit: string; + }; + }; + }; +}; + +/** + * Default state of the {@link BalancesController}. + */ +export const defaultState: BalancesControllerState = { balances: {} }; + +/** + * Returns the state of the {@link BalancesController}. + */ +export type BalancesControllerGetStateAction = ControllerGetStateAction< + typeof controllerName, + BalancesControllerState +>; + +/** + * Updates the balances of all supported accounts. + */ +export type BalancesControllerUpdateBalancesAction = { + type: `${typeof controllerName}:updateBalances`; + handler: BalancesController['updateBalances']; +}; + +/** + * Event emitted when the state of the {@link BalancesController} changes. + */ +export type BalancesControllerStateChange = ControllerStateChangeEvent< + typeof controllerName, + BalancesControllerState +>; + +/** + * Actions exposed by the {@link BalancesController}. + */ +export type BalancesControllerActions = + | BalancesControllerGetStateAction + | BalancesControllerUpdateBalancesAction; + +/** + * Events emitted by {@link BalancesController}. + */ +export type BalancesControllerEvents = BalancesControllerStateChange; + +/** + * Actions that this controller is allowed to call. + */ +export type AllowedActions = HandleSnapRequest; + +/** + * Messenger type for the BalancesController. + */ +export type BalancesControllerMessenger = RestrictedControllerMessenger< + typeof controllerName, + BalancesControllerActions | AllowedActions, + BalancesControllerEvents, + AllowedActions['type'], + never +>; + +/** + * {@link BalancesController}'s metadata. + * + * This allows us to choose if fields of the state should be persisted or not + * using the `persist` flag; and if they can be sent to Sentry or not, using + * the `anonymous` flag. + */ +const balancesControllerMetadata = { + balances: { + persist: true, + anonymous: false, + }, +}; + +const BTC_TESTNET_ASSETS = ['bip122:000000000933ea01ad0ee984209779ba/slip44:0']; +const BTC_MAINNET_ASSETS = ['bip122:000000000019d6689c085ae165831e93/slip44:0']; +export const BTC_AVG_BLOCK_TIME = 600000; // 10 minutes in milliseconds + +/** + * Returns whether an address is on the Bitcoin mainnet. + * + * This function only checks the prefix of the address to determine if it's on + * the mainnet or not. It doesn't validate the address itself, and should only + * be used as a temporary solution until this information is included in the + * account object. + * + * @param address - The address to check. + * @returns `true` if the address is on the Bitcoin mainnet, `false` otherwise. + */ +function isBtcMainnet(address: string): boolean { + return address.startsWith('bc1') || address.startsWith('1'); +} + +/** + * The BalancesController is responsible for fetching and caching account + * balances. + */ +export class BalancesController extends BaseController< + typeof controllerName, + BalancesControllerState, + BalancesControllerMessenger +> { + #poller: Poller; + + // TODO: remove once action is implemented + #listMultichainAccounts: () => InternalAccount[]; + + constructor({ + messenger, + state, + listMultichainAccounts, + }: { + messenger: BalancesControllerMessenger; + state: BalancesControllerState; + listMultichainAccounts: () => InternalAccount[]; + }) { + super({ + messenger, + name: controllerName, + metadata: balancesControllerMetadata, + state: { + ...defaultState, + ...state, + }, + }); + + this.#listMultichainAccounts = listMultichainAccounts; + this.#poller = new Poller(() => this.updateBalances(), BTC_AVG_BLOCK_TIME); + } + + /** + * Starts the polling process. + */ + async start(): Promise { + this.#poller.start(); + } + + /** + * Stops the polling process. + */ + async stop(): Promise { + this.#poller.stop(); + } + + /** + * Lists the accounts that we should get balances for. + * + * Currently, we only get balances for P2WPKH accounts, but this will change + * in the future when we start support other non-EVM account types. + * + * @returns A list of accounts that we should get balances for. + */ + async #listAccounts(): Promise { + const accounts = this.#listMultichainAccounts(); + + return accounts.filter((account) => account.type === BtcAccountType.P2wpkh); + } + + /** + * Updates the balances of all supported accounts. This method doesn't return + * anything, but it updates the state of the controller. + */ + async updateBalances() { + const accounts = await this.#listAccounts(); + const partialState: BalancesControllerState = { balances: {} }; + + for (const account of accounts) { + if (account.metadata.snap) { + partialState.balances[account.id] = await this.#getBalances( + account.id, + account.metadata.snap.id, + isBtcMainnet(account.address) + ? BTC_MAINNET_ASSETS + : BTC_TESTNET_ASSETS, + ); + } + } + + this.update((state: Draft) => ({ + ...state, + ...partialState, + })); + } + + /** + * Get the balances for an account. + * + * @param accountId - ID of the account to get balances for. + * @param snapId - ID of the Snap which manages the account. + * @param assetTypes - Array of asset types to get balances for. + * @returns A map of asset types to balances. + */ + async #getBalances( + accountId: string, + snapId: string, + assetTypes: CaipAssetType[], + ): Promise> { + return await this.#getClient(snapId).getAccountBalances( + accountId, + assetTypes, + ); + } + + /** + * Gets a `KeyringClient` for a Snap. + * + * @param snapId - ID of the Snap to get the client for. + * @returns A `KeyringClient` for the Snap. + */ + #getClient(snapId: string): KeyringClient { + return new KeyringClient({ + send: async (request: JsonRpcRequest) => + (await this.messagingSystem.call('SnapController:handleRequest', { + snapId: snapId as SnapId, + origin: 'metamask', + handler: HandlerType.OnKeyringRequest, + request, + })) as Promise, + }); + } +} diff --git a/app/scripts/lib/accounts/Poller.test.ts b/app/scripts/lib/accounts/Poller.test.ts new file mode 100644 index 000000000000..e79d4961a0c8 --- /dev/null +++ b/app/scripts/lib/accounts/Poller.test.ts @@ -0,0 +1,59 @@ +import { Poller } from './Poller'; + +jest.useFakeTimers(); + +const interval = 1000; +const intervalPlus100ms = interval + 100; + +describe('Poller', () => { + let callback: jest.Mock; + + beforeEach(() => { + callback = jest.fn(); + }); + + afterEach(() => { + jest.clearAllMocks(); + }); + + it('calls the callback function after the specified interval', async () => { + const poller = new Poller(callback, interval); + poller.start(); + jest.advanceTimersByTime(intervalPlus100ms); + poller.stop(); + + expect(callback).toHaveBeenCalledTimes(1); + }); + + it('does not call the callback function if stopped before the interval', async () => { + const poller = new Poller(callback, interval); + poller.start(); + poller.stop(); + jest.advanceTimersByTime(intervalPlus100ms); + + expect(callback).not.toHaveBeenCalled(); + }); + + it('calls the callback function multiple times if started and stopped multiple times', async () => { + const poller = new Poller(callback, interval); + poller.start(); + jest.advanceTimersByTime(intervalPlus100ms); + poller.stop(); + jest.advanceTimersByTime(intervalPlus100ms); + poller.start(); + jest.advanceTimersByTime(intervalPlus100ms); + poller.stop(); + + expect(callback).toHaveBeenCalledTimes(2); + }); + + it('does not call the callback if the poller is stopped before the interval has passed', async () => { + const poller = new Poller(callback, interval); + poller.start(); + // Wait for some time, but resumes before reaching out + // the `interval` timeout + jest.advanceTimersByTime(interval / 2); + poller.stop(); + expect(callback).not.toHaveBeenCalled(); + }); +}); diff --git a/app/scripts/lib/accounts/Poller.ts b/app/scripts/lib/accounts/Poller.ts new file mode 100644 index 000000000000..600e2ea615d7 --- /dev/null +++ b/app/scripts/lib/accounts/Poller.ts @@ -0,0 +1,28 @@ +export class Poller { + #interval: number; + + #callback: () => void; + + #handle: NodeJS.Timeout | undefined = undefined; + + constructor(callback: () => void, interval: number) { + this.#interval = interval; + this.#callback = callback; + } + + start() { + if (this.#handle) { + return; + } + + this.#handle = setInterval(this.#callback, this.#interval); + } + + stop() { + if (!this.#handle) { + return; + } + clearInterval(this.#handle); + this.#handle = undefined; + } +} diff --git a/app/scripts/lib/setupSentry.js b/app/scripts/lib/setupSentry.js index fd3e771bcff5..92173fc5f72c 100644 --- a/app/scripts/lib/setupSentry.js +++ b/app/scripts/lib/setupSentry.js @@ -117,6 +117,9 @@ export const SENTRY_BACKGROUND_STATE = { trezorModel: true, usedNetworks: true, }, + MultichainBalancesController: { + balances: false, + }, CronjobController: { jobs: false, }, diff --git a/app/scripts/metamask-controller.js b/app/scripts/metamask-controller.js index 646676c3c8f7..11463edf3430 100644 --- a/app/scripts/metamask-controller.js +++ b/app/scripts/metamask-controller.js @@ -222,6 +222,7 @@ import { getCurrentChainSupportsSmartTransactions, } from '../../shared/modules/selectors'; import { BaseUrl } from '../../shared/constants/urls'; +import { BalancesController as MultichainBalancesController } from './lib/accounts/BalancesController'; import { ///: BEGIN:ONLY_INCLUDE_IF(build-mmi) handleMMITransactionUpdate, @@ -926,6 +927,23 @@ export default class MetamaskController extends EventEmitter { state: initState.AccountOrderController, }); + const multichainBalancesControllerMessenger = + this.controllerMessenger.getRestricted({ + name: 'BalancesController', + allowedEvents: [], + allowedActions: ['SnapController:handleRequest'], + }); + + this.multichainBalancesController = new MultichainBalancesController({ + messenger: multichainBalancesControllerMessenger, + state: {}, + // TODO: remove when listMultichainAccounts action is available + listMultichainAccounts: + this.accountsController.listMultichainAccounts.bind( + this.accountsController, + ), + }); + const multichainRatesControllerMessenger = this.controllerMessenger.getRestricted({ name: 'RatesController', @@ -2186,6 +2204,7 @@ export default class MetamaskController extends EventEmitter { AccountsController: this.accountsController, AppStateController: this.appStateController.store, AppMetadataController: this.appMetadataController.store, + MultichainBalancesController: this.multichainBalancesController, TransactionController: this.txController, KeyringController: this.keyringController, PreferencesController: this.preferencesController.store, @@ -2243,6 +2262,7 @@ export default class MetamaskController extends EventEmitter { AccountsController: this.accountsController, AppStateController: this.appStateController.store, AppMetadataController: this.appMetadataController.store, + MultichainBalancesController: this.multichainBalancesController, NetworkController: this.networkController, KeyringController: this.keyringController, PreferencesController: this.preferencesController.store, @@ -2855,6 +2875,8 @@ export default class MetamaskController extends EventEmitter { this.multichainRatesController.start(); }, ); + this.multichainBalancesController.start(); + this.multichainBalancesController.updateBalances(); } /** diff --git a/app/scripts/metamask-controller.test.js b/app/scripts/metamask-controller.test.js index cf80709cfea1..f6c7e4a25b3b 100644 --- a/app/scripts/metamask-controller.test.js +++ b/app/scripts/metamask-controller.test.js @@ -37,6 +37,10 @@ import * as tokenUtils from '../../shared/lib/token-util'; import { flushPromises } from '../../test/lib/timer-helpers'; import { ETH_EOA_METHODS } from '../../shared/constants/eth-methods'; import { createMockInternalAccount } from '../../test/jest/mocks'; +import { + BalancesController as MultichainBalancesController, + BTC_AVG_BLOCK_TIME, +} from './lib/accounts/BalancesController'; import { deferredPromise } from './lib/util'; import MetaMaskController from './metamask-controller'; @@ -2141,6 +2145,66 @@ describe('MetaMaskController', () => { ).toHaveBeenCalled(); }); }); + + describe('MultichainBalancesController', () => { + const mockEvmAccount = createMockInternalAccount(); + const mockNonEvmAccount = { + ...mockEvmAccount, + id: '21690786-6abd-45d8-a9f0-9ff1d8ca76a1', + type: BtcAccountType.P2wpkh, + methods: [BtcMethod.SendMany], + address: 'bc1qar0srrr7xfkvy5l643lydnw9re59gtzzwf5mdq', + }; + let localMetamaskController; + + beforeEach(() => { + jest.useFakeTimers(); + jest.spyOn(MultichainBalancesController.prototype, 'updateBalances'); + localMetamaskController = new MetaMaskController({ + showUserConfirmation: noop, + encryptor: mockEncryptor, + initState: { + ...cloneDeep(firstTimeState), + AccountsController: { + internalAccounts: { + accounts: { + [mockNonEvmAccount.id]: mockNonEvmAccount, + [mockEvmAccount.id]: mockEvmAccount, + }, + selectedAccount: mockNonEvmAccount.id, + }, + }, + }, + initLangCode: 'en_US', + platform: { + showTransactionNotification: () => undefined, + getVersion: () => 'foo', + }, + browser: browserPolyfillMock, + infuraProjectId: 'foo', + isFirstMetaMaskControllerSetup: true, + }); + }); + + afterEach(() => { + jest.clearAllMocks(); + jest.useRealTimers(); + }); + + it('calls updateBalances during startup', async () => { + expect( + localMetamaskController.multichainBalancesController.updateBalances, + ).toHaveBeenCalled(); + }); + + it('calls updateBalances after the interval has passed', async () => { + jest.advanceTimersByTime(BTC_AVG_BLOCK_TIME); + // 2 calls because 1 is during startup + expect( + localMetamaskController.multichainBalancesController.updateBalances, + ).toHaveBeenCalledTimes(2); + }); + }); }); describe('MV3 Specific behaviour', () => { diff --git a/test/e2e/tests/metrics/state-snapshots/errors-after-init-opt-in-background-state.json b/test/e2e/tests/metrics/state-snapshots/errors-after-init-opt-in-background-state.json index e9be8ea5b8c7..534d5e4d606a 100644 --- a/test/e2e/tests/metrics/state-snapshots/errors-after-init-opt-in-background-state.json +++ b/test/e2e/tests/metrics/state-snapshots/errors-after-init-opt-in-background-state.json @@ -119,6 +119,7 @@ "isUpdatingMetamaskNotificationsAccount": "object", "isCheckingAccountsPresence": "boolean" }, + "MultichainBalancesController": { "balances": "object" }, "MultichainRatesController": { "fiatCurrency": "usd", "rates": { "btc": { "conversionDate": 0, "conversionRate": "0" } }, diff --git a/test/e2e/tests/metrics/state-snapshots/errors-after-init-opt-in-ui-state.json b/test/e2e/tests/metrics/state-snapshots/errors-after-init-opt-in-ui-state.json index 2e932a982fef..7bcbd81bf539 100644 --- a/test/e2e/tests/metrics/state-snapshots/errors-after-init-opt-in-ui-state.json +++ b/test/e2e/tests/metrics/state-snapshots/errors-after-init-opt-in-ui-state.json @@ -94,6 +94,7 @@ "previousMigrationVersion": 0, "currentMigrationVersion": "number", "showTokenAutodetectModalOnUpgrade": "object", + "balances": "object", "selectedNetworkClientId": "string", "networksMetadata": { "networkConfigurationId": { diff --git a/test/jest/mocks.js b/test/jest/mocks.js index 3170a8dbb8f9..b52a0d984df3 100644 --- a/test/jest/mocks.js +++ b/test/jest/mocks.js @@ -1,4 +1,9 @@ -import { EthAccountType, EthMethod } from '@metamask/keyring-api'; +import { + EthAccountType, + EthMethod, + BtcMethod, + BtcAccountType, +} from '@metamask/keyring-api'; import { KeyringTypes } from '@metamask/keyring-controller'; import { v4 as uuidv4 } from 'uuid'; import { keyringTypeToName } from '@metamask/accounts-controller'; @@ -168,10 +173,37 @@ export const getInitialSendStateWithExistingTxState = (draftTxState) => ({ export function createMockInternalAccount({ address = MOCK_DEFAULT_ADDRESS, name, - is4337 = false, + type = EthAccountType.Eoa, keyringType = KeyringTypes.hd, snapOptions, } = {}) { + let methods; + + switch (type) { + case EthAccountType.Eoa: + methods = [ + EthMethod.PersonalSign, + EthMethod.Sign, + EthMethod.SignTransaction, + EthMethod.SignTypedDataV1, + EthMethod.SignTypedDataV3, + EthMethod.SignTypedDataV4, + ]; + break; + case EthAccountType.Erc4337: + methods = [ + EthMethod.PatchUserOperation, + EthMethod.PrepareUserOperation, + EthMethod.SignUserOperation, + ]; + break; + case BtcAccountType.P2wpkh: + methods = [BtcMethod.SendMany]; + break; + default: + throw new Error(`Unknown account type: ${type}`); + } + return { address, id: uuidv4(), @@ -184,21 +216,8 @@ export function createMockInternalAccount({ snap: snapOptions, }, options: {}, - methods: is4337 - ? [ - EthMethod.PrepareUserOperation, - EthMethod.PatchUserOperation, - EthMethod.SignUserOperation, - ] - : [ - EthMethod.PersonalSign, - EthMethod.Sign, - EthMethod.SignTransaction, - EthMethod.SignTypedDataV1, - EthMethod.SignTypedDataV3, - EthMethod.SignTypedDataV4, - ], - type: is4337 ? EthAccountType.Erc4337 : EthAccountType.Eoa, + methods, + type, }; }