From 5cffffdfcbef0982aceb0e4e32d93759bc550f33 Mon Sep 17 00:00:00 2001 From: Charly Chevalier Date: Thu, 17 Sep 2026 22:31:46 +0200 Subject: [PATCH] refactor: invalidate cache on unlock --- .../src/SnapAccountService.test.ts | 74 +++++++++++++++++++ .../src/SnapAccountService.ts | 59 +++++++-------- 2 files changed, 104 insertions(+), 29 deletions(-) diff --git a/packages/snap-account-service/src/SnapAccountService.test.ts b/packages/snap-account-service/src/SnapAccountService.test.ts index 891b914fbf3..d3182ed7bae 100644 --- a/packages/snap-account-service/src/SnapAccountService.test.ts +++ b/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', () => { diff --git a/packages/snap-account-service/src/SnapAccountService.ts b/packages/snap-account-service/src/SnapAccountService.ts index 715c4c929f7..b3626ed6509 100644 --- a/packages/snap-account-service/src/SnapAccountService.ts +++ b/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, ): Record { - this.#initAccountSnapCache(); + this.#ensureAccountSnapCacheIsReady(); const filtered: Record = {}; 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(); - 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(); + 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