Skip to content

Commit c3b90a5

Browse files
gantunesrccharlycryptodev-2s
authored
feat: add MultichainBalancesController (#4965)
This change adds a new controller to track the balance from non-EVM accounts by using the snaps as the data source. from non-EVM accounts by using the snaps as the data source Co-authored-by: Charly Chevalier <charly.chevalier@consensys.net> Co-authored-by: cryptodev-2s <109512101+cryptodev-2s@users.noreply.github.com>
1 parent 3a7e6cb commit c3b90a5

File tree

15 files changed

+1506
-16
lines changed

15 files changed

+1506
-16
lines changed

‎packages/assets-controllers/package.json‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,10 +61,12 @@
6161
"@metamask/metamask-eth-abis": "^3.1.1",
6262
"@metamask/polling-controller": "^12.0.2",
6363
"@metamask/rpc-errors": "^7.0.2",
64+
"@metamask/snaps-utils": "^8.3.0",
6465
"@metamask/utils": "^11.0.1",
6566
"@types/bn.js": "^5.1.5",
6667
"@types/uuid": "^8.3.0",
6768
"async-mutex": "^0.5.0",
69+
"bitcoin-address-validation": "^2.2.3",
6870
"bn.js": "^5.2.1",
6971
"cockatiel": "^3.1.2",
7072
"immer": "^9.0.6",
@@ -79,11 +81,15 @@
7981
"@metamask/approval-controller": "^7.1.1",
8082
"@metamask/auto-changelog": "^3.4.4",
8183
"@metamask/ethjs-provider-http": "^0.3.0",
84+
"@metamask/keyring-api": "^13.0.0",
8285
"@metamask/keyring-controller": "^19.0.2",
8386
"@metamask/keyring-internal-api": "^1.1.0",
87+
"@metamask/keyring-snap-client": "^1.0.0",
8488
"@metamask/network-controller": "^22.1.1",
8589
"@metamask/preferences-controller": "^15.0.1",
8690
"@metamask/providers": "^18.1.1",
91+
"@metamask/snaps-controllers": "^9.10.0",
92+
"@metamask/snaps-sdk": "^6.7.0",
8793
"@types/jest": "^27.4.1",
8894
"@types/lodash": "^4.14.191",
8995
"@types/node": "^16.18.54",

‎packages/assets-controllers/src/MultichainBalancesController/BalancesTracker.test.ts‎

Lines changed: 143 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,143 @@
1+
import { BtcAccountType, BtcMethod } from '@metamask/keyring-api';
2+
import { KeyringTypes } from '@metamask/keyring-controller';
3+
import { v4 as uuidv4 } from 'uuid';
4+
5+
import { BalancesTracker } from './BalancesTracker';
6+
import { Poller } from './Poller';
7+
8+
const MOCK_TIMESTAMP = 1709983353;
9+
10+
const mockBtcAccount = {
11+
address: '',
12+
id: uuidv4(),
13+
metadata: {
14+
name: 'Bitcoin Account 1',
15+
importTime: Date.now(),
16+
keyring: {
17+
type: KeyringTypes.snap,
18+
},
19+
snap: {
20+
id: 'mock-btc-snap',
21+
name: 'mock-btc-snap',
22+
enabled: true,
23+
},
24+
lastSelected: 0,
25+
},
26+
options: {},
27+
methods: [BtcMethod.SendBitcoin],
28+
type: BtcAccountType.P2wpkh,
29+
};
30+
31+
/**
32+
* Sets up a BalancesTracker instance for testing.
33+
* @returns The BalancesTracker instance and a mock update balance function.
34+
*/
35+
function setupTracker() {
36+
const mockUpdateBalance = jest.fn();
37+
const tracker = new BalancesTracker(mockUpdateBalance);
38+
39+
return {
40+
tracker,
41+
mockUpdateBalance,
42+
};
43+
}
44+
45+
describe('BalancesTracker', () => {
46+
it('starts polling when calling start', async () => {
47+
const { tracker } = setupTracker();
48+
const spyPoller = jest.spyOn(Poller.prototype, 'start');
49+
50+
tracker.start();
51+
expect(spyPoller).toHaveBeenCalledTimes(1);
52+
});
53+
54+
it('stops polling when calling stop', async () => {
55+
const { tracker } = setupTracker();
56+
const spyPoller = jest.spyOn(Poller.prototype, 'stop');
57+
58+
tracker.start();
59+
tracker.stop();
60+
expect(spyPoller).toHaveBeenCalledTimes(1);
61+
});
62+
63+
it('is not tracking if none accounts have been registered', async () => {
64+
const { tracker, mockUpdateBalance } = setupTracker();
65+
66+
tracker.start();
67+
await tracker.updateBalances();
68+
69+
expect(mockUpdateBalance).not.toHaveBeenCalled();
70+
});
71+
72+
it('tracks account balances', async () => {
73+
const { tracker, mockUpdateBalance } = setupTracker();
74+
75+
tracker.start();
76+
// We must track account IDs explicitly
77+
tracker.track(mockBtcAccount.id, 0);
78+
// Trigger balances refresh (not waiting for the Poller here)
79+
await tracker.updateBalances();
80+
81+
expect(mockUpdateBalance).toHaveBeenCalledWith(mockBtcAccount.id);
82+
});
83+
84+
it('untracks account balances', async () => {
85+
const { tracker, mockUpdateBalance } = setupTracker();
86+
87+
tracker.start();
88+
tracker.track(mockBtcAccount.id, 0);
89+
await tracker.updateBalances();
90+
expect(mockUpdateBalance).toHaveBeenCalledWith(mockBtcAccount.id);
91+
92+
tracker.untrack(mockBtcAccount.id);
93+
await tracker.updateBalances();
94+
expect(mockUpdateBalance).toHaveBeenCalledTimes(1); // No second call after untracking
95+
});
96+
97+
it('tracks account after being registered', async () => {
98+
const { tracker } = setupTracker();
99+
100+
tracker.start();
101+
tracker.track(mockBtcAccount.id, 0);
102+
expect(tracker.isTracked(mockBtcAccount.id)).toBe(true);
103+
});
104+
105+
it('does not track account if not registered', async () => {
106+
const { tracker } = setupTracker();
107+
108+
tracker.start();
109+
expect(tracker.isTracked(mockBtcAccount.id)).toBe(false);
110+
});
111+
112+
it('does not refresh balance if they are considered up-to-date', async () => {
113+
const { tracker, mockUpdateBalance } = setupTracker();
114+
115+
const blockTime = 10 * 60 * 1000; // 10 minutes in milliseconds.
116+
jest
117+
.spyOn(global.Date, 'now')
118+
.mockImplementation(() => new Date(MOCK_TIMESTAMP).getTime());
119+
120+
tracker.start();
121+
tracker.track(mockBtcAccount.id, blockTime);
122+
await tracker.updateBalances();
123+
expect(mockUpdateBalance).toHaveBeenCalledTimes(1);
124+
125+
await tracker.updateBalances();
126+
expect(mockUpdateBalance).toHaveBeenCalledTimes(1); // No second call since the balances is already still up-to-date
127+
128+
jest
129+
.spyOn(global.Date, 'now')
130+
.mockImplementation(() => new Date(MOCK_TIMESTAMP + blockTime).getTime());
131+
132+
await tracker.updateBalances();
133+
expect(mockUpdateBalance).toHaveBeenCalledTimes(2); // Now the balance will update
134+
});
135+
136+
it('throws an error if trying to update balance of an untracked account', async () => {
137+
const { tracker } = setupTracker();
138+
139+
await expect(tracker.updateBalance(mockBtcAccount.id)).rejects.toThrow(
140+
`Account is not being tracked: ${mockBtcAccount.id}`,
141+
);
142+
});
143+
});

‎packages/assets-controllers/src/MultichainBalancesController/BalancesTracker.ts‎

Lines changed: 139 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,139 @@
1+
import { Poller } from './Poller';
2+
3+
type BalanceInfo = {
4+
lastUpdated: number;
5+
blockTime: number;
6+
};
7+
8+
const BALANCES_TRACKING_INTERVAL = 5000; // Every 5s in milliseconds.
9+
10+
export class BalancesTracker {
11+
#poller: Poller;
12+
13+
#updateBalance: (accountId: string) => Promise<void>;
14+
15+
#balances: Record<string, BalanceInfo> = {};
16+
17+
constructor(updateBalanceCallback: (accountId: string) => Promise<void>) {
18+
this.#updateBalance = updateBalanceCallback;
19+
20+
this.#poller = new Poller(
21+
() => this.updateBalances(),
22+
BALANCES_TRACKING_INTERVAL,
23+
);
24+
}
25+
26+
/**
27+
* Starts the tracking process.
28+
*/
29+
start(): void {
30+
this.#poller.start();
31+
}
32+
33+
/**
34+
* Stops the tracking process.
35+
*/
36+
stop(): void {
37+
this.#poller.stop();
38+
}
39+
40+
/**
41+
* Checks if an account ID is being tracked.
42+
*
43+
* @param accountId - The account ID.
44+
* @returns True if the account is being tracked, false otherwise.
45+
*/
46+
isTracked(accountId: string) {
47+
return Object.prototype.hasOwnProperty.call(this.#balances, accountId);
48+
}
49+
50+
/**
51+
* Asserts that an account ID is being tracked.
52+
*
53+
* @param accountId - The account ID.
54+
* @throws If the account ID is not being tracked.
55+
*/
56+
assertBeingTracked(accountId: string) {
57+
if (!this.isTracked(accountId)) {
58+
throw new Error(`Account is not being tracked: ${accountId}`);
59+
}
60+
}
61+
62+
/**
63+
* Starts tracking a new account ID. This method has no effect on already tracked
64+
* accounts.
65+
*
66+
* @param accountId - The account ID.
67+
* @param blockTime - The block time (used when refreshing the account balances).
68+
*/
69+
track(accountId: string, blockTime: number) {
70+
// Do not overwrite current info if already being tracked!
71+
if (!this.isTracked(accountId)) {
72+
this.#balances[accountId] = {
73+
lastUpdated: 0,
74+
blockTime,
75+
};
76+
}
77+
}
78+
79+
/**
80+
* Stops tracking a tracked account ID.
81+
*
82+
* @param accountId - The account ID.
83+
* @throws If the account ID is not being tracked.
84+
*/
85+
untrack(accountId: string) {
86+
this.assertBeingTracked(accountId);
87+
delete this.#balances[accountId];
88+
}
89+
90+
/**
91+
* Update the balances for a tracked account ID.
92+
*
93+
* @param accountId - The account ID.
94+
* @throws If the account ID is not being tracked.
95+
*/
96+
async updateBalance(accountId: string) {
97+
this.assertBeingTracked(accountId);
98+
99+
// We check if the balance is outdated (by comparing to the block time associated
100+
// with this kind of account).
101+
//
102+
// This might not be super accurate, but we could probably compute this differently
103+
// and try to sync with the "real block time"!
104+
const info = this.#balances[accountId];
105+
if (this.#isBalanceOutdated(info)) {
106+
await this.#updateBalance(accountId);
107+
this.#balances[accountId].lastUpdated = Date.now();
108+
}
109+
}
110+
111+
/**
112+
* Update the balances of all tracked accounts (only if the balances
113+
* is considered outdated).
114+
*/
115+
async updateBalances() {
116+
await Promise.allSettled(
117+
Object.keys(this.#balances).map(async (accountId) => {
118+
await this.updateBalance(accountId);
119+
}),
120+
);
121+
}
122+
123+
/**
124+
* Checks if the balance is outdated according to the provided data.
125+
*
126+
* @param param - The balance info.
127+
* @param param.lastUpdated - The last updated timestamp.
128+
* @param param.blockTime - The block time.
129+
* @returns True if the balance is outdated, false otherwise.
130+
*/
131+
#isBalanceOutdated({ lastUpdated, blockTime }: BalanceInfo): boolean {
132+
return (
133+
// Never been updated:
134+
lastUpdated === 0 ||
135+
// Outdated:
136+
Date.now() - lastUpdated >= blockTime
137+
);
138+
}
139+
}

0 commit comments

Comments
 (0)