fix(assets-controller): clean v5 native balance state update

packages/assets-controller/CHANGELOG.md

@@ -12,6 +12,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Bump `lodash-es` from `^4.17.21` to `^4.18.1` ([#10447](https://github.com/MetaMask/core/pull/10447))
 - Bump `@ethersproject/providers` from `^5.7.0` to `^5.8.0` ([#10482](https://github.com/MetaMask/core/pull/10482))
 
+### Fixed
+
+- On the v5 balance path, stop seeding every enabled chain's native onto an account when the update arrives after that account is no longer selected ([#10567](https://github.com/MetaMask/core/pull/10567))
+- On unlock, drop stored balances whose chain namespace is outside the selected account's scopes ([#10567](https://github.com/MetaMask/core/pull/10567))
+
 ## [17.0.0]
 
 ### Added

packages/assets-controller/src/AssetsController.ts

@@ -118,6 +118,7 @@ import {
   createParallelMiddleware,
 } from './middlewares/ParallelMiddleware.js';
 import { RpcFallbackMiddleware } from './middlewares/RpcFallbackMiddleware.js';
+import { dropBalancesOutsideAccountScopes } from './migrations/dropBalancesOutsideAccountScopes.js';
 import type { Assets3346MigrationState } from './migrations/healAssetsInfoMetadata.js';
 import {
   cleanSpamAssets,
@@ -1261,6 +1262,11 @@ export class AssetsController extends BaseController<
 
     // eslint-disable-next-line @typescript-eslint/no-misused-promises
     this.messenger.subscribe('KeyringController:unlock', async () => {
+      try {
+        this.#runOutOfScopeBalanceCleanup();
+      } catch {
+        /* Do nothing */
+      }
       await this.#runSpamCleanup().catch(() => {
         /* Do nothing */
       });
@@ -1339,6 +1345,26 @@ export class AssetsController extends BaseController<
     });
   }
 
+  /**
+   * One-time cleanup of balances a v5 update stamped onto accounts after a
+   * group switch. Runs on unlock, the same point as spam cleanup, and only
+   * touches the selected accounts (their scopes are what define ownership).
+   */
+  #runOutOfScopeBalanceCleanup(): void {
+    const accounts = this.#getSelectedAccounts();
+    if (accounts.length === 0) {
+      return;
+    }
+
+    this.update((state) => {
+      const balances = state.assetsBalance as Record<
+        string,
+        Record<string, AssetBalance>
+      >;
+      dropBalancesOutsideAccountScopes(balances, accounts);
+    });
+  }
+
   async #runSpamCleanup(): Promise<void> {
     try {
       const shouldRun =
@@ -3032,7 +3058,7 @@ export class AssetsController extends BaseController<
             );
             const nativeAssetIdsForAccount = account
               ? this.#getNativeAssetIdsForAccount(account)
-              : this.#getNativeAssetIdsForEnabledChains();
+              : [];
             for (const nativeAssetId of nativeAssetIdsForAccount) {
               if (
                 !Object.prototype.hasOwnProperty.call(

packages/assets-controller/src/migrations/dropBalancesOutsideAccountScopes.test.ts

@@ -0,0 +1,216 @@
+import type { InternalAccount } from '@metamask/keyring-internal-api';
+
+import { createMockInternalAccount } from '../__fixtures__/MockAssetControllerMessenger.js';
+import type { Caip19AssetId } from '../types.js';
+import { dropBalancesOutsideAccountScopes } from './dropBalancesOutsideAccountScopes.js';
+
+const BITCOIN_ACCOUNT_ID = 'bitcoin-account-id';
+const EVM_ACCOUNT_ID = 'evm-account-id';
+const BITCOIN_NATIVE =
+  'bip122:000000000019d6689c085ae165831e93/slip44:0' as Caip19AssetId;
+const SOLANA_NATIVE =
+  'solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp/slip44:501' as Caip19AssetId;
+const MAINNET_NATIVE = 'eip155:1/slip44:60' as Caip19AssetId;
+const LINEA_NATIVE = 'eip155:59144/slip44:60' as Caip19AssetId;
+
+const bitcoinAccount = (scopes: InternalAccount['scopes']): InternalAccount =>
+  createMockInternalAccount({
+    id: BITCOIN_ACCOUNT_ID,
+    address: 'bc1qdy9vjsg26rrd8tkqr7f26hlhd9524wyjkynemj',
+    type: 'bip122:p2wpkh',
+    scopes,
+  });
+
+const evmAccount = (scopes: InternalAccount['scopes']): InternalAccount =>
+  createMockInternalAccount({ id: EVM_ACCOUNT_ID, scopes });
+
+describe('dropBalancesOutsideAccountScopes', () => {
+  it('drops balances whose chain namespace is outside the account scopes', () => {
+    const balances = {
+      [BITCOIN_ACCOUNT_ID]: {
+        [BITCOIN_NATIVE]: { amount: '0.00000404' },
+        [SOLANA_NATIVE]: { amount: '0' },
+        [MAINNET_NATIVE]: { amount: '0' },
+      },
+      [EVM_ACCOUNT_ID]: {
+        [MAINNET_NATIVE]: { amount: '0.02' },
+        [LINEA_NATIVE]: { amount: '0' },
+        [SOLANA_NATIVE]: { amount: '0' },
+      },
+    };
+
+    const removed = dropBalancesOutsideAccountScopes(balances, [
+      bitcoinAccount(['bip122:000000000019d6689c085ae165831e93']),
+      evmAccount(['eip155:0']),
+    ]);
+
+    expect(removed).toBe(true);
+    expect(balances[BITCOIN_ACCOUNT_ID]).toStrictEqual({
+      [BITCOIN_NATIVE]: { amount: '0.00000404' },
+    });
+    expect(balances[EVM_ACCOUNT_ID]).toStrictEqual({
+      [MAINNET_NATIVE]: { amount: '0.02' },
+      [LINEA_NATIVE]: { amount: '0' },
+    });
+  });
+
+  it('keeps every EVM chain when the account scope is the eip155 wildcard', () => {
+    const balances = {
+      [EVM_ACCOUNT_ID]: {
+        [MAINNET_NATIVE]: { amount: '1' },
+        [LINEA_NATIVE]: { amount: '2' },
+      },
+    };
+
+    const removed = dropBalancesOutsideAccountScopes(balances, [
+      evmAccount(['eip155:0']),
+    ]);
+
+    expect(removed).toBe(false);
+    expect(balances[EVM_ACCOUNT_ID]).toStrictEqual({
+      [MAINNET_NATIVE]: { amount: '1' },
+      [LINEA_NATIVE]: { amount: '2' },
+    });
+  });
+
+  it('keeps balances for every namespace the account scopes include', () => {
+    const balances = {
+      [EVM_ACCOUNT_ID]: {
+        [MAINNET_NATIVE]: { amount: '1' },
+        [SOLANA_NATIVE]: { amount: '2' },
+        [BITCOIN_NATIVE]: { amount: '3' },
+      },
+    };
+
+    const removed = dropBalancesOutsideAccountScopes(balances, [
+      evmAccount(['eip155:1', 'solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp']),
+    ]);
+
+    expect(removed).toBe(true);
+    expect(balances[EVM_ACCOUNT_ID]).toStrictEqual({
+      [MAINNET_NATIVE]: { amount: '1' },
+      [SOLANA_NATIVE]: { amount: '2' },
+    });
+  });
+
+  it('leaves malformed scopes and balance keys alone', () => {
+    const balances = {
+      [EVM_ACCOUNT_ID]: {
+        [MAINNET_NATIVE]: { amount: '1' },
+        [SOLANA_NATIVE]: { amount: '0' },
+        'not-a-caip-id': { amount: '0' },
+      },
+    } as Record<string, Record<Caip19AssetId, { amount: string }>>;
+
+    const removed = dropBalancesOutsideAccountScopes(balances, [
+      evmAccount(['eip155:1', 'garbage' as `${string}:${string}`]),
+    ]);
+
+    expect(removed).toBe(true);
+    expect(balances[EVM_ACCOUNT_ID]).toStrictEqual({
+      [MAINNET_NATIVE]: { amount: '1' },
+      'not-a-caip-id': { amount: '0' },
+    });
+  });
+
+  it('leaves an account untouched when none of its scopes are valid CAIP-2 ids', () => {
+    const balances = {
+      [EVM_ACCOUNT_ID]: {
+        [MAINNET_NATIVE]: { amount: '1' },
+        [SOLANA_NATIVE]: { amount: '0' },
+      },
+    };
+
+    const removed = dropBalancesOutsideAccountScopes(balances, [
+      evmAccount(['garbage' as `${string}:${string}`]),
+    ]);
+
+    expect(removed).toBe(false);
+    expect(balances[EVM_ACCOUNT_ID]).toStrictEqual({
+      [MAINNET_NATIVE]: { amount: '1' },
+      [SOLANA_NATIVE]: { amount: '0' },
+    });
+  });
+
+  it('leaves an account whose scopes are missing untouched', () => {
+    const balances = {
+      [BITCOIN_ACCOUNT_ID]: {
+        [BITCOIN_NATIVE]: { amount: '1' },
+        [SOLANA_NATIVE]: { amount: '0' },
+      },
+    };
+    const { scopes: _scopes, ...accountWithoutScopes } = bitcoinAccount([
+      'bip122:000000000019d6689c085ae165831e93',
+    ]);
+
+    const removed = dropBalancesOutsideAccountScopes(balances, [
+      accountWithoutScopes as InternalAccount,
+    ]);
+
+    expect(removed).toBe(false);
+    expect(balances[BITCOIN_ACCOUNT_ID]).toStrictEqual({
+      [BITCOIN_NATIVE]: { amount: '1' },
+      [SOLANA_NATIVE]: { amount: '0' },
+    });
+  });
+
+  it('leaves an account with no scopes untouched', () => {
+    const balances = {
+      [BITCOIN_ACCOUNT_ID]: {
+        [BITCOIN_NATIVE]: { amount: '1' },
+        [SOLANA_NATIVE]: { amount: '0' },
+      },
+    };
+
+    const removed = dropBalancesOutsideAccountScopes(balances, [
+      bitcoinAccount([]),
+    ]);
+
+    expect(removed).toBe(false);
+    expect(balances[BITCOIN_ACCOUNT_ID]).toStrictEqual({
+      [BITCOIN_NATIVE]: { amount: '1' },
+      [SOLANA_NATIVE]: { amount: '0' },
+    });
+  });
+
+  it('leaves an account that has no stored balances untouched', () => {
+    const balances = {
+      [EVM_ACCOUNT_ID]: {
+        [MAINNET_NATIVE]: { amount: '1' },
+      },
+    };
+
+    const removed = dropBalancesOutsideAccountScopes(balances, [
+      createMockInternalAccount({
+        id: 'missing-account',
+        scopes: ['eip155:1'],
+      }),
+    ]);
+
+    expect(removed).toBe(false);
+    expect(balances).toStrictEqual({
+      [EVM_ACCOUNT_ID]: {
+        [MAINNET_NATIVE]: { amount: '1' },
+      },
+    });
+  });
+
+  it('leaves accounts that were not passed in untouched', () => {
+    const balances = {
+      [BITCOIN_ACCOUNT_ID]: {
+        [BITCOIN_NATIVE]: { amount: '1' },
+        [SOLANA_NATIVE]: { amount: '0' },
+      },
+    };
+
+    const removed = dropBalancesOutsideAccountScopes(balances, [
+      evmAccount(['eip155:1']),
+    ]);
+
+    expect(removed).toBe(false);
+    expect(balances[BITCOIN_ACCOUNT_ID]).toStrictEqual({
+      [BITCOIN_NATIVE]: { amount: '1' },
+      [SOLANA_NATIVE]: { amount: '0' },
+    });
+  });
+});

packages/assets-controller/src/migrations/dropBalancesOutsideAccountScopes.ts

@@ -0,0 +1,84 @@
+import type { InternalAccount } from '@metamask/keyring-internal-api';
+import {
+  isCaipAssetType,
+  isCaipChainId,
+  parseCaipAssetType,
+  parseCaipChainId,
+} from '@metamask/utils';
+
+import type { AccountId, AssetBalance, Caip19AssetId } from '../types.js';
+
+/**
+ * Remove balance entries the given accounts cannot own.
+ *
+ * Ownership is the chain namespace of the account's scopes, not the
+ * currently enabled chains. `eip155:0` keeps every EVM chain, including
+ * ones the user has turned off. An account with no scopes is left untouched,
+ * as is any account that is not in `accounts`.
+ *
+ * Mutates `balances`. Intended to run inside an Immer `update` draft, once
+ * per unlock, the same way spam cleanup is applied.
+ *
+ * @param balances - `assetsBalance` being updated.
+ * @param accounts - Accounts whose scopes define the namespaces to keep.
+ * @returns True when at least one balance entry was removed.
+ */
+export function dropBalancesOutsideAccountScopes(
+  balances: Record<AccountId, Record<Caip19AssetId, AssetBalance>>,
+  accounts: readonly InternalAccount[],
+): boolean {
+  let removed = false;
+  for (const account of accounts) {
+    if (dropAccountBalancesOutsideScopes(balances, account)) {
+      removed = true;
+    }
+  }
+  return removed;
+}
+
+/**
+ * Remove one account's balances whose chain namespace is outside its scopes.
+ *
+ * @param balances - `assetsBalance` being updated.
+ * @param account - Account whose scopes define the namespaces to keep.
+ * @returns True when at least one balance entry was removed.
+ */
+function dropAccountBalancesOutsideScopes(
+  balances: Record<AccountId, Record<Caip19AssetId, AssetBalance>>,
+  account: InternalAccount,
+): boolean {
+  // Persisted accounts can omit `scopes` even though the type requires it.
+  const scopes = account.scopes ?? [];
+  if (scopes.length === 0) {
+    return false;
+  }
+
+  const accountBalances = balances[account.id];
+  if (!accountBalances) {
+    return false;
+  }
+
+  const namespaces = new Set<string>();
+  for (const scope of scopes) {
+    if (isCaipChainId(scope)) {
+      namespaces.add(parseCaipChainId(scope).namespace);
+    }
+  }
+  if (namespaces.size === 0) {
+    return false;
+  }
+
+  let removed = false;
+  for (const assetId of Object.keys(accountBalances) as Caip19AssetId[]) {
+    // Malformed keys are left alone; this cleanup only targets valid CAIP-19
+    // entries stamped onto the wrong account.
+    if (!isCaipAssetType(assetId)) {
+      continue;
+    }
+    if (!namespaces.has(parseCaipAssetType(assetId).chain.namespace)) {
+      delete accountBalances[assetId];
+      removed = true;
+    }
+  }
+  return removed;
+}