fix: filter account data updates by Snap ownership

packages/snap-account-service/src/SnapAccountService.test.ts

@@ -1467,6 +1467,80 @@ describe('SnapAccountService', () => {
       expect(result).toBeNull();
       expect(listener).not.toHaveBeenCalled();
     });
+
+    it('ignores accountsAdded events while cache is not initialized', async () => {
+      // Start with the account owned by MOCK_SNAP_ID in AccountsController state.
+      const { service, rootMessenger } = await setup({
+        accounts: [{ id: MOCK_ACCOUNT_ID, snapId: MOCK_SNAP_ID as string }],
+      });
+      const listener = jest.fn();
+      rootMessenger.subscribe(
+        'SnapAccountService:accountBalancesUpdated',
+        listener,
+      );
+
+      const payload = {
+        balances: {
+          [MOCK_ACCOUNT_ID]: {
+            'eip155:1/slip44:60': { amount: '1', unit: 'ETH' },
+          },
+        },
+      } satisfies AccountBalancesUpdatedEventPayload;
+
+      // Unlock invalidates the cache — accountsAdded fired before the cache is
+      // rebuilt should be ignored (the rebuild on next use will read fresh state).
+      publishUnlock(rootMessenger);
+      await flushMicrotasks();
+      publishAccountsAdded(rootMessenger, [
+        { id: MOCK_ACCOUNT_ID, snapId: MOCK_OTHER_SNAP_ID as string },
+      ]);
+
+      // The first use after unlock rebuilds from AccountsController state, which
+      // still maps MOCK_ACCOUNT_ID to MOCK_SNAP_ID — the ignored accountsAdded
+      // (which would have reassigned it to MOCK_OTHER_SNAP_ID) was correctly dropped.
+      const result = await service.handleKeyringSnapMessage(MOCK_SNAP_ID, {
+        method: KeyringEvent.AccountBalancesUpdated,
+        params: payload,
+      } as unknown as SnapMessage);
+      expect(result).toBeNull();
+      expect(listener).toHaveBeenCalledTimes(1);
+    });
+
+    it('ignores accountsRemoved events while cache is not initialized', async () => {
+      // Start with the account owned by MOCK_SNAP_ID in AccountsController state.
+      const { service, rootMessenger } = await setup({
+        accounts: [{ id: MOCK_ACCOUNT_ID, snapId: MOCK_SNAP_ID as string }],
+      });
+      const listener = jest.fn();
+      rootMessenger.subscribe(
+        'SnapAccountService:accountBalancesUpdated',
+        listener,
+      );
+
+      const payload = {
+        balances: {
+          [MOCK_ACCOUNT_ID]: {
+            'eip155:1/slip44:60': { amount: '1', unit: 'ETH' },
+          },
+        },
+      } satisfies AccountBalancesUpdatedEventPayload;
+
+      // Unlock invalidates the cache — accountsRemoved fired before the cache is
+      // rebuilt should be ignored (the rebuild on next use will read fresh state).
+      publishUnlock(rootMessenger);
+      await flushMicrotasks();
+      publishAccountsRemoved(rootMessenger, [MOCK_ACCOUNT_ID]);
+
+      // The first use after unlock rebuilds from AccountsController state, which
+      // still has MOCK_ACCOUNT_ID owned by MOCK_SNAP_ID — the ignored
+      // accountsRemoved was correctly dropped.
+      const result = await service.handleKeyringSnapMessage(MOCK_SNAP_ID, {
+        method: KeyringEvent.AccountBalancesUpdated,
+        params: payload,
+      } as unknown as SnapMessage);
+      expect(result).toBeNull();
+      expect(listener).toHaveBeenCalledTimes(1);
+    });
   });
 
   describe('on AccountTreeController:selectedAccountGroupChange', () => {

packages/snap-account-service/src/SnapAccountService.ts

@@ -3,7 +3,6 @@ import type {
   AccountsControllerAccountsAddedEvent,
   AccountsControllerAccountsRemovedEvent,
   AccountsControllerGetStateAction,
-  AccountsControllerState,
 } from '@metamask/accounts-controller';
 import {
   SnapKeyring as LegacySnapKeyring,
@@ -379,6 +378,12 @@ export class SnapAccountService {
    * keyring.
    */
   #handleUnlock(): void {
+    // Invalidate the cache so the next use rebuilds it from fresh state.
+    // Changes that occurred while the wallet was locked are not reflected by
+    // the incremental accountsAdded/accountsRemoved events, so a full rebuild
+    // is needed on first use after unlock.
+    this.#accountSnapCacheInitialized = false;
+
     // eslint-disable-next-line no-void
     void this.ensureMigrated().then(
       async () => {
@@ -981,7 +986,7 @@ export class SnapAccountService {
     event: AccountDataUpdatedKeyringEvent,
     entries: Record<string, Value>,
   ): Record<string, Value> {
-    this.#initAccountSnapCache();
+    this.#ensureAccountSnapCacheIsReady();
     const filtered: Record<string, Value> = {};
     for (const [accountId, value] of Object.entries(entries)) {
       if (this.#accountSnapIds.get(accountId) === snapId) {
@@ -995,27 +1000,6 @@ export class SnapAccountService {
     return filtered;
   }
 
-  /**
-   * Rebuilds the Snap-ownership cache from `AccountsController` state.
-   *
-   * Used for lazy initialization on first use; subsequent updates are applied
-   * incrementally by {@link SnapAccountService.#addAccountsToCache} and
-   * {@link SnapAccountService.#removeAccountsFromCache}.
-   *
-   * @param state - The current `AccountsController` state.
-   */
-  #rebuildAccountSnapCache(state: AccountsControllerState): void {
-    const cache = new Map<AccountId, SnapId>();
-    for (const account of Object.values(state.internalAccounts.accounts)) {
-      const snapId = account.metadata?.snap?.id;
-      if (snapId) {
-        cache.set(account.id, snapId as SnapId);
-      }
-    }
-    this.#accountSnapIds = cache;
-    this.#accountSnapCacheInitialized = true;
-  }
-
   /**
    * Adds the given accounts to the Snap-ownership cache.
    *
@@ -1024,6 +1008,9 @@ export class SnapAccountService {
   #addAccountsToCache(
     accounts: AccountsControllerAccountsAddedEvent['payload'][0],
   ): void {
+    if (!this.#accountSnapCacheInitialized) {
+      return;
+    }
     for (const account of accounts) {
       const snapId = account.metadata?.snap?.id;
       if (snapId) {
@@ -1040,20 +1027,34 @@ export class SnapAccountService {
   #removeAccountsFromCache(
     accountIds: AccountsControllerAccountsRemovedEvent['payload'][0],
   ): void {
+    if (!this.#accountSnapCacheInitialized) {
+      return;
+    }
     for (const accountId of accountIds) {
       this.#accountSnapIds.delete(accountId);
     }
   }
 
   /**
-   * Lazily builds the Snap-ownership cache on first use.
+   * Ensures the Snap-ownership cache is populated. If not yet initialized,
+   * builds it from scratch from `AccountsController` state. Once initialized,
+   * incremental {@link SnapAccountService.#addAccountsToCache} /
+   * {@link SnapAccountService.#removeAccountsFromCache} events keep it current.
    */
-  #initAccountSnapCache(): void {
-    if (!this.#accountSnapCacheInitialized) {
-      this.#rebuildAccountSnapCache(
-        this.#messenger.call('AccountsController:getState'),
-      );
+  #ensureAccountSnapCacheIsReady(): void {
+    if (this.#accountSnapCacheInitialized) {
+      return;
     }
+    const state = this.#messenger.call('AccountsController:getState');
+    const cache = new Map<AccountId, SnapId>();
+    for (const account of Object.values(state.internalAccounts.accounts)) {
+      const snapId = account.metadata?.snap?.id;
+      if (snapId) {
+        cache.set(account.id, snapId as SnapId);
+      }
+    }
+    this.#accountSnapIds = cache;
+    this.#accountSnapCacheInitialized = true;
   }
 
   // eslint-disable-next-line jsdoc/require-returns