-
Notifications
You must be signed in to change notification settings - Fork 4.9k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat(multichain): add BalancesController (#25257)
## **Description** This PR adds the Balances Controller that retrieves non EVM balances of accounts. ## **Related issues** Fixes MetaMask/accounts-planning#452 ## **Manual testing steps** **This tests requires a btc account for the balances to appear in the state** 1. Start the extension 2. Create a non evm account 3. Wait for interval to pass and see the balances in the state `state.metamask.balances` ## **Screenshots/Recordings** Not Applicable ## **Pre-merge author checklist** - [x] I've followed [MetaMask Contributor Docs](https://github.com/MetaMask/contributor-docs) and [MetaMask Extension Coding Standards](https://github.com/MetaMask/metamask-extension/blob/develop/.github/guidelines/CODING_GUIDELINES.md). - [x] I've completed the PR template to the best of my ability - [x] I’ve included tests if applicable - [ ] I’ve documented my code using [JSDoc](https://jsdoc.app/) format if applicable - [ ] I’ve applied the right labels on the PR (see [labeling guidelines](https://github.com/MetaMask/metamask-extension/blob/develop/.github/guidelines/LABELING_GUIDELINES.md)). Not required for external contributors. ## **Pre-merge reviewer checklist** - [ ] I've manually tested the PR (e.g. pull and build branch, run the app, test code being changed). - [ ] I confirm that this PR addresses all acceptance criteria described in the ticket it closes and includes the necessary testing evidence such as recordings and or screenshots. --------- Co-authored-by: Monte Lai <[email protected]> Co-authored-by: MetaMask Bot <[email protected]> Co-authored-by: Charly Chevalier <[email protected]>
- Loading branch information
1 parent
7c87e58
commit a4dc9d3
Showing
11 changed files
with
591 additions
and
18 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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<CaipAssetType, Balance>; | ||
}; | ||
} = {}) => { | ||
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', | ||
}, | ||
}, | ||
}, | ||
}); | ||
}); | ||
}); |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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<void> { | ||
this.#poller.start(); | ||
} | ||
|
||
/** | ||
* Stops the polling process. | ||
*/ | ||
async stop(): Promise<void> { | ||
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<InternalAccount[]> { | ||
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<BalancesControllerState>) => ({ | ||
...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<Record<CaipAssetType, Balance>> { | ||
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<Json>, | ||
}); | ||
} | ||
} |
Oops, something went wrong.