From 5cffffdfcbef0982aceb0e4e32d93759bc550f33 Mon Sep 17 00:00:00 2001 From: Charly Chevalier Date: Thu, 17 Sep 2026 22:31:46 +0200 Subject: [PATCH 1/2] 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 From 3a46764769b9c30e36fa0281efc3572b8da55539 Mon Sep 17 00:00:00 2001 From: Charly Chevalier Date: Thu, 17 Sep 2026 22:50:03 +0200 Subject: [PATCH 2/2] refactor(snap-account-service): add SnapAccountCache --- .../src/SnapAccountCache.test.ts | 305 ++++++++++++++++++ .../src/SnapAccountCache.ts | 117 +++++++ .../src/SnapAccountService.test.ts | 136 -------- .../src/SnapAccountService.ts | 100 +----- 4 files changed, 431 insertions(+), 227 deletions(-) create mode 100644 packages/snap-account-service/src/SnapAccountCache.test.ts create mode 100644 packages/snap-account-service/src/SnapAccountCache.ts diff --git a/packages/snap-account-service/src/SnapAccountCache.test.ts b/packages/snap-account-service/src/SnapAccountCache.test.ts new file mode 100644 index 00000000000..47ccd9ea726 --- /dev/null +++ b/packages/snap-account-service/src/SnapAccountCache.test.ts @@ -0,0 +1,305 @@ +import type { + AccountsControllerAccountsAddedEvent, + AccountsControllerAccountsRemovedEvent, + AccountsControllerGetStateAction, + AccountsControllerState, +} from '@metamask/accounts-controller'; +import { Messenger, MOCK_ANY_NAMESPACE } from '@metamask/messenger'; +import type { + MockAnyNamespace, + MessengerActions, + MessengerEvents, +} from '@metamask/messenger'; +import type { SnapId } from '@metamask/snaps-sdk'; + +import { SnapAccountCache } from './SnapAccountCache.js'; +import type { SnapAccountServiceMessenger } from './SnapAccountService.js'; + +type RootMessenger = Messenger< + MockAnyNamespace, + MessengerActions, + MessengerEvents +>; + +/** + * Constructs the root messenger for the cache under test. + * + * @returns The root messenger. + */ +function getRootMessenger(): RootMessenger { + return new Messenger({ namespace: MOCK_ANY_NAMESPACE }); +} + +/** + * Constructs the messenger for the cache under test, and delegates all + * required external actions and events from the root messenger to it. + * + * @param rootMessenger - The root messenger. + * @returns The cache messenger. + */ +function getMessenger( + rootMessenger: RootMessenger, +): SnapAccountServiceMessenger { + const messenger = new Messenger({ + namespace: 'SnapAccountService', + parent: rootMessenger, + }); + rootMessenger.delegate({ + messenger, + actions: ['AccountsController:getState'], + events: [ + 'AccountsController:accountsAdded', + 'AccountsController:accountsRemoved', + ], + }); + return messenger; +} + +/** + * Builds a minimal `AccountsControllerState` whose `internalAccounts.accounts` + * maps each given account ID to an account owned by the given Snap ID. + * + * @param accounts - The accounts to include. + * @returns A minimal `AccountsControllerState`. + */ +function buildAccountsState( + accounts: { id: string; snapId?: string }[], +): AccountsControllerState { + const accountsRecord = Object.fromEntries( + accounts.map(({ id, snapId }) => [ + id, + { id, metadata: snapId ? { snap: { id: snapId } } : {} }, + ]), + ); + return { + internalAccounts: { accounts: accountsRecord }, + } as unknown as AccountsControllerState; +} + +/** + * Publishes an `AccountsController:accountsAdded` event. + * + * @param rootMessenger - The root messenger. + * @param accounts - The accounts that were added. + */ +function publishAccountsAdded( + rootMessenger: RootMessenger, + accounts: { id: string; snapId?: string }[], +): void { + rootMessenger.publish( + 'AccountsController:accountsAdded', + accounts.map(({ id, snapId }) => ({ + id, + metadata: snapId ? { snap: { id: snapId } } : {}, + })) as AccountsControllerAccountsAddedEvent['payload'][0], + ); +} + +/** + * Publishes an `AccountsController:accountsRemoved` event. + * + * @param rootMessenger - The root messenger. + * @param accountIds - The IDs of the accounts that were removed. + */ +function publishAccountsRemoved( + rootMessenger: RootMessenger, + accountIds: string[], +): void { + rootMessenger.publish( + 'AccountsController:accountsRemoved', + accountIds as AccountsControllerAccountsRemovedEvent['payload'][0], + ); +} + +type Mocks = { + // eslint-disable-next-line @typescript-eslint/naming-convention + AccountsController: { + getState: jest.MockedFunction<() => AccountsControllerState>; + }; +}; + +/** + * Constructs the cache under test with sensible defaults. + * + * @param args - The arguments to this function. + * @param args.accounts - Accounts to seed into the `AccountsController:getState` mock. + * @returns The new cache, root messenger, cache messenger, and mocks. + */ +function setup({ + accounts = [], +}: { + accounts?: { id: string; snapId?: string }[]; +} = {}): { + cache: SnapAccountCache; + rootMessenger: RootMessenger; + mocks: Mocks; +} { + const rootMessenger = getRootMessenger(); + const messenger = getMessenger(rootMessenger); + + const mocks: Mocks = { + AccountsController: { + getState: jest.fn().mockReturnValue(buildAccountsState(accounts)), + }, + }; + + rootMessenger.registerActionHandler( + 'AccountsController:getState', + mocks.AccountsController.getState as never, + ); + + const cache = new SnapAccountCache(messenger); + + return { cache, rootMessenger, mocks }; +} + +const MOCK_SNAP_ID = 'npm:@metamask/mock-snap' as SnapId; +const MOCK_OTHER_SNAP_ID = 'npm:@metamask/other-snap' as SnapId; +const MOCK_ACCOUNT_ID = '00000000-0000-0000-0000-000000000001'; +const MOCK_OTHER_ACCOUNT_ID = '00000000-0000-0000-0000-000000000002'; +const MOCK_NO_SNAP_ACCOUNT_ID = '00000000-0000-0000-0000-000000000003'; + +describe('SnapAccountCache', () => { + describe('getSnapId', () => { + it('returns undefined for an unknown account before the cache is built', () => { + const { cache } = setup(); + + expect(cache.getSnapId(MOCK_ACCOUNT_ID)).toBeUndefined(); + }); + + it('lazily builds the cache from AccountsController state on first use', () => { + const { cache, mocks } = setup({ + accounts: [{ id: MOCK_ACCOUNT_ID, snapId: MOCK_SNAP_ID as string }], + }); + + expect(mocks.AccountsController.getState).not.toHaveBeenCalled(); + + expect(cache.getSnapId(MOCK_ACCOUNT_ID)).toBe(MOCK_SNAP_ID); + + expect(mocks.AccountsController.getState).toHaveBeenCalledTimes(1); + }); + + it('does not rebuild the cache on subsequent calls', () => { + const { cache, mocks } = setup({ + accounts: [{ id: MOCK_ACCOUNT_ID, snapId: MOCK_SNAP_ID as string }], + }); + + cache.getSnapId(MOCK_ACCOUNT_ID); + cache.getSnapId(MOCK_ACCOUNT_ID); + + expect(mocks.AccountsController.getState).toHaveBeenCalledTimes(1); + }); + + it('skips accounts without a snap ID when building the cache', () => { + const { cache } = setup({ + accounts: [ + { id: MOCK_ACCOUNT_ID, snapId: MOCK_SNAP_ID as string }, + { id: MOCK_NO_SNAP_ACCOUNT_ID }, + ], + }); + + expect(cache.getSnapId(MOCK_ACCOUNT_ID)).toBe(MOCK_SNAP_ID); + expect(cache.getSnapId(MOCK_NO_SNAP_ACCOUNT_ID)).toBeUndefined(); + }); + }); + + describe('invalidate', () => { + it('causes the next getSnapId call to rebuild from fresh state', () => { + const { cache, mocks } = setup({ + accounts: [{ id: MOCK_ACCOUNT_ID, snapId: MOCK_SNAP_ID as string }], + }); + + expect(cache.getSnapId(MOCK_ACCOUNT_ID)).toBe(MOCK_SNAP_ID); + expect(mocks.AccountsController.getState).toHaveBeenCalledTimes(1); + + mocks.AccountsController.getState.mockReturnValue( + buildAccountsState([ + { id: MOCK_ACCOUNT_ID, snapId: MOCK_OTHER_SNAP_ID as string }, + ]), + ); + + cache.invalidate(); + + expect(cache.getSnapId(MOCK_ACCOUNT_ID)).toBe(MOCK_OTHER_SNAP_ID); + expect(mocks.AccountsController.getState).toHaveBeenCalledTimes(2); + }); + }); + + describe('on AccountsController:accountsAdded', () => { + it('adds Snap-owned accounts to an initialized cache', () => { + const { cache, rootMessenger } = setup(); + + cache.getSnapId(MOCK_ACCOUNT_ID); // initialize + + publishAccountsAdded(rootMessenger, [ + { id: MOCK_ACCOUNT_ID, snapId: MOCK_SNAP_ID as string }, + ]); + + expect(cache.getSnapId(MOCK_ACCOUNT_ID)).toBe(MOCK_SNAP_ID); + }); + + it('skips accounts without a snap ID', () => { + const { cache, rootMessenger } = setup(); + + cache.getSnapId(MOCK_ACCOUNT_ID); // initialize + + publishAccountsAdded(rootMessenger, [{ id: MOCK_NO_SNAP_ACCOUNT_ID }]); + + expect(cache.getSnapId(MOCK_NO_SNAP_ACCOUNT_ID)).toBeUndefined(); + }); + + it('is a no-op when the cache is not initialized', () => { + const { cache, rootMessenger, mocks } = setup({ + accounts: [{ id: MOCK_ACCOUNT_ID, snapId: MOCK_SNAP_ID as string }], + }); + + publishAccountsAdded(rootMessenger, [ + { id: MOCK_ACCOUNT_ID, snapId: MOCK_OTHER_SNAP_ID as string }, + ]); + + // The cache is still uninitialized — the next getSnapId call rebuilds + // from state, which still maps MOCK_ACCOUNT_ID to MOCK_SNAP_ID. + expect(cache.getSnapId(MOCK_ACCOUNT_ID)).toBe(MOCK_SNAP_ID); + expect(mocks.AccountsController.getState).toHaveBeenCalledTimes(1); + }); + }); + + describe('on AccountsController:accountsRemoved', () => { + it('removes accounts from an initialized cache', () => { + const { cache, rootMessenger } = setup({ + accounts: [{ id: MOCK_ACCOUNT_ID, snapId: MOCK_SNAP_ID as string }], + }); + + cache.getSnapId(MOCK_ACCOUNT_ID); // initialize + + publishAccountsRemoved(rootMessenger, [MOCK_ACCOUNT_ID]); + + expect(cache.getSnapId(MOCK_ACCOUNT_ID)).toBeUndefined(); + }); + + it('is a no-op for unknown account IDs', () => { + const { cache, rootMessenger } = setup({ + accounts: [{ id: MOCK_ACCOUNT_ID, snapId: MOCK_SNAP_ID as string }], + }); + + cache.getSnapId(MOCK_ACCOUNT_ID); // initialize + + publishAccountsRemoved(rootMessenger, [MOCK_OTHER_ACCOUNT_ID]); + + expect(cache.getSnapId(MOCK_ACCOUNT_ID)).toBe(MOCK_SNAP_ID); + }); + + it('is a no-op when the cache is not initialized', () => { + const { cache, rootMessenger, mocks } = setup({ + accounts: [{ id: MOCK_ACCOUNT_ID, snapId: MOCK_SNAP_ID as string }], + }); + + publishAccountsRemoved(rootMessenger, [MOCK_ACCOUNT_ID]); + + // The cache is still uninitialized — the next getSnapId call rebuilds + // from state, which still has MOCK_ACCOUNT_ID mapped to MOCK_SNAP_ID. + expect(cache.getSnapId(MOCK_ACCOUNT_ID)).toBe(MOCK_SNAP_ID); + expect(mocks.AccountsController.getState).toHaveBeenCalledTimes(1); + }); + }); +}); diff --git a/packages/snap-account-service/src/SnapAccountCache.ts b/packages/snap-account-service/src/SnapAccountCache.ts new file mode 100644 index 00000000000..a65ffbef21c --- /dev/null +++ b/packages/snap-account-service/src/SnapAccountCache.ts @@ -0,0 +1,117 @@ +import { KeyringAccount } from '@metamask/keyring-api'; +import type { AccountId } from '@metamask/keyring-utils'; +import { SnapId } from '@metamask/snaps-sdk'; + +import type { SnapAccountServiceMessenger } from './SnapAccountService.js'; + +// Re-define it here to avoid pulling dependency. +type InternalAccount = KeyringAccount & { + metadata?: { + snap?: { + id?: string; + }; + }; +}; + +/** + * Cache mapping each Snap-owned account ID to the ID of the Snap that owns + * it, derived from `AccountsController` state. + * + * The cache is built lazily on first use and kept in sync incrementally via + * `AccountsController:accountsAdded` / `AccountsController:accountsRemoved` + * events. Calling {@link SnapAccountCache.invalidate} resets the initialized + * flag so the next use triggers a full rebuild — used on wallet unlock, when + * changes that occurred while locked are not reflected by those events. + */ +export class SnapAccountCache { + readonly #messenger: SnapAccountServiceMessenger; + + #entries: Map = new Map(); + + #initialized = false; + + constructor(messenger: SnapAccountServiceMessenger) { + this.#messenger = messenger; + + messenger.subscribe('AccountsController:accountsAdded', (accounts) => + this.#handleAccountsAdded(accounts), + ); + messenger.subscribe('AccountsController:accountsRemoved', (accountIds) => + this.#handleAccountsRemoved(accountIds), + ); + } + + /** + * Invalidates the cache so the next {@link SnapAccountCache.getSnapId} call + * triggers a full rebuild from `AccountsController` state. + */ + invalidate(): void { + this.#initialized = false; + } + + /** + * Returns the Snap ID that owns the given account, or `undefined` if the + * account is not owned by any Snap. Lazily builds the cache on first call + * after construction or invalidation. + * + * @param accountId - The account ID to look up. + * @returns The Snap ID that owns the account, or `undefined`. + */ + getSnapId(accountId: AccountId): SnapId | undefined { + this.#ensureReady(); + return this.#entries.get(accountId); + } + + /** + * Builds the cache from `AccountsController` state if it is not already + * initialized. + */ + #ensureReady(): void { + if (this.#initialized) { + 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.#entries = cache; + this.#initialized = true; + } + + /** + * Adds the given accounts to the cache. No-op when the cache is not yet + * initialized — the next rebuild will read fresh state instead. + * + * @param accounts - The accounts that were added. + */ + #handleAccountsAdded(accounts: InternalAccount[]): void { + if (!this.#initialized) { + return; + } + for (const account of accounts) { + const snapId = account.metadata?.snap?.id; + if (snapId) { + this.#entries.set(account.id, snapId as SnapId); + } + } + } + + /** + * Removes the given account IDs from the cache. No-op when the cache is not + * yet initialized — the next rebuild will read fresh state instead. + * + * @param accountIds - The IDs of the accounts that were removed. + */ + #handleAccountsRemoved(accountIds: AccountId[]): void { + if (!this.#initialized) { + return; + } + for (const accountId of accountIds) { + this.#entries.delete(accountId); + } + } +} diff --git a/packages/snap-account-service/src/SnapAccountService.test.ts b/packages/snap-account-service/src/SnapAccountService.test.ts index d3182ed7bae..6f05e3a22f1 100644 --- a/packages/snap-account-service/src/SnapAccountService.test.ts +++ b/packages/snap-account-service/src/SnapAccountService.test.ts @@ -1405,142 +1405,6 @@ describe('SnapAccountService', () => { expect(listener).not.toHaveBeenCalled(); }, ); - - it('picks up added/removed accounts from AccountsController:accountsAdded and :accountsRemoved', async () => { - // Initially the Snap does not own the account, so the update is dropped. - const { service, rootMessenger } = await setup({ - accounts: [ - { id: MOCK_ACCOUNT_ID, snapId: MOCK_OTHER_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; - - let result = await service.handleKeyringSnapMessage(MOCK_SNAP_ID, { - method: KeyringEvent.AccountBalancesUpdated, - params: payload, - } as unknown as SnapMessage); - expect(result).toBeNull(); - expect(listener).not.toHaveBeenCalled(); - - // The account is added for this Snap — the cache picks it up from - // `accountsAdded` and the next update is forwarded. A no-Snap account - // is included to verify such accounts are skipped when updating the - // cache incrementally. - publishAccountsAdded(rootMessenger, [ - { id: MOCK_ACCOUNT_ID, snapId: MOCK_SNAP_ID as string }, - { id: MOCK_NO_SNAP_ACCOUNT_ID }, - ]); - - result = await service.handleKeyringSnapMessage(MOCK_SNAP_ID, { - method: KeyringEvent.AccountBalancesUpdated, - params: payload, - } as unknown as SnapMessage); - expect(result).toBeNull(); - expect(listener).toHaveBeenCalledTimes(1); - expect(listener).toHaveBeenCalledWith({ - balances: { - [MOCK_ACCOUNT_ID]: payload.balances[MOCK_ACCOUNT_ID], - }, - }); - - // The account is removed — the cache drops it and the next update is - // dropped again (fail closed). - publishAccountsRemoved(rootMessenger, [MOCK_ACCOUNT_ID]); - - listener.mockClear(); - result = await service.handleKeyringSnapMessage(MOCK_SNAP_ID, { - method: KeyringEvent.AccountBalancesUpdated, - params: payload, - } as unknown as SnapMessage); - 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 b3626ed6509..039af5f5f87 100644 --- a/packages/snap-account-service/src/SnapAccountService.ts +++ b/packages/snap-account-service/src/SnapAccountService.ts @@ -72,6 +72,7 @@ import { assertStruct } from '@metamask/utils'; import { reportError, withSafeError } from './errors.js'; import { projectLogger as log } from './logger.js'; +import { SnapAccountCache } from './SnapAccountCache.js'; import type { SnapAccountServiceEnsureReadyAction, SnapAccountServiceEnsureMigratedAction, @@ -273,23 +274,14 @@ export class SnapAccountService { readonly #tracker: SnapTracker; + readonly #cache: SnapAccountCache; + readonly #client: KeyringInternalSnapClient; #migrated = false; #migratePromise: Promise | null = null; - /** - * Cache mapping each Snap-owned account ID to the ID of the Snap that owns - * it, derived from `AccountsController` state. - */ - #accountSnapIds: Map = new Map(); - - /** - * Whether `#accountSnapIds` has been populated yet. - */ - #accountSnapCacheInitialized = false; - /** * Constructs a new {@link SnapAccountService}. * @@ -305,6 +297,7 @@ export class SnapAccountService { config?.snapPlatformWatcher, ); this.#tracker = new SnapTracker(messenger); + this.#cache = new SnapAccountCache(messenger); this.#client = new KeyringInternalSnapClient({ messenger: messenger.buildChild({ namespace: 'KeyringInternalSnapClient', @@ -317,22 +310,6 @@ export class SnapAccountService { MESSENGER_EXPOSED_METHODS, ); - // Keep the Snap-ownership cache in sync as accounts are added/removed. - // The initial cache is built lazily on first use (see - // `#initAccountSnapCache`) rather than in the constructor, so that this - // service does not force clients to instantiate `AccountsController` - // before it. This keeps the account data update event path synchronous — - // the cache is a plain `Map` read. The granular `accountsAdded` / - // `accountsRemoved` events (batch-compatible) update the cache - // incrementally instead of rebuilding it from full state on every change. - this.#messenger.subscribe('AccountsController:accountsAdded', (accounts) => - this.#addAccountsToCache(accounts), - ); - this.#messenger.subscribe( - 'AccountsController:accountsRemoved', - (accountIds) => this.#removeAccountsFromCache(accountIds), - ); - this.#messenger.subscribe( 'AccountTreeController:selectedAccountGroupChange', (groupId) => this.#handleSelectedAccountGroupChange(groupId), @@ -378,11 +355,10 @@ 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; + // Invalidate the Snap account cache to ensure it will be rebuilt on + // next access. This allows us to always rebuild it after the `AccountsController` + // has re-synced with the `KeyringController`. + this.#cache.invalidate(); // eslint-disable-next-line no-void void this.ensureMigrated().then( @@ -986,10 +962,9 @@ export class SnapAccountService { event: AccountDataUpdatedKeyringEvent, entries: Record, ): Record { - this.#ensureAccountSnapCacheIsReady(); const filtered: Record = {}; for (const [accountId, value] of Object.entries(entries)) { - if (this.#accountSnapIds.get(accountId) === snapId) { + if (this.#cache.getSnapId(accountId) === snapId) { filtered[accountId] = value; } else { log( @@ -1000,63 +975,6 @@ export class SnapAccountService { return filtered; } - /** - * Adds the given accounts to the Snap-ownership cache. - * - * @param accounts - The accounts that were added. - */ - #addAccountsToCache( - accounts: AccountsControllerAccountsAddedEvent['payload'][0], - ): void { - if (!this.#accountSnapCacheInitialized) { - return; - } - for (const account of accounts) { - const snapId = account.metadata?.snap?.id; - if (snapId) { - this.#accountSnapIds.set(account.id, snapId as SnapId); - } - } - } - - /** - * Removes the given account IDs from the Snap-ownership cache. - * - * @param accountIds - The IDs of the accounts that were removed. - */ - #removeAccountsFromCache( - accountIds: AccountsControllerAccountsRemovedEvent['payload'][0], - ): void { - if (!this.#accountSnapCacheInitialized) { - return; - } - for (const accountId of accountIds) { - this.#accountSnapIds.delete(accountId); - } - } - - /** - * 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. - */ - #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 /** * Forwards the accounts of the given account group to the Snap keyring.