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 <[email protected]>
Co-authored-by: cryptodev-2s <[email protected]>

packages/assets-controllers/package.json

@@ -61,10 +61,12 @@
     "@metamask/metamask-eth-abis": "^3.1.1",
     "@metamask/polling-controller": "^12.0.2",
     "@metamask/rpc-errors": "^7.0.2",
+    "@metamask/snaps-utils": "^8.3.0",
     "@metamask/utils": "^11.0.1",
     "@types/bn.js": "^5.1.5",
     "@types/uuid": "^8.3.0",
     "async-mutex": "^0.5.0",
+    "bitcoin-address-validation": "^2.2.3",
     "bn.js": "^5.2.1",
     "cockatiel": "^3.1.2",
     "immer": "^9.0.6",
@@ -79,11 +81,15 @@
     "@metamask/approval-controller": "^7.1.1",
     "@metamask/auto-changelog": "^3.4.4",
     "@metamask/ethjs-provider-http": "^0.3.0",
+    "@metamask/keyring-api": "^13.0.0",
     "@metamask/keyring-controller": "^19.0.2",
     "@metamask/keyring-internal-api": "^1.1.0",
+    "@metamask/keyring-snap-client": "^1.0.0",
     "@metamask/network-controller": "^22.1.1",
     "@metamask/preferences-controller": "^15.0.1",
     "@metamask/providers": "^18.1.1",
+    "@metamask/snaps-controllers": "^9.10.0",
+    "@metamask/snaps-sdk": "^6.7.0",
     "@types/jest": "^27.4.1",
     "@types/lodash": "^4.14.191",
     "@types/node": "^16.18.54",

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

@@ -0,0 +1,143 @@
+import { BtcAccountType, BtcMethod } from '@metamask/keyring-api';
+import { KeyringTypes } from '@metamask/keyring-controller';
+import { v4 as uuidv4 } from 'uuid';
+
+import { BalancesTracker } from './BalancesTracker';
+import { Poller } from './Poller';
+
+const MOCK_TIMESTAMP = 1709983353;
+
+const mockBtcAccount = {
+  address: '',
+  id: uuidv4(),
+  metadata: {
+    name: 'Bitcoin Account 1',
+    importTime: Date.now(),
+    keyring: {
+      type: KeyringTypes.snap,
+    },
+    snap: {
+      id: 'mock-btc-snap',
+      name: 'mock-btc-snap',
+      enabled: true,
+    },
+    lastSelected: 0,
+  },
+  options: {},
+  methods: [BtcMethod.SendBitcoin],
+  type: BtcAccountType.P2wpkh,
+};
+
+/**
+ * Sets up a BalancesTracker instance for testing.
+ * @returns The BalancesTracker instance and a mock update balance function.
+ */
+function setupTracker() {
+  const mockUpdateBalance = jest.fn();
+  const tracker = new BalancesTracker(mockUpdateBalance);
+
+  return {
+    tracker,
+    mockUpdateBalance,
+  };
+}
+
+describe('BalancesTracker', () => {
+  it('starts polling when calling start', async () => {
+    const { tracker } = setupTracker();
+    const spyPoller = jest.spyOn(Poller.prototype, 'start');
+
+    tracker.start();
+    expect(spyPoller).toHaveBeenCalledTimes(1);
+  });
+
+  it('stops polling when calling stop', async () => {
+    const { tracker } = setupTracker();
+    const spyPoller = jest.spyOn(Poller.prototype, 'stop');
+
+    tracker.start();
+    tracker.stop();
+    expect(spyPoller).toHaveBeenCalledTimes(1);
+  });
+
+  it('is not tracking if none accounts have been registered', async () => {
+    const { tracker, mockUpdateBalance } = setupTracker();
+
+    tracker.start();
+    await tracker.updateBalances();
+
+    expect(mockUpdateBalance).not.toHaveBeenCalled();
+  });
+
+  it('tracks account balances', async () => {
+    const { tracker, mockUpdateBalance } = setupTracker();
+
+    tracker.start();
+    // We must track account IDs explicitly
+    tracker.track(mockBtcAccount.id, 0);
+    // Trigger balances refresh (not waiting for the Poller here)
+    await tracker.updateBalances();
+
+    expect(mockUpdateBalance).toHaveBeenCalledWith(mockBtcAccount.id);
+  });
+
+  it('untracks account balances', async () => {
+    const { tracker, mockUpdateBalance } = setupTracker();
+
+    tracker.start();
+    tracker.track(mockBtcAccount.id, 0);
+    await tracker.updateBalances();
+    expect(mockUpdateBalance).toHaveBeenCalledWith(mockBtcAccount.id);
+
+    tracker.untrack(mockBtcAccount.id);
+    await tracker.updateBalances();
+    expect(mockUpdateBalance).toHaveBeenCalledTimes(1); // No second call after untracking
+  });
+
+  it('tracks account after being registered', async () => {
+    const { tracker } = setupTracker();
+
+    tracker.start();
+    tracker.track(mockBtcAccount.id, 0);
+    expect(tracker.isTracked(mockBtcAccount.id)).toBe(true);
+  });
+
+  it('does not track account if not registered', async () => {
+    const { tracker } = setupTracker();
+
+    tracker.start();
+    expect(tracker.isTracked(mockBtcAccount.id)).toBe(false);
+  });
+
+  it('does not refresh balance if they are considered up-to-date', async () => {
+    const { tracker, mockUpdateBalance } = setupTracker();
+
+    const blockTime = 10 * 60 * 1000; // 10 minutes in milliseconds.
+    jest
+      .spyOn(global.Date, 'now')
+      .mockImplementation(() => new Date(MOCK_TIMESTAMP).getTime());
+
+    tracker.start();
+    tracker.track(mockBtcAccount.id, blockTime);
+    await tracker.updateBalances();
+    expect(mockUpdateBalance).toHaveBeenCalledTimes(1);
+
+    await tracker.updateBalances();
+    expect(mockUpdateBalance).toHaveBeenCalledTimes(1); // No second call since the balances is already still up-to-date
+
+    jest
+      .spyOn(global.Date, 'now')
+      .mockImplementation(() => new Date(MOCK_TIMESTAMP + blockTime).getTime());
+
+    await tracker.updateBalances();
+    expect(mockUpdateBalance).toHaveBeenCalledTimes(2); // Now the balance will update
+  });
+
+  it('throws an error if trying to update balance of an untracked account', async () => {
+    const { tracker } = setupTracker();
+
+    await expect(tracker.updateBalance(mockBtcAccount.id)).rejects.toThrow(
+      `Account is not being tracked: ${mockBtcAccount.id}`,
+    );
+  });
+});

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

@@ -0,0 +1,139 @@
+import { Poller } from './Poller';
+
+type BalanceInfo = {
+  lastUpdated: number;
+  blockTime: number;
+};
+
+const BALANCES_TRACKING_INTERVAL = 5000; // Every 5s in milliseconds.
+
+export class BalancesTracker {
+  #poller: Poller;
+
+  #updateBalance: (accountId: string) => Promise<void>;
+
+  #balances: Record<string, BalanceInfo> = {};
+
+  constructor(updateBalanceCallback: (accountId: string) => Promise<void>) {
+    this.#updateBalance = updateBalanceCallback;
+
+    this.#poller = new Poller(
+      () => this.updateBalances(),
+      BALANCES_TRACKING_INTERVAL,
+    );
+  }
+
+  /**
+   * Starts the tracking process.
+   */
+  start(): void {
+    this.#poller.start();
+  }
+
+  /**
+   * Stops the tracking process.
+   */
+  stop(): void {
+    this.#poller.stop();
+  }
+
+  /**
+   * Checks if an account ID is being tracked.
+   *
+   * @param accountId - The account ID.
+   * @returns True if the account is being tracked, false otherwise.
+   */
+  isTracked(accountId: string) {
+    return Object.prototype.hasOwnProperty.call(this.#balances, accountId);
+  }
+
+  /**
+   * Asserts that an account ID is being tracked.
+   *
+   * @param accountId - The account ID.
+   * @throws If the account ID is not being tracked.
+   */
+  assertBeingTracked(accountId: string) {
+    if (!this.isTracked(accountId)) {
+      throw new Error(`Account is not being tracked: ${accountId}`);
+    }
+  }
+
+  /**
+   * Starts tracking a new account ID. This method has no effect on already tracked
+   * accounts.
+   *
+   * @param accountId - The account ID.
+   * @param blockTime - The block time (used when refreshing the account balances).
+   */
+  track(accountId: string, blockTime: number) {
+    // Do not overwrite current info if already being tracked!
+    if (!this.isTracked(accountId)) {
+      this.#balances[accountId] = {
+        lastUpdated: 0,
+        blockTime,
+      };
+    }
+  }
+
+  /**
+   * Stops tracking a tracked account ID.
+   *
+   * @param accountId - The account ID.
+   * @throws If the account ID is not being tracked.
+   */
+  untrack(accountId: string) {
+    this.assertBeingTracked(accountId);
+    delete this.#balances[accountId];
+  }
+
+  /**
+   * Update the balances for a tracked account ID.
+   *
+   * @param accountId - The account ID.
+   * @throws If the account ID is not being tracked.
+   */
+  async updateBalance(accountId: string) {
+    this.assertBeingTracked(accountId);
+
+    // We check if the balance is outdated (by comparing to the block time associated
+    // with this kind of account).
+    //
+    // This might not be super accurate, but we could probably compute this differently
+    // and try to sync with the "real block time"!
+    const info = this.#balances[accountId];
+    if (this.#isBalanceOutdated(info)) {
+      await this.#updateBalance(accountId);
+      this.#balances[accountId].lastUpdated = Date.now();
+    }
+  }
+
+  /**
+   * Update the balances of all tracked accounts (only if the balances
+   * is considered outdated).
+   */
+  async updateBalances() {
+    await Promise.allSettled(
+      Object.keys(this.#balances).map(async (accountId) => {
+        await this.updateBalance(accountId);
+      }),
+    );
+  }
+
+  /**
+   * Checks if the balance is outdated according to the provided data.
+   *
+   * @param param - The balance info.
+   * @param param.lastUpdated - The last updated timestamp.
+   * @param param.blockTime - The block time.
+   * @returns True if the balance is outdated, false otherwise.
+   */
+  #isBalanceOutdated({ lastUpdated, blockTime }: BalanceInfo): boolean {
+    return (
+      // Never been updated:
+      lastUpdated === 0 ||
+      // Outdated:
+      Date.now() - lastUpdated >= blockTime
+    );
+  }
+}