diff --git a/packages/multichain-account-service/CHANGELOG.md b/packages/multichain-account-service/CHANGELOG.md index 27129ce778..7d671f0431 100644 --- a/packages/multichain-account-service/CHANGELOG.md +++ b/packages/multichain-account-service/CHANGELOG.md @@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Changed +- **BREAKING:** Replace `KeyringController:withKeyring` with `KeyringController:withKeyringV2` for the Snap account providers ([#8732](https://github.com/MetaMask/core/pull/8732)) - **BREAKING:** The service messenger now requires the `SnapAccountService:ensureReady` action to be declared ([#8715](https://github.com/MetaMask/core/pull/8715)) - **BREAKING:** Delegate Snap platform readiness to `@metamask/snap-account-service` ([#8715](https://github.com/MetaMask/core/pull/8715)), ([#8752](https://github.com/MetaMask/core/pull/8752)) - Removed `MultichainAccountService.ensureCanUseSnapPlatform()` method and the corresponding `MultichainAccountService:ensureCanUseSnapPlatform` messenger action. diff --git a/packages/multichain-account-service/src/providers/BaseBip44AccountProvider.ts b/packages/multichain-account-service/src/providers/BaseBip44AccountProvider.ts index 5ac23e82d7..e7a63366bc 100644 --- a/packages/multichain-account-service/src/providers/BaseBip44AccountProvider.ts +++ b/packages/multichain-account-service/src/providers/BaseBip44AccountProvider.ts @@ -11,7 +11,6 @@ import type { } from '@metamask/keyring-api/v2'; import type { KeyringMetadata, - KeyringSelector, KeyringSelectorV2, } from '@metamask/keyring-controller'; import type { InternalAccount } from '@metamask/keyring-internal-api'; @@ -159,40 +158,6 @@ export abstract class BaseBip44AccountProvider< ) as unknown as Account; } - /** - * Run an operation against a V1 keyring selected by `selector`. - * - * Forwards to `KeyringController:withKeyring`. Use this for keyrings that - * have not yet migrated to the unified V2 `Keyring` interface (e.g. the - * snap keyring). - * - * @param selector - The selector identifying the keyring. - * @param operation - The operation to run with the selected keyring. - * @returns The result of the operation. - */ - protected async withKeyring( - selector: KeyringSelector, - operation: ({ - keyring, - metadata, - }: { - keyring: SelectedKeyring; - metadata: KeyringMetadata; - }) => Promise, - ): Promise { - const result = await this.messenger.call( - 'KeyringController:withKeyring', - selector, - ({ keyring, metadata }) => - operation({ - keyring: keyring as SelectedKeyring, - metadata, - }), - ); - - return result as CallbackResult; - } - /** * Run an operation against a V2 keyring selected by `selector`. * diff --git a/packages/multichain-account-service/src/providers/BtcAccountProvider.test.ts b/packages/multichain-account-service/src/providers/BtcAccountProvider.test.ts index c1e7dc05b1..bde6eba947 100644 --- a/packages/multichain-account-service/src/providers/BtcAccountProvider.test.ts +++ b/packages/multichain-account-service/src/providers/BtcAccountProvider.test.ts @@ -1,11 +1,7 @@ import { isBip44Account } from '@metamask/account-api'; -import type { SnapKeyring } from '@metamask/eth-snap-keyring'; import { AccountCreationType, BtcAccountType } from '@metamask/keyring-api'; import type { KeyringMetadata } from '@metamask/keyring-controller'; -import type { - EthKeyring, - InternalAccount, -} from '@metamask/keyring-internal-api'; +import type { InternalAccount } from '@metamask/keyring-internal-api'; import { SnapControllerState } from '@metamask/snaps-controllers'; import deepmerge from 'deepmerge'; @@ -69,9 +65,9 @@ class MockBtcKeyring { return Number(index); } - createAccount: SnapKeyring['createAccount'] = jest + createAccount = jest .fn() - .mockImplementation((_, { derivationPath, index, ...options }) => { + .mockImplementation(({ derivationPath, index, ...options }) => { // Determine the group index to use - either from derivationPath parsing, explicit index, or fallback let groupIndex: number; @@ -110,34 +106,32 @@ class MockBtcKeyring { return account; }); - createAccounts: SnapKeyring['createAccounts'] = jest - .fn() - .mockImplementation((_, options) => { - const groupIndices = - options.type === 'bip44:derive-index' - ? [options.groupIndex] - : toGroupIndexRangeArray(options.range); - - return groupIndices.map((groupIndex) => { - const found = this.accounts.find( - (account) => - isBip44Account(account) && - account.options.entropy.groupIndex === groupIndex, - ); - - if (found) { - return found; // Idempotent. - } - - const account = MockAccountBuilder.from(MOCK_BTC_P2WPKH_ACCOUNT_1) - .withUuid() - .withAddressSuffix(`${groupIndex}`) - .withGroupIndex(groupIndex) - .get(); - this.accounts.push(account); - return account; - }); + createAccounts = jest.fn().mockImplementation((options) => { + const groupIndices = + options.type === 'bip44:derive-index' + ? [options.groupIndex] + : toGroupIndexRangeArray(options.range); + + return groupIndices.map((groupIndex) => { + const found = this.accounts.find( + (account) => + isBip44Account(account) && + account.options.entropy.groupIndex === groupIndex, + ); + + if (found) { + return found; // Idempotent. + } + + const account = MockAccountBuilder.from(MOCK_BTC_P2WPKH_ACCOUNT_1) + .withUuid() + .withAddressSuffix(`${groupIndex}`) + .withGroupIndex(groupIndex) + .get(); + this.accounts.push(account); + return account; }); + }); } class MockBtcAccountProvider extends BtcAccountProvider { @@ -212,12 +206,10 @@ function setup({ ); messenger.registerActionHandler( - 'KeyringController:withKeyring', + 'KeyringController:withKeyringV2', async (_, operation) => operation({ - // We type-cast here, since `withKeyring` defaults to `EthKeyring` and the - // Snap keyring doesn't really implement this interface (this is expected). - keyring: keyring as unknown as EthKeyring, + keyring, metadata: keyring.metadata, }), ); @@ -243,8 +235,8 @@ function setup({ mocks: { handleRequest: mockHandleRequest, keyring: { - createAccount: keyring.createAccount as jest.Mock, - createAccounts: keyring.createAccounts as jest.Mock, + createAccount: keyring.createAccount, + createAccounts: keyring.createAccounts, }, trace: mockTrace, }, @@ -483,14 +475,11 @@ describe('BtcAccountProvider', () => { }); expect(newAccounts).toHaveLength(1); // Batch endpoint must be called, NOT the singular one. - expect(mocks.keyring.createAccounts).toHaveBeenCalledWith( - BtcAccountProvider.BTC_SNAP_ID, - { - type: AccountCreationType.Bip44DeriveIndex, - entropySource: MOCK_HD_KEYRING_1.metadata.id, - groupIndex: newGroupIndex, - }, - ); + expect(mocks.keyring.createAccounts).toHaveBeenCalledWith({ + type: AccountCreationType.Bip44DeriveIndex, + entropySource: MOCK_HD_KEYRING_1.metadata.id, + groupIndex: newGroupIndex, + }); expect(mocks.keyring.createAccount).not.toHaveBeenCalled(); }); diff --git a/packages/multichain-account-service/src/providers/SnapAccountProvider.test.ts b/packages/multichain-account-service/src/providers/SnapAccountProvider.test.ts index 70a8931ff6..6d4f4e972e 100644 --- a/packages/multichain-account-service/src/providers/SnapAccountProvider.test.ts +++ b/packages/multichain-account-service/src/providers/SnapAccountProvider.test.ts @@ -16,6 +16,7 @@ import type { KeyringCapabilities, } from '@metamask/keyring-api'; import type { EntropySourceId, KeyringAccount } from '@metamask/keyring-api'; +import type { KeyringMetadata } from '@metamask/keyring-controller'; import type { InternalAccount } from '@metamask/keyring-internal-api'; import type { JsonRpcRequest, SnapId } from '@metamask/snaps-sdk'; import deepmerge from 'deepmerge'; @@ -153,10 +154,12 @@ const setup = ({ config: configOverride = {}, messenger = getRootMessenger(), accounts = [], + keyring: keyringOverrides = {}, }: { config?: DeepPartial; messenger?: RootMessenger; accounts?: InternalAccount[]; + keyring?: { type?: string; snapId?: SnapId }; } = {}) => { const mocks = { AccountsController: { @@ -165,6 +168,9 @@ const setup = ({ ErrorReportingService: { captureException: jest.fn(), }, + KeyringController: { + withKeyringV2: jest.fn(), + }, SnapController: { handleKeyringRequest: { getAccount: jest.fn(), @@ -223,18 +229,30 @@ const setup = ({ ); const keyring = { + type: keyringOverrides.type ?? 'snap', + snapId: keyringOverrides.snapId ?? TEST_SNAP_ID, createAccount: jest.fn(), createAccounts: jest.fn(), - removeAccount: jest.fn(), + deleteAccount: jest.fn().mockResolvedValue(undefined), + lookupByAddress: jest + .fn() + .mockImplementation((address: string) => + accounts.map(asKeyringAccount).find((a) => a.address === address), + ), }; + const metadata = { id: 'mock-keyring-id', name: '' } as KeyringMetadata; + mocks.KeyringController.withKeyringV2.mockImplementation( + async (selector, operation) => { + if (selector.filter && !selector.filter(keyring, metadata)) { + throw new Error('No keyring matches the selector'); + } + return await operation({ keyring, metadata }); + }, + ); messenger.registerActionHandler( - 'KeyringController:withKeyring', - jest - .fn() - .mockImplementation( - async (_ /* selector */, operation) => await operation({ keyring }), - ), + 'KeyringController:withKeyringV2', + mocks.KeyringController.withKeyringV2, ); const serviceMessenger = getMultichainAccountServiceMessenger(messenger); @@ -865,10 +883,8 @@ describe('SnapAccountProvider', () => { mocks.SnapController.handleKeyringRequest.deleteAccount, ).toHaveBeenCalledWith(extraSnapAccount2.id); - // Should remove from keyring and recreate the missing account - expect(keyring.removeAccount).toHaveBeenCalledWith( - mockAccounts[1].address, - ); + // Should delete the missing account from the keyring (by id) before recreating it. + expect(keyring.deleteAccount).toHaveBeenCalledWith(mockAccounts[1].id); expect(createAccountsSpy).toHaveBeenCalledWith({ entropySource: mockAccounts[1].options.entropy.id, groupIndex: mockAccounts[1].options.entropy.groupIndex, @@ -949,6 +965,37 @@ describe('SnapAccountProvider', () => { }); }); + describe('withKeyringV2 selector', () => { + const mockAccounts = [ + MockAccountBuilder.from(MOCK_HD_ACCOUNT_1) + .withUuid() + .withSnapId(TEST_SNAP_ID) + .get(), + ].filter(isBip44Account); + + it('rejects when the keyring type is not a Snap keyring', async () => { + const { provider } = setup({ + accounts: mockAccounts, + keyring: { type: 'not-a-snap-keyring' }, + }); + + await expect(provider.resyncAccounts(mockAccounts)).rejects.toThrow( + 'No keyring matches the selector', + ); + }); + + it('rejects when the Snap keyring is for a different Snap ID', async () => { + const { provider } = setup({ + accounts: mockAccounts, + keyring: { snapId: 'npm:@metamask/other-snap' as SnapId }, + }); + + await expect(provider.resyncAccounts(mockAccounts)).rejects.toThrow( + 'No keyring matches the selector', + ); + }); + }); + describe('ensureReady', () => { it('delegates Snap platform readiness check to SnapAccountService:ensureReady', async () => { const { provider, mocks } = setup(); diff --git a/packages/multichain-account-service/src/providers/SnapAccountProvider.ts b/packages/multichain-account-service/src/providers/SnapAccountProvider.ts index 0d3ae60c4d..405df61ab5 100644 --- a/packages/multichain-account-service/src/providers/SnapAccountProvider.ts +++ b/packages/multichain-account-service/src/providers/SnapAccountProvider.ts @@ -1,7 +1,7 @@ import { assertIsBip44Account } from '@metamask/account-api'; import type { Bip44Account } from '@metamask/account-api'; import type { TraceCallback, TraceRequest } from '@metamask/controller-utils'; -import type { SnapKeyring } from '@metamask/eth-snap-keyring'; +import type { SnapKeyring as SnapKeyringV2 } from '@metamask/eth-snap-keyring/v2'; import { AccountCreationType, assertCreateAccountOptionIsSupported, @@ -13,8 +13,8 @@ import type { EntropySourceId, KeyringAccount, } from '@metamask/keyring-api'; +import { Keyring, KeyringType } from '@metamask/keyring-api/v2'; import type { KeyringMetadata } from '@metamask/keyring-controller'; -import { KeyringTypes } from '@metamask/keyring-controller'; import type { InternalAccount } from '@metamask/keyring-internal-api'; import { KeyringClient } from '@metamask/keyring-snap-client'; import type { Json, JsonRpcRequest, SnapId } from '@metamask/snaps-sdk'; @@ -33,9 +33,9 @@ import { BaseBip44AccountProvider } from './BaseBip44AccountProvider'; import { withTimeout } from './utils'; export type RestrictedSnapKeyring = { - createAccount: (options: Record) => Promise; - createAccounts: (options: CreateAccountOptions) => Promise; - removeAccount: (address: string) => Promise; + createAccount: SnapKeyringV2['createAccount']; + createAccounts: SnapKeyringV2['createAccounts']; + deleteAccount: SnapKeyringV2['deleteAccount']; }; export type SnapAccountProviderConfig = { @@ -72,6 +72,18 @@ export type SnapAccountProviderConfig = { }; }; +// TODO: Move this in `eth-snap-keyring/v2` package? +/** + * Checks if a given keyring is a Snap keyring (v2). + * + * @param keyring - The keyring to check. + * @returns `true` if the keyring is a Snap keyring (v2), `false` otherwise. + */ +function isSnapKeyring(keyring: Keyring): keyring is SnapKeyringV2 { + // Using `KeyringType.Snap` (used for v2). + return keyring.type === KeyringType.Snap; +} + export abstract class SnapAccountProvider extends BaseBip44AccountProvider { readonly snapId: SnapId; @@ -148,15 +160,15 @@ export abstract class SnapAccountProvider extends BaseBip44AccountProvider { } async #getRestrictedSnapKeyring(): Promise { - // NOTE: We're not supposed to make the keyring instance escape `withKeyring` but - // we have to use the `SnapKeyring` instance to be able to create Solana account + // NOTE: We're not supposed to make the keyring instance escape `withKeyringV2` but + // we have to use the `SnapKeyringV2` instance to be able to create Solana account // without triggering UI confirmation. // Also, creating account that way won't invalidate the Snap keyring state. The // account will get created and persisted properly with the Snap account creation // flow "asynchronously" (with `notify:accountCreated`). const { createAccount, createAccounts } = await this.#withSnapKeyring<{ - createAccount: SnapKeyring['createAccount']; - createAccounts: SnapKeyring['createAccounts']; + createAccount: SnapKeyringV2['createAccount']; + createAccounts: SnapKeyringV2['createAccounts']; }>(async ({ keyring }) => ({ createAccount: keyring.createAccount.bind(keyring), createAccounts: keyring.createAccounts.bind(keyring), @@ -165,17 +177,16 @@ export abstract class SnapAccountProvider extends BaseBip44AccountProvider { return { createAccount: async (options) => // We use the "unguarded" account creation here (see explanation above). - await createAccount(this.snapId, options, { + await createAccount(options, { displayAccountNameSuggestion: false, displayConfirmation: false, setSelectedAccount: false, }), - createAccounts: async (options) => - await createAccounts(this.snapId, options), - removeAccount: async (address: string) => + createAccounts: async (options) => await createAccounts(options), + deleteAccount: async (id: string) => // Though, when removing account, we can use the normal flow. await this.#withSnapKeyring(async ({ keyring }) => { - await keyring.removeAccount(address); + await keyring.deleteAccount(id); }), }; } @@ -265,7 +276,7 @@ export abstract class SnapAccountProvider extends BaseBip44AccountProvider { // We still need to remove the accounts from the Snap keyring since we're // about to create the same account again, which will use a new ID, but will // keep using the same address, and the Snap keyring does not allow this. - await keyring.removeAccount(account.address); + await keyring.deleteAccount(account.id); // The Snap has no account in its state for this one, we re-create it. await this.createAccounts({ type: AccountCreationType.Bip44DeriveIndex, @@ -290,15 +301,16 @@ export abstract class SnapAccountProvider extends BaseBip44AccountProvider { keyring, metadata, }: { - keyring: SnapKeyring; + keyring: SnapKeyringV2; metadata: KeyringMetadata; }) => Promise, ): Promise { - return this.withKeyring( - { type: KeyringTypes.snap }, - (args) => { - return operation(args); + return this.withKeyringV2( + { + filter: (keyring) => + isSnapKeyring(keyring) && keyring.snapId === this.snapId, }, + (args) => operation(args), ); } diff --git a/packages/multichain-account-service/src/providers/SolAccountProvider.test.ts b/packages/multichain-account-service/src/providers/SolAccountProvider.test.ts index 6419707e3b..b2c424af5f 100644 --- a/packages/multichain-account-service/src/providers/SolAccountProvider.test.ts +++ b/packages/multichain-account-service/src/providers/SolAccountProvider.test.ts @@ -1,11 +1,7 @@ import { isBip44Account } from '@metamask/account-api'; -import type { SnapKeyring } from '@metamask/eth-snap-keyring'; import { AccountCreationType } from '@metamask/keyring-api'; import type { KeyringMetadata } from '@metamask/keyring-controller'; -import type { - EthKeyring, - InternalAccount, -} from '@metamask/keyring-internal-api'; +import type { InternalAccount } from '@metamask/keyring-internal-api'; import { SnapControllerState } from '@metamask/snaps-controllers'; import deepmerge from 'deepmerge'; @@ -68,60 +64,56 @@ class MockSolanaKeyring { return Number(index); } - createAccount: SnapKeyring['createAccount'] = jest - .fn() - .mockImplementation((_, { derivationPath }) => { - if (derivationPath !== undefined) { - const index = this.#getIndexFromDerivationPath(derivationPath); - const found = this.accounts.find( - (account) => - isBip44Account(account) && - account.options.entropy.groupIndex === index, - ); - - if (found) { - return found; // Idempotent. - } + createAccount = jest.fn().mockImplementation(({ derivationPath }) => { + if (derivationPath !== undefined) { + const index = this.#getIndexFromDerivationPath(derivationPath); + const found = this.accounts.find( + (account) => + isBip44Account(account) && + account.options.entropy.groupIndex === index, + ); + + if (found) { + return found; // Idempotent. + } + } + + const account = MockAccountBuilder.from(MOCK_SOL_ACCOUNT_1) + .withUuid() + .withAddressSuffix(`${this.accounts.length}`) + .withGroupIndex(this.accounts.length) + .get(); + this.accounts.push(account); + + return account; + }); + + createAccounts = jest.fn().mockImplementation((options) => { + const groupIndices = + options.type === 'bip44:derive-index' + ? [options.groupIndex] + : toGroupIndexRangeArray(options.range); + + return groupIndices.map((groupIndex) => { + const found = this.accounts.find( + (account) => + isBip44Account(account) && + account.options.entropy.groupIndex === groupIndex, + ); + + if (found) { + return found; // Idempotent. } const account = MockAccountBuilder.from(MOCK_SOL_ACCOUNT_1) .withUuid() - .withAddressSuffix(`${this.accounts.length}`) - .withGroupIndex(this.accounts.length) + .withAddressSuffix(`${groupIndex}`) + .withGroupIndex(groupIndex) .get(); this.accounts.push(account); - return account; }); - - createAccounts: SnapKeyring['createAccounts'] = jest - .fn() - .mockImplementation((_, options) => { - const groupIndices = - options.type === 'bip44:derive-index' - ? [options.groupIndex] - : toGroupIndexRangeArray(options.range); - - return groupIndices.map((groupIndex) => { - const found = this.accounts.find( - (account) => - isBip44Account(account) && - account.options.entropy.groupIndex === groupIndex, - ); - - if (found) { - return found; // Idempotent. - } - - const account = MockAccountBuilder.from(MOCK_SOL_ACCOUNT_1) - .withUuid() - .withAddressSuffix(`${groupIndex}`) - .withGroupIndex(groupIndex) - .get(); - this.accounts.push(account); - return account; - }); - }); + }); } class MockSolAccountProvider extends SolAccountProvider { @@ -201,12 +193,10 @@ function setup({ ); messenger.registerActionHandler( - 'KeyringController:withKeyring', + 'KeyringController:withKeyringV2', async (_, operation) => operation({ - // We type-cast here, since `withKeyring` defaults to `EthKeyring` and the - // Snap keyring doesn't really implement this interface (this is expected). - keyring: keyring as unknown as EthKeyring, + keyring, metadata: keyring.metadata, }), ); @@ -228,8 +218,8 @@ function setup({ mocks: { handleRequest: mockHandleRequest, keyring: { - createAccount: keyring.createAccount as jest.Mock, - createAccounts: keyring.createAccounts as jest.Mock, + createAccount: keyring.createAccount, + createAccounts: keyring.createAccounts, }, trace: mockTrace, }, @@ -465,14 +455,11 @@ describe('SolAccountProvider', () => { }); expect(newAccounts).toHaveLength(1); // Batch endpoint must be called, NOT the singular one. - expect(mocks.keyring.createAccounts).toHaveBeenCalledWith( - SolAccountProvider.SOLANA_SNAP_ID, - { - type: AccountCreationType.Bip44DeriveIndex, - entropySource: MOCK_HD_KEYRING_1.metadata.id, - groupIndex: newGroupIndex, - }, - ); + expect(mocks.keyring.createAccounts).toHaveBeenCalledWith({ + type: AccountCreationType.Bip44DeriveIndex, + entropySource: MOCK_HD_KEYRING_1.metadata.id, + groupIndex: newGroupIndex, + }); expect(mocks.keyring.createAccount).not.toHaveBeenCalled(); }); diff --git a/packages/multichain-account-service/src/providers/TrxAccountProvider.test.ts b/packages/multichain-account-service/src/providers/TrxAccountProvider.test.ts index c755f8e5ce..d8b617ad1e 100644 --- a/packages/multichain-account-service/src/providers/TrxAccountProvider.test.ts +++ b/packages/multichain-account-service/src/providers/TrxAccountProvider.test.ts @@ -1,11 +1,7 @@ import { isBip44Account } from '@metamask/account-api'; -import type { SnapKeyring } from '@metamask/eth-snap-keyring'; import { AccountCreationType } from '@metamask/keyring-api'; import type { KeyringMetadata } from '@metamask/keyring-controller'; -import type { - EthKeyring, - InternalAccount, -} from '@metamask/keyring-internal-api'; +import type { InternalAccount } from '@metamask/keyring-internal-api'; import { SnapControllerState } from '@metamask/snaps-controllers'; import deepmerge from 'deepmerge'; @@ -52,13 +48,39 @@ class MockTronKeyring { this.accounts = accounts; } - createAccount: SnapKeyring['createAccount'] = jest - .fn() - .mockImplementation((_, { index }) => { - // Use the provided index or fallback to accounts length - const groupIndex = index ?? this.accounts.length; + createAccount = jest.fn().mockImplementation(({ index }) => { + // Use the provided index or fallback to accounts length + const groupIndex = index ?? this.accounts.length; - // Check if an account already exists for this group index (idempotent behavior) + // Check if an account already exists for this group index (idempotent behavior) + const found = this.accounts.find( + (account) => + isBip44Account(account) && + account.options.entropy.groupIndex === groupIndex, + ); + + if (found) { + return found; // Idempotent. + } + + // Create new account with the correct group index + const account = MockAccountBuilder.from(MOCK_TRX_ACCOUNT_1) + .withUuid() + .withAddressSuffix(`${this.accounts.length}`) + .withGroupIndex(groupIndex) + .get(); + this.accounts.push(account); + + return account; + }); + + createAccounts = jest.fn().mockImplementation((options) => { + const groupIndices = + options.type === 'bip44:derive-index' + ? [options.groupIndex] + : toGroupIndexRangeArray(options.range); + + return groupIndices.map((groupIndex) => { const found = this.accounts.find( (account) => isBip44Account(account) && @@ -69,45 +91,15 @@ class MockTronKeyring { return found; // Idempotent. } - // Create new account with the correct group index const account = MockAccountBuilder.from(MOCK_TRX_ACCOUNT_1) .withUuid() - .withAddressSuffix(`${this.accounts.length}`) + .withAddressSuffix(`${groupIndex}`) .withGroupIndex(groupIndex) .get(); this.accounts.push(account); - return account; }); - - createAccounts: SnapKeyring['createAccounts'] = jest - .fn() - .mockImplementation((_, options) => { - const groupIndices = - options.type === 'bip44:derive-index' - ? [options.groupIndex] - : toGroupIndexRangeArray(options.range); - - return groupIndices.map((groupIndex) => { - const found = this.accounts.find( - (account) => - isBip44Account(account) && - account.options.entropy.groupIndex === groupIndex, - ); - - if (found) { - return found; // Idempotent. - } - - const account = MockAccountBuilder.from(MOCK_TRX_ACCOUNT_1) - .withUuid() - .withAddressSuffix(`${groupIndex}`) - .withGroupIndex(groupIndex) - .get(); - this.accounts.push(account); - return account; - }); - }); + }); // Add discoverAccounts method to match the provider's usage discoverAccounts = jest.fn().mockResolvedValue([]); @@ -192,12 +184,10 @@ function setup({ ); messenger.registerActionHandler( - 'KeyringController:withKeyring', + 'KeyringController:withKeyringV2', async (_, operation) => operation({ - // We type-cast here, since `withKeyring` defaults to `EthKeyring` and the - // Snap keyring doesn't really implement this interface (this is expected). - keyring: keyring as unknown as EthKeyring, + keyring, metadata: keyring.metadata, }), ); @@ -223,8 +213,8 @@ function setup({ mocks: { handleRequest: mockHandleRequest, keyring: { - createAccount: keyring.createAccount as jest.Mock, - createAccounts: keyring.createAccounts as jest.Mock, + createAccount: keyring.createAccount, + createAccounts: keyring.createAccounts, discoverAccounts: keyring.discoverAccounts, }, trace: mockTrace, @@ -466,14 +456,11 @@ describe('TrxAccountProvider', () => { }); expect(newAccounts).toHaveLength(1); // Batch endpoint must be called, NOT the singular one. - expect(mocks.keyring.createAccounts).toHaveBeenCalledWith( - TrxAccountProvider.TRX_SNAP_ID, - { - type: AccountCreationType.Bip44DeriveIndex, - entropySource: MOCK_HD_KEYRING_1.metadata.id, - groupIndex: newGroupIndex, - }, - ); + expect(mocks.keyring.createAccounts).toHaveBeenCalledWith({ + type: AccountCreationType.Bip44DeriveIndex, + entropySource: MOCK_HD_KEYRING_1.metadata.id, + groupIndex: newGroupIndex, + }); expect(mocks.keyring.createAccount).not.toHaveBeenCalled(); }); diff --git a/packages/snap-account-service/CHANGELOG.md b/packages/snap-account-service/CHANGELOG.md index 4f7d12cebc..a38ea7d267 100644 --- a/packages/snap-account-service/CHANGELOG.md +++ b/packages/snap-account-service/CHANGELOG.md @@ -12,11 +12,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added - Add `SnapAccountService` ([#8414](https://github.com/MetaMask/core/pull/8414)) -- Add `SnapPlatformWatcher` and `SnapAccountService.ensureReady` ([#8715](https://github.com/MetaMask/core/pull/8715)), ([#8725](https://github.com/MetaMask/core/pull/8725)) +- Add `SnapPlatformWatcher` and `SnapAccountService.ensureReady` ([#8715](https://github.com/MetaMask/core/pull/8715)), ([#8725](https://github.com/MetaMask/core/pull/8725)), ([#8732](https://github.com/MetaMask/core/pull/8732)) - Waits for the Snap platform to be ready and for a Snap keyring to appear in `KeyringController` state before allowing Snap account operations. - Callers must ensure `init()` has run and the Snap is currently installed, enabled, non-blocked, and declares `endowment:keyring`. - `SnapAccountService.ensureReady` now awaits the watcher, so it only resolves once both conditions hold (or rejects if the Snap keyring does not appear within the configured timeout). - `SnapAccountService.ensureReady` now throws `Unknown snap: ""` when called with a Snap ID that isn't tracked as an account-management Snap. + - `SnapAccountService.ensureReady` now automatically creates the Snap keyring (v2) for a given Snap ID if it was not available. - Add `config` option to `SnapAccountService` constructor with a `snapPlatformWatcher` field exposing `ensureOnboardingComplete` and `snapKeyringWaitTimeoutMs` ([#8715](https://github.com/MetaMask/core/pull/8715)) - Export `SnapAccountServiceConfig` and `SnapPlatformWatcherConfig` types. - Add `@metamask/keyring-controller` dependency ([#8715](https://github.com/MetaMask/core/pull/8715)) @@ -34,6 +35,18 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - This logic used to live on the clients. - The service messenger now requires the `KeyringController:unlock`, `AccountTreeController:selectedAccountGroupChange`, `AccountTreeController:accountGroup{Created,Updated,Removed}` events. - The service messenger now requires the `AccountTreeController:getSelectedAccountGroup` and `AccountTreeController:getAccountGroupObject` actions. +- Add `migrate` ([#8732](https://github.com/MetaMask/core/pull/8732)) + - The migration is guaranteed to be run when using `ensureReady`. + - If the migration is not successful, it will get retried everytime we need to interact with any Snap keyring (v2) instances. + - It is conccurent-free and can safely be called by multiple execution flows. + - Once the migration has ran, the legacy Snap keyring will be emptied, thus, consumers are expected to use the new per-Snap keyring (v2) instances instead. + - Selected-account forwarding now targets v2 Snap keyrings. + - The service messenger now requires the `KeyringController:withKeyringV2Unsafe`. + +### Removed + +- **BREAKING:** Removed `getLegacySnapKeyring` ([#8732](https://github.com/MetaMask/core/pull/8732)) + - The legacy Snap keyring should not be used anymore after the migration has completed. ### Changed diff --git a/packages/snap-account-service/package.json b/packages/snap-account-service/package.json index 8ecbc9933f..2053ff7249 100644 --- a/packages/snap-account-service/package.json +++ b/packages/snap-account-service/package.json @@ -56,7 +56,9 @@ "@metamask/account-api": "^1.0.4", "@metamask/account-tree-controller": "^7.4.0", "@metamask/eth-snap-keyring": "^22.0.1", + "@metamask/keyring-api": "^23.1.0", "@metamask/keyring-controller": "^25.5.0", + "@metamask/keyring-snap-sdk": "^9.0.1", "@metamask/messenger": "^1.2.0", "@metamask/snaps-controllers": "^19.0.0", "@metamask/snaps-sdk": "^11.0.0", diff --git a/packages/snap-account-service/src/SnapAccountService-method-action-types.ts b/packages/snap-account-service/src/SnapAccountService-method-action-types.ts index 15af5c939f..1c08a74df2 100644 --- a/packages/snap-account-service/src/SnapAccountService-method-action-types.ts +++ b/packages/snap-account-service/src/SnapAccountService-method-action-types.ts @@ -20,7 +20,11 @@ export type SnapAccountServiceGetSnapsAction = { /** * Ensures everything is ready to use Snap accounts for the given Snap. * 1. Validates that `snapId` is a tracked account-management Snap. - * 2. Waits for the Snap platform to be fully started. + * 2. Runs the legacy -> v2 Snap keyring migration (cached — no-op if + * already done). + * 3. Atomically creates the v2 keyring for this Snap if it doesn't exist + * yet. + * 4. Waits for the Snap platform to be fully started. * * Safe to call concurrently — each step is idempotent or mutex-protected. * @@ -33,14 +37,15 @@ export type SnapAccountServiceEnsureReadyAction = { }; /** - * Atomically gets-or-creates the legacy (v1) Snap keyring — the keyring - * associated with {@link KeyringTypes.snap}. + * Migrate the legacy Snap keyring to the new (per-snap) Snap keyring v2. + * Safe to call concurrently — the migration runs only once; all callers + * await the same promise. * - * @returns The existing or newly-created Snap keyring instance. + * @returns A promise that resolves when the migration is complete. */ -export type SnapAccountServiceGetLegacySnapKeyringAction = { - type: `SnapAccountService:getLegacySnapKeyring`; - handler: SnapAccountService['getLegacySnapKeyring']; +export type SnapAccountServiceMigrateAction = { + type: `SnapAccountService:migrate`; + handler: SnapAccountService['migrate']; }; /** @@ -61,5 +66,5 @@ export type SnapAccountServiceHandleKeyringSnapMessageAction = { export type SnapAccountServiceMethodActions = | SnapAccountServiceGetSnapsAction | SnapAccountServiceEnsureReadyAction - | SnapAccountServiceGetLegacySnapKeyringAction + | SnapAccountServiceMigrateAction | SnapAccountServiceHandleKeyringSnapMessageAction; diff --git a/packages/snap-account-service/src/SnapAccountService.test.ts b/packages/snap-account-service/src/SnapAccountService.test.ts index e1569d63a0..89ee9bb318 100644 --- a/packages/snap-account-service/src/SnapAccountService.test.ts +++ b/packages/snap-account-service/src/SnapAccountService.test.ts @@ -1,13 +1,22 @@ import type { AccountGroupId } from '@metamask/account-api'; -import type { SnapKeyring, SnapMessage } from '@metamask/eth-snap-keyring'; +import { SNAP_KEYRING_TYPE } from '@metamask/eth-snap-keyring'; +import type { SnapMessage } from '@metamask/eth-snap-keyring'; +import type { SnapKeyring as SnapKeyringV2 } from '@metamask/eth-snap-keyring/v2'; +import { KeyringEvent } from '@metamask/keyring-api'; +import { KeyringType } from '@metamask/keyring-api/v2'; import { + KeyringControllerError, + KeyringControllerErrorMessage, KeyringControllerState, KeyringTypes, } from '@metamask/keyring-controller'; import type { KeyringEntry, + KeyringMetadata, + KeyringSelectorV2, RestrictedController, } from '@metamask/keyring-controller'; +import { SnapManageAccountsMethod } from '@metamask/keyring-snap-sdk'; import { Messenger, MOCK_ANY_NAMESPACE } from '@metamask/messenger'; import type { MockAnyNamespace, @@ -43,15 +52,6 @@ type MockTruncatedSnap = Pick< /** Mock account group type for tests. */ type MockAccountGroup = Pick; -/** Mock Snap keyring type for tests. */ -type MockSnapKeyring = { - type: KeyringTypes.snap; - handleKeyringSnapMessage?: jest.MockedFunction< - SnapKeyring['handleKeyringSnapMessage'] - >; - setSelectedAccounts?: jest.MockedFunction; -}; - type Mocks = { // eslint-disable-next-line @typescript-eslint/naming-convention SnapController: { @@ -62,6 +62,8 @@ type Mocks = { KeyringController: { getState: jest.MockedFunction<() => { keyrings: { type: string }[] }>; withController: jest.Mock; + withKeyringV2: jest.Mock; + withKeyringV2Unsafe: jest.Mock; }; // eslint-disable-next-line @typescript-eslint/naming-convention AccountTreeController: { @@ -103,6 +105,8 @@ function getMessenger( 'SnapController:getRunnableSnaps', 'KeyringController:getState', 'KeyringController:withController', + 'KeyringController:withKeyringV2', + 'KeyringController:withKeyringV2Unsafe', 'AccountTreeController:getAccountGroupObject', 'AccountTreeController:getSelectedAccountGroup', ], @@ -202,10 +206,10 @@ function publishUnlock(rootMessenger: RootMessenger): void { * Builds a minimal `TruncatedSnap` for tests. * * @param id - The Snap ID. - * @param hasKeyring - Whether the Snap declares the `endowment:keyring` initial permission. + * @param hasKeyring - Whether the Snap declares the `endowment:keyring` initial permission (default: `true`). * @returns A minimal `TruncatedSnap`. */ -function buildSnap(id: string, hasKeyring: boolean): TruncatedSnap { +function buildSnap(id: string, hasKeyring = true): TruncatedSnap { return { id: id as SnapId, initialPermissions: hasKeyring ? { 'endowment:keyring': {} } : {}, @@ -316,75 +320,141 @@ function mockWithController( } /** - * Configures `mocks.KeyringController.withController` to expose a single - * legacy Snap keyring with the provided mocked methods. + * Configures `mocks.KeyringController.withKeyringV2` so that the operation + * receives a Snap keyring v2 matching the given selector, or throws + * `KeyringNotFound` when none matches. * * @param mocks - The mocks object from {@link setup}. - * @param keyring - The mocked Snap keyring methods. - * @param keyring.handleKeyringSnapMessage - The mocked implementation. - * @param keyring.setSelectedAccounts - The mocked implementation. + * @param keyrings - The available v2 Snap keyrings, keyed by snap ID. */ -function mockLegacySnapKeyring( +function mockWithKeyringV2( mocks: Mocks, - { - handleKeyringSnapMessage, - setSelectedAccounts, - }: { - handleKeyringSnapMessage?: jest.MockedFunction< - SnapKeyring['handleKeyringSnapMessage'] - >; - setSelectedAccounts?: jest.MockedFunction< - SnapKeyring['setSelectedAccounts'] - >; - }, + keyrings: Record>, ): void { - const snapKeyring: MockSnapKeyring = { - type: KeyringTypes.snap, - handleKeyringSnapMessage, - setSelectedAccounts, - }; - mocks.KeyringController.withController.mockImplementation(async (operation) => - operation({ - get keyrings() { - return Object.freeze([ + mocks.KeyringController.withKeyringV2.mockImplementation( + async ( + selector: KeyringSelectorV2, + operation: (args: { + keyring: SnapKeyringV2; + metadata: KeyringMetadata; + }) => Promise, + ) => { + const entry = Object.entries(keyrings).find(([snapId, snapKeyring]) => + // The selector's filter expects a v2 keyring object; we synthesise + // a minimal shape (`type` + `snapId`) so the production filter + // function can identify it. + selector.filter?.( { - keyring: snapKeyring as KeyringEntry['keyring'], - metadata: { id: 'id-snap', name: KeyringTypes.snap }, - }, - ]); - }, - addNewKeyring: jest.fn(), - removeKeyring: jest.fn(), - }), + type: KeyringType.Snap, + snapId, + ...snapKeyring, + } as unknown, + { id: `id-${snapId}`, name: 'snap' } as KeyringMetadata, + ), + ); + if (!entry) { + throw new KeyringControllerError( + KeyringControllerErrorMessage.KeyringNotFound, + ); + } + const [snapId, snapKeyring] = entry; + return operation({ + keyring: { + type: KeyringType.Snap, + snapId, + ...snapKeyring, + } as unknown as SnapKeyringV2, + metadata: { id: `id-${snapId}`, name: 'snap' } as KeyringMetadata, + }); + }, + ); +} + +/** + * Configures `mocks.KeyringController.withKeyringV2Unsafe` so that the + * operation receives a Snap keyring v2 matching the given selector, or + * throws `KeyringNotFound` when none matches. + * + * @param mocks - The mocks object from {@link setup}. + * @param keyrings - The available v2 Snap keyrings, keyed by snap ID. + */ +function mockWithKeyringV2Unsafe( + mocks: Mocks, + keyrings: Record< + string, + { + setSelectedAccounts?: jest.Mock; + hasAccount?: (id: string) => boolean; + } + >, +): void { + mocks.KeyringController.withKeyringV2Unsafe.mockImplementation( + async ( + selector: KeyringSelectorV2, + operation: (args: { + keyring: SnapKeyringV2; + metadata: KeyringMetadata; + }) => Promise, + ) => { + const entry = Object.entries(keyrings).find(([snapId, kr]) => + selector.filter?.( + { + type: KeyringType.Snap, + snapId, + ...kr, + } as unknown, + { id: `id-${snapId}`, name: 'snap' } as KeyringMetadata, + ), + ); + if (!entry) { + throw new KeyringControllerError( + KeyringControllerErrorMessage.KeyringNotFound, + ); + } + const [snapId, kr] = entry; + return operation({ + keyring: { + type: KeyringType.Snap, + snapId, + hasAccount: () => true, + ...kr, + } as unknown as SnapKeyringV2, + metadata: { id: `id-${snapId}`, name: 'snap' } as KeyringMetadata, + }); + }, ); } /** - * Constructs the service under test with sensible defaults. + * Constructs the service under test with sensible defaults. Calls + * `service.init()` before returning unless `init: false` is passed. * * @param args - The arguments to this function. * @param args.snapIsReady - Initial value of `SnapController.isReady`. * @param args.keyrings - Initial keyrings returned by `KeyringController:getState`. * @param args.runnableSnaps - Snaps returned by `SnapController:getRunnableSnaps`. * @param args.config - Optional service config. + * @param args.init - Whether to call `service.init()` before returning (default: `true`). * @returns The new service, root messenger, service messenger, and mocks. */ -function setup({ +async function setup({ snapIsReady = true, keyrings = [{ type: KeyringTypes.snap }], runnableSnaps = [], config, + init = true, }: { snapIsReady?: boolean; keyrings?: { type: string }[]; runnableSnaps?: TruncatedSnap[]; config?: SnapAccountServiceOptions['config']; -} = {}): { + init?: boolean; +} = {}): Promise<{ service: SnapAccountService; rootMessenger: RootMessenger; messenger: SnapAccountServiceMessenger; mocks: Mocks; -} { +}> { const rootMessenger = getRootMessenger(); const messenger = getMessenger(rootMessenger); @@ -398,6 +468,8 @@ function setup({ KeyringController: { getState: jest.fn().mockReturnValue({ keyrings }), withController: jest.fn(), + withKeyringV2: jest.fn(), + withKeyringV2Unsafe: jest.fn(), }, AccountTreeController: { getAccountGroupObject: jest.fn().mockReturnValue(undefined), @@ -421,6 +493,14 @@ function setup({ 'KeyringController:withController', mocks.KeyringController.withController, ); + rootMessenger.registerActionHandler( + 'KeyringController:withKeyringV2', + mocks.KeyringController.withKeyringV2, + ); + rootMessenger.registerActionHandler( + 'KeyringController:withKeyringV2Unsafe', + mocks.KeyringController.withKeyringV2Unsafe, + ); rootMessenger.registerActionHandler( 'AccountTreeController:getAccountGroupObject', mocks.AccountTreeController.getAccountGroupObject, @@ -432,15 +512,20 @@ function setup({ const service = new SnapAccountService({ messenger, config }); + if (init) { + await service.init(); + } + return { service, rootMessenger, messenger, mocks }; } const MOCK_SNAP_ID = 'npm:@metamask/mock-snap' as SnapId; +const MOCK_OTHER_SNAP_ID = 'npm:@metamask/other-snap' as SnapId; describe('SnapAccountService', () => { describe('init', () => { it('resolves without throwing', async () => { - const { service } = setup(); + const { service } = await setup({ init: false }); expect(await service.init()).toBeUndefined(); }); @@ -448,21 +533,146 @@ describe('SnapAccountService', () => { describe('getSnaps', () => { it('exposes tracked Snaps seeded by init', async () => { - const { service } = setup({ - runnableSnaps: [buildSnap(MOCK_SNAP_ID, true)], + const { service } = await setup({ + runnableSnaps: [buildSnap(MOCK_SNAP_ID)], }); - await service.init(); - expect(service.getSnaps()).toStrictEqual([MOCK_SNAP_ID]); }); }); + describe('migrate', () => { + it('runs the migration only once when called concurrently', async () => { + const { service, mocks } = await setup(); + mocks.KeyringController.withController.mockResolvedValue(undefined); + + await Promise.all([service.migrate(), service.migrate()]); + + expect(mocks.KeyringController.withController).toHaveBeenCalledTimes(1); + }); + + it('is a no-op when no legacy Snap keyring is present', async () => { + const addNewKeyring = jest.fn().mockResolvedValue(undefined); + const removeKeyring = jest.fn().mockResolvedValue(undefined); + const { service, mocks } = await setup(); + mocks.KeyringController.withController.mockImplementation( + async (operation) => + operation({ keyrings: [], addNewKeyring, removeKeyring }), + ); + + await service.migrate(); + + expect(addNewKeyring).not.toHaveBeenCalled(); + expect(removeKeyring).not.toHaveBeenCalled(); + }); + + it('migrates accounts from the legacy Snap keyring to per-Snap v2 keyrings and removes the legacy entry', async () => { + const addNewKeyring = jest.fn().mockResolvedValue(undefined); + const removeKeyring = jest.fn().mockResolvedValue(undefined); + const legacyKeyringId = 'legacy-keyring-id'; + const account1 = { + id: 'account-1', + address: '0x1', + metadata: { snap: { id: MOCK_SNAP_ID } }, + }; + const account2 = { + id: 'account-2', + address: '0x2', + metadata: { snap: { id: MOCK_SNAP_ID } }, + }; + const account3 = { + id: 'account-3', + address: '0x3', + metadata: { snap: { id: MOCK_OTHER_SNAP_ID } }, + }; + const orphanAccount = { + id: 'orphan', + address: '0x4', + metadata: {}, + }; + const legacyKeyring = { + type: SNAP_KEYRING_TYPE, + listAccounts: jest + .fn() + .mockReturnValue([account1, account2, account3, orphanAccount]), + }; + const { service, mocks } = await setup(); + mocks.KeyringController.withController.mockImplementation( + async (operation) => + operation({ + keyrings: [ + { keyring: legacyKeyring, metadata: { id: legacyKeyringId } }, + ], + addNewKeyring, + removeKeyring, + }), + ); + + await service.migrate(); + + expect(addNewKeyring).toHaveBeenCalledTimes(2); + expect(addNewKeyring).toHaveBeenCalledWith(KeyringType.Snap, { + snapId: MOCK_SNAP_ID, + accounts: { + [account1.id]: { id: account1.id, address: account1.address }, + [account2.id]: { id: account2.id, address: account2.address }, + }, + }); + expect(addNewKeyring).toHaveBeenCalledWith(KeyringType.Snap, { + snapId: MOCK_OTHER_SNAP_ID, + accounts: { + [account3.id]: { id: account3.id, address: account3.address }, + }, + }); + expect(removeKeyring).toHaveBeenCalledWith(legacyKeyringId); + }); + + it('does not re-run after a successful migration', async () => { + const { service, mocks } = await setup(); + mocks.KeyringController.withController.mockResolvedValue(undefined); + + await service.migrate(); + await service.migrate(); + + expect(mocks.KeyringController.withController).toHaveBeenCalledTimes(1); + }); + + it('retries on a subsequent call after a failed migration', async () => { + const { service, mocks } = await setup(); + const error = new Error('migration boom'); + mocks.KeyringController.withController + .mockRejectedValueOnce(error) + .mockResolvedValueOnce(undefined); + + await expect(service.migrate()).rejects.toThrow(error); + expect(await service.migrate()).toBeUndefined(); + + expect(mocks.KeyringController.withController).toHaveBeenCalledTimes(2); + }); + + it('shares the rejection across concurrent callers but allows a later retry', async () => { + const { service, mocks } = await setup(); + const error = new Error('migration boom'); + mocks.KeyringController.withController + .mockRejectedValueOnce(error) + .mockResolvedValueOnce(undefined); + + const [first, second] = await Promise.allSettled([ + service.migrate(), + service.migrate(), + ]); + expect(first).toStrictEqual({ status: 'rejected', reason: error }); + expect(second).toStrictEqual({ status: 'rejected', reason: error }); + expect(mocks.KeyringController.withController).toHaveBeenCalledTimes(1); + + expect(await service.migrate()).toBeUndefined(); + expect(mocks.KeyringController.withController).toHaveBeenCalledTimes(2); + }); + }); + describe('ensureReady', () => { it('throws when the Snap is not tracked', async () => { - const { service } = setup(); - - await service.init(); + const { service } = await setup(); await expect(service.ensureReady(MOCK_SNAP_ID)).rejects.toThrow( `Unknown snap: "${MOCK_SNAP_ID}"`, @@ -470,8 +680,9 @@ describe('SnapAccountService', () => { }); it('throws before init even for runnable Snaps', async () => { - const { service } = setup({ - runnableSnaps: [buildSnap(MOCK_SNAP_ID, true)], + const { service } = await setup({ + runnableSnaps: [buildSnap(MOCK_SNAP_ID)], + init: false, }); await expect(service.ensureReady(MOCK_SNAP_ID)).rejects.toThrow( @@ -480,23 +691,83 @@ describe('SnapAccountService', () => { }); it('resolves when platform is already ready', async () => { - const { service } = setup({ - runnableSnaps: [buildSnap(MOCK_SNAP_ID, true)], + const { service } = await setup({ + runnableSnaps: [buildSnap(MOCK_SNAP_ID)], }); - await service.init(); - expect(await service.ensureReady(MOCK_SNAP_ID)).toBeUndefined(); }); + it('runs the migration before checking platform readiness', async () => { + const { service, mocks } = await setup({ + runnableSnaps: [buildSnap(MOCK_SNAP_ID)], + }); + mocks.KeyringController.withController.mockResolvedValue(undefined); + + // `migrate` is invoked once + `#createKeyringForSnap` is invoked once + // (the cached migrate call is a no-op on subsequent calls). + await service.ensureReady(MOCK_SNAP_ID); + + expect(mocks.KeyringController.withController).toHaveBeenCalledTimes(2); + }); + + it('creates a v2 keyring for the Snap when one does not exist yet', async () => { + const addNewKeyring = jest.fn().mockResolvedValue(undefined); + const removeKeyring = jest.fn().mockResolvedValue(undefined); + const { service, mocks } = await setup({ + runnableSnaps: [buildSnap(MOCK_SNAP_ID)], + }); + mocks.KeyringController.withController.mockImplementation( + async (operation) => + operation({ keyrings: [], addNewKeyring, removeKeyring }), + ); + + await service.ensureReady(MOCK_SNAP_ID); + + expect(addNewKeyring).toHaveBeenCalledWith(KeyringType.Snap, { + snapId: MOCK_SNAP_ID, + accounts: {}, + }); + }); + + it('does not create a v2 keyring when one already exists for the Snap', async () => { + const addNewKeyring = jest.fn().mockResolvedValue(undefined); + const removeKeyring = jest.fn().mockResolvedValue(undefined); + const { service, mocks } = await setup({ + runnableSnaps: [buildSnap(MOCK_SNAP_ID)], + }); + // First call: migration (no legacy snap keyring → early return). + // Second call: ensureReady keyring check (snap keyring already exists). + mocks.KeyringController.withController + .mockImplementationOnce(async (operation) => + operation({ keyrings: [], addNewKeyring, removeKeyring }), + ) + .mockImplementationOnce(async (operation) => + operation({ + keyrings: [ + { + keyringV2: { + type: KeyringType.Snap, + snapId: MOCK_SNAP_ID, + }, + }, + ], + addNewKeyring, + removeKeyring, + }), + ); + + await service.ensureReady(MOCK_SNAP_ID); + + expect(addNewKeyring).not.toHaveBeenCalled(); + }); + it('waits for the Snap platform to become ready', async () => { - const { service, rootMessenger } = setup({ + const { service, rootMessenger } = await setup({ snapIsReady: false, - runnableSnaps: [buildSnap(MOCK_SNAP_ID, true)], + runnableSnaps: [buildSnap(MOCK_SNAP_ID)], }); - await service.init(); - let resolved = false; const ensurePromise = service.ensureReady(MOCK_SNAP_ID).then(() => { resolved = true; @@ -512,22 +783,20 @@ describe('SnapAccountService', () => { }); it('waits for the Snap keyring to appear via KeyringController:stateChange', async () => { - const { service, rootMessenger } = setup({ + const { service, rootMessenger } = await setup({ keyrings: [], - runnableSnaps: [buildSnap(MOCK_SNAP_ID, true)], + runnableSnaps: [buildSnap(MOCK_SNAP_ID)], }); - await service.init(); - let resolved = false; const ensurePromise = service.ensureReady(MOCK_SNAP_ID).then(() => { resolved = true; return undefined; }); - // Flush microtasks so #waitForSnapKeyring subscribes. - await Promise.resolve(); - await Promise.resolve(); + // Flush microtasks so migration completes and #waitForSnapKeyring + // subscribes. + await flushMicrotasks(); expect(resolved).toBe(false); @@ -538,16 +807,14 @@ describe('SnapAccountService', () => { }); it('rejects if the Snap keyring does not appear within snapKeyringWaitTimeoutMs', async () => { - const { service } = setup({ + const { service } = await setup({ keyrings: [], - runnableSnaps: [buildSnap(MOCK_SNAP_ID, true)], + runnableSnaps: [buildSnap(MOCK_SNAP_ID)], config: { snapPlatformWatcher: { snapKeyringWaitTimeoutMs: 1_000 }, }, }); - await service.init(); - jest.useFakeTimers(); const ensurePromise = service.ensureReady(MOCK_SNAP_ID); // Attach rejection handler before advancing timers to avoid unhandled rejection. @@ -571,20 +838,18 @@ describe('SnapAccountService', () => { }), ); - const { service } = setup({ - runnableSnaps: [buildSnap(MOCK_SNAP_ID, true)], + const { service } = await setup({ + runnableSnaps: [buildSnap(MOCK_SNAP_ID)], config: { snapPlatformWatcher: { ensureOnboardingComplete } }, }); - await service.init(); - let resolved = false; const ensurePromise = service.ensureReady(MOCK_SNAP_ID).then(() => { resolved = true; return undefined; }); - await Promise.resolve(); + await flushMicrotasks(); expect(ensureOnboardingComplete).toHaveBeenCalledTimes(1); expect(resolved).toBe(false); @@ -595,83 +860,133 @@ describe('SnapAccountService', () => { }); }); - describe('getLegacySnapKeyring', () => { - it('returns the existing Snap keyring when one is already present', async () => { - const { service, mocks } = setup(); - const existing = buildKeyringEntry(KeyringTypes.snap); - const { addNewKeyring } = mockWithController(mocks, [ - buildKeyringEntry(KeyringTypes.hd), - existing, - ]); - - const result = await service.getLegacySnapKeyring(); - - expect(result).toBe(existing.keyring as unknown as SnapKeyring); - expect(addNewKeyring).not.toHaveBeenCalled(); - }); - - it('creates a new Snap keyring when none exists', async () => { - const { service, mocks } = setup(); - const { addNewKeyring } = mockWithController(mocks, [ - buildKeyringEntry(KeyringTypes.hd), - ]); - - const result = await service.getLegacySnapKeyring(); - - expect(addNewKeyring).toHaveBeenCalledWith(KeyringTypes.snap); - expect(result.type).toBe(KeyringTypes.snap); - }); - - it('propagates errors thrown by withController', async () => { - const { service, mocks } = setup(); - mocks.KeyringController.withController.mockImplementation(async () => { - throw new Error('boom'); - }); - - await expect(service.getLegacySnapKeyring()).rejects.toThrow('boom'); - }); - }); - describe('handleKeyringSnapMessage', () => { const MOCK_MESSAGE = { - method: 'keyring_listAccounts', + method: KeyringEvent.AccountUpdated, params: {}, } as unknown as SnapMessage; + const MOCK_ACCOUNT_CREATED_MESSAGE = { + method: KeyringEvent.AccountCreated, + params: {}, + } as unknown as SnapMessage; + const MOCK_GROUP_ID = 'keyring:01JABC/group-1' as AccountGroupId; + const MOCK_ACCOUNTS = [ + '00000000-0000-0000-0000-000000000001', + '00000000-0000-0000-0000-000000000002', + ]; - it('forwards the call to the legacy Snap keyring and returns its result', async () => { - const { service, mocks } = setup(); + it('forwards the message to the matching v2 Snap keyring and returns its result', async () => { + const { service, mocks } = await setup(); const handleKeyringSnapMessage = jest .fn() .mockResolvedValue({ ok: true }); - mockLegacySnapKeyring(mocks, { handleKeyringSnapMessage }); + mockWithKeyringV2(mocks, { + [MOCK_SNAP_ID]: { handleKeyringSnapMessage }, + }); const result = await service.handleKeyringSnapMessage( MOCK_SNAP_ID, MOCK_MESSAGE, ); - expect(handleKeyringSnapMessage).toHaveBeenCalledWith( - MOCK_SNAP_ID, - MOCK_MESSAGE, - ); + expect(handleKeyringSnapMessage).toHaveBeenCalledWith(MOCK_MESSAGE); expect(result).toStrictEqual({ ok: true }); }); - it('propagates errors thrown by the Snap keyring', async () => { - const { service, mocks } = setup(); + it('short-circuits GetSelectedAccounts by returning only the selected group accounts the Snap actually owns', async () => { + const { service, mocks } = await setup(); + mocks.AccountTreeController.getSelectedAccountGroup.mockReturnValue( + MOCK_GROUP_ID, + ); + mocks.AccountTreeController.getAccountGroupObject.mockReturnValue( + buildGroup(MOCK_GROUP_ID, MOCK_ACCOUNTS), + ); + // The Snap only owns the first account of the selected group. + mockWithKeyringV2Unsafe(mocks, { + [MOCK_SNAP_ID]: { + hasAccount: (id) => id === MOCK_ACCOUNTS[0], + }, + }); + + const result = await service.handleKeyringSnapMessage(MOCK_SNAP_ID, { + method: SnapManageAccountsMethod.GetSelectedAccounts, + params: {}, + } as unknown as SnapMessage); + + expect(result).toStrictEqual([MOCK_ACCOUNTS[0]]); + expect(mocks.KeyringController.withKeyringV2).not.toHaveBeenCalled(); + }); + + it('returns an empty array for GetSelectedAccounts when no account group is selected', async () => { + const { service, mocks } = await setup(); + mockWithKeyringV2Unsafe(mocks, { + [MOCK_SNAP_ID]: { hasAccount: () => true }, + }); + + const result = await service.handleKeyringSnapMessage(MOCK_SNAP_ID, { + method: SnapManageAccountsMethod.GetSelectedAccounts, + params: {}, + } as unknown as SnapMessage); + + expect(result).toStrictEqual([]); + }); + + it('throws a dedicated error when no v2 Snap keyring exists for the given Snap', async () => { + const { service, mocks } = await setup(); + mockWithKeyringV2(mocks, {}); + + await expect( + service.handleKeyringSnapMessage(MOCK_SNAP_ID, MOCK_MESSAGE), + ).rejects.toThrow( + `Cannot delegate keyring Snap message, keyring does not exist yet for Snap "${MOCK_SNAP_ID}".`, + ); + }); + + it('propagates non-KeyringNotFound errors thrown by the Snap keyring', async () => { + const { service, mocks } = await setup(); const error = new Error('snap boom'); const handleKeyringSnapMessage = jest.fn().mockRejectedValue(error); - mockLegacySnapKeyring(mocks, { handleKeyringSnapMessage }); + mockWithKeyringV2(mocks, { + [MOCK_SNAP_ID]: { handleKeyringSnapMessage }, + }); await expect( service.handleKeyringSnapMessage(MOCK_SNAP_ID, MOCK_MESSAGE), ).rejects.toThrow(error); }); + it('ensures the v2 keyring exists before forwarding an AccountCreated event', async () => { + const { service, mocks } = await setup(); + // `#ensureKeyringIsReady` uses `withController` — start with no keyring + // so it must create one. + const { addNewKeyring } = mockWithController(mocks, []); + const handleKeyringSnapMessage = jest.fn().mockResolvedValue(null); + // `withController` mock takes precedence over `withKeyringV2`; configure + // `withKeyringV2` separately for the forwarding step. + mockWithKeyringV2(mocks, { + [MOCK_SNAP_ID]: { handleKeyringSnapMessage }, + }); + + await service.handleKeyringSnapMessage( + MOCK_SNAP_ID, + MOCK_ACCOUNT_CREATED_MESSAGE, + ); + + expect(addNewKeyring).toHaveBeenCalledWith(KeyringType.Snap, { + snapId: MOCK_SNAP_ID, + accounts: {}, + }); + expect(handleKeyringSnapMessage).toHaveBeenCalledWith( + MOCK_ACCOUNT_CREATED_MESSAGE, + ); + }); + it('is exposed as a messenger action', async () => { - const { service, mocks, messenger } = setup(); + const { service, mocks, messenger } = await setup(); const handleKeyringSnapMessage = jest.fn().mockResolvedValue('pong'); - mockLegacySnapKeyring(mocks, { handleKeyringSnapMessage }); + mockWithKeyringV2(mocks, { + [MOCK_SNAP_ID]: { handleKeyringSnapMessage }, + }); // Reference `service` so it isn't flagged as unused; constructing it // registers the messenger action under test. @@ -683,10 +998,7 @@ describe('SnapAccountService', () => { MOCK_MESSAGE, ); - expect(handleKeyringSnapMessage).toHaveBeenCalledWith( - MOCK_SNAP_ID, - MOCK_MESSAGE, - ); + expect(handleKeyringSnapMessage).toHaveBeenCalledWith(MOCK_MESSAGE); expect(result).toBe('pong'); }); }); @@ -698,14 +1010,26 @@ describe('SnapAccountService', () => { '00000000-0000-0000-0000-000000000002', ]; - it('forwards the selected accounts to the Snap keyring', async () => { - const { service, rootMessenger, mocks } = setup(); - const setSelectedAccounts = jest.fn().mockResolvedValue(undefined); - mockLegacySnapKeyring(mocks, { setSelectedAccounts }); + it('forwards owned accounts to every tracked v2 Snap keyring in parallel', async () => { + const { rootMessenger, mocks } = await setup({ + runnableSnaps: [buildSnap(MOCK_SNAP_ID), buildSnap(MOCK_OTHER_SNAP_ID)], + }); + const setSelectedAccounts1 = jest.fn().mockResolvedValue(undefined); + const setSelectedAccounts2 = jest.fn().mockResolvedValue(undefined); + // Snap A owns the first account; Snap B owns the second. + mockWithKeyringV2Unsafe(mocks, { + [MOCK_SNAP_ID]: { + setSelectedAccounts: setSelectedAccounts1, + hasAccount: (id) => id === MOCK_ACCOUNTS[0], + }, + [MOCK_OTHER_SNAP_ID]: { + setSelectedAccounts: setSelectedAccounts2, + hasAccount: (id) => id === MOCK_ACCOUNTS[1], + }, + }); mocks.AccountTreeController.getAccountGroupObject.mockReturnValue( buildGroup(MOCK_GROUP_ID, MOCK_ACCOUNTS), ); - expect(service).toBeDefined(); publishSelectedAccountGroupChange(rootMessenger, MOCK_GROUP_ID); await flushMicrotasks(); @@ -713,14 +1037,74 @@ describe('SnapAccountService', () => { expect( mocks.AccountTreeController.getAccountGroupObject, ).toHaveBeenCalledWith(MOCK_GROUP_ID); - expect(setSelectedAccounts).toHaveBeenCalledWith(MOCK_ACCOUNTS); + expect(setSelectedAccounts1).toHaveBeenCalledWith([MOCK_ACCOUNTS[0]]); + expect(setSelectedAccounts2).toHaveBeenCalledWith([MOCK_ACCOUNTS[1]]); }); - it('does nothing when the new group ID is empty', async () => { - const { service, rootMessenger, mocks } = setup(); + it('forwards an empty list to a tracked Snap that owns none of the selected accounts', async () => { + const { rootMessenger, mocks } = await setup({ + runnableSnaps: [buildSnap(MOCK_SNAP_ID)], + }); const setSelectedAccounts = jest.fn().mockResolvedValue(undefined); - mockLegacySnapKeyring(mocks, { setSelectedAccounts }); + mockWithKeyringV2Unsafe(mocks, { + [MOCK_SNAP_ID]: { + setSelectedAccounts, + hasAccount: () => false, + }, + }); + mocks.AccountTreeController.getAccountGroupObject.mockReturnValue( + buildGroup(MOCK_GROUP_ID, MOCK_ACCOUNTS), + ); + + publishSelectedAccountGroupChange(rootMessenger, MOCK_GROUP_ID); + await flushMicrotasks(); + + expect(setSelectedAccounts).toHaveBeenCalledWith([]); + }); + + it('does nothing when no Snap is tracked', async () => { + const { service, rootMessenger, mocks } = await setup(); + mocks.AccountTreeController.getAccountGroupObject.mockReturnValue( + buildGroup(MOCK_GROUP_ID, MOCK_ACCOUNTS), + ); + + publishSelectedAccountGroupChange(rootMessenger, MOCK_GROUP_ID); + await flushMicrotasks(); + + expect( + mocks.KeyringController.withKeyringV2Unsafe, + ).not.toHaveBeenCalled(); expect(service).toBeDefined(); + }); + + it('silently skips a tracked Snap that has no v2 keyring yet', async () => { + const { service, rootMessenger, mocks } = await setup({ + runnableSnaps: [buildSnap(MOCK_SNAP_ID)], + }); + // No keyrings configured → withKeyringV2Unsafe throws KeyringNotFound. + mockWithKeyringV2Unsafe(mocks, {}); + mocks.AccountTreeController.getAccountGroupObject.mockReturnValue( + buildGroup(MOCK_GROUP_ID, MOCK_ACCOUNTS), + ); + const consoleErrorSpy = jest + .spyOn(console, 'error') + .mockImplementation(() => undefined); + + publishSelectedAccountGroupChange(rootMessenger, MOCK_GROUP_ID); + await flushMicrotasks(); + + expect(consoleErrorSpy).not.toHaveBeenCalled(); + expect(service).toBeDefined(); + consoleErrorSpy.mockRestore(); + }); + + it('does nothing when the new group ID is empty', async () => { + const { service, rootMessenger, mocks } = await setup({ + runnableSnaps: [buildSnap(MOCK_SNAP_ID)], + }); + mockWithKeyringV2Unsafe(mocks, { + [MOCK_SNAP_ID]: { setSelectedAccounts: jest.fn() }, + }); publishSelectedAccountGroupChange(rootMessenger, ''); await flushMicrotasks(); @@ -728,17 +1112,22 @@ describe('SnapAccountService', () => { expect( mocks.AccountTreeController.getAccountGroupObject, ).not.toHaveBeenCalled(); - expect(setSelectedAccounts).not.toHaveBeenCalled(); + expect( + mocks.KeyringController.withKeyringV2Unsafe, + ).not.toHaveBeenCalled(); + expect(service).toBeDefined(); }); it('does nothing when the account group is not found', async () => { - const { service, rootMessenger, mocks } = setup(); - const setSelectedAccounts = jest.fn().mockResolvedValue(undefined); - mockLegacySnapKeyring(mocks, { setSelectedAccounts }); + const { service, rootMessenger, mocks } = await setup({ + runnableSnaps: [buildSnap(MOCK_SNAP_ID)], + }); + mockWithKeyringV2Unsafe(mocks, { + [MOCK_SNAP_ID]: { setSelectedAccounts: jest.fn() }, + }); mocks.AccountTreeController.getAccountGroupObject.mockReturnValue( undefined, ); - expect(service).toBeDefined(); publishSelectedAccountGroupChange(rootMessenger, MOCK_GROUP_ID); await flushMicrotasks(); @@ -746,31 +1135,75 @@ describe('SnapAccountService', () => { expect( mocks.AccountTreeController.getAccountGroupObject, ).toHaveBeenCalledWith(MOCK_GROUP_ID); - expect(setSelectedAccounts).not.toHaveBeenCalled(); + expect( + mocks.KeyringController.withKeyringV2Unsafe, + ).not.toHaveBeenCalled(); + expect(service).toBeDefined(); }); - it('logs an error when forwarding to the Snap keyring fails', async () => { - const { service, rootMessenger, mocks } = setup(); + it('logs an error when forwarding to a v2 Snap keyring fails, but still forwards to the others', async () => { + const { service, rootMessenger, mocks } = await setup({ + runnableSnaps: [buildSnap(MOCK_SNAP_ID), buildSnap(MOCK_OTHER_SNAP_ID)], + }); const error = new Error('forward boom'); - const setSelectedAccounts = jest.fn().mockRejectedValue(error); - mockLegacySnapKeyring(mocks, { setSelectedAccounts }); + const setSelectedAccounts1 = jest.fn().mockRejectedValue(error); + const setSelectedAccounts2 = jest.fn().mockResolvedValue(undefined); + mockWithKeyringV2Unsafe(mocks, { + [MOCK_SNAP_ID]: { setSelectedAccounts: setSelectedAccounts1 }, + [MOCK_OTHER_SNAP_ID]: { setSelectedAccounts: setSelectedAccounts2 }, + }); mocks.AccountTreeController.getAccountGroupObject.mockReturnValue( buildGroup(MOCK_GROUP_ID, MOCK_ACCOUNTS), ); const consoleErrorSpy = jest .spyOn(console, 'error') .mockImplementation(() => undefined); - expect(service).toBeDefined(); publishSelectedAccountGroupChange(rootMessenger, MOCK_GROUP_ID); await flushMicrotasks(); - expect(setSelectedAccounts).toHaveBeenCalledWith(MOCK_ACCOUNTS); + expect(setSelectedAccounts1).toHaveBeenCalled(); + expect(setSelectedAccounts2).toHaveBeenCalledWith(MOCK_ACCOUNTS); expect(consoleErrorSpy).toHaveBeenCalledWith( - 'Error forwarding selected accounts:', + `Error forwarding selected accounts to Snap "${MOCK_SNAP_ID}":`, error, ); + expect(service).toBeDefined(); + consoleErrorSpy.mockRestore(); + }); + it('logs a top-level error if forwarding itself rejects unexpectedly', async () => { + const { service, rootMessenger, mocks } = await setup({ + runnableSnaps: [buildSnap(MOCK_SNAP_ID)], + }); + const innerError = new Error('inner boom'); + mockWithKeyringV2Unsafe(mocks, { + [MOCK_SNAP_ID]: { + setSelectedAccounts: jest.fn().mockRejectedValue(innerError), + }, + }); + mocks.AccountTreeController.getAccountGroupObject.mockReturnValue( + buildGroup(MOCK_GROUP_ID, MOCK_ACCOUNTS), + ); + // Force the per-Snap error handler itself to throw on its first + // invocation, so the rejection escapes the inner try/catch and reaches + // the outer `.catch` (the top-level fallback). Subsequent calls + // (including the top-level one) no-op. + const consoleErrorSpy = jest + .spyOn(console, 'error') + .mockImplementation(() => undefined); + consoleErrorSpy.mockImplementationOnce(() => { + throw new Error('logger boom'); + }); + + publishSelectedAccountGroupChange(rootMessenger, MOCK_GROUP_ID); + await flushMicrotasks(); + + expect(consoleErrorSpy).toHaveBeenCalledWith( + 'Error forwarding selected accounts:', + expect.any(Error), + ); + expect(service).toBeDefined(); consoleErrorSpy.mockRestore(); }); }); @@ -782,17 +1215,20 @@ describe('SnapAccountService', () => { '00000000-0000-0000-0000-000000000002', ]; - it('forwards the currently selected account group to the Snap keyring', async () => { - const { service, rootMessenger, mocks } = setup(); + it('forwards the currently selected account group to every tracked v2 Snap keyring', async () => { + const { service, rootMessenger, mocks } = await setup({ + runnableSnaps: [buildSnap(MOCK_SNAP_ID)], + }); const setSelectedAccounts = jest.fn().mockResolvedValue(undefined); - mockLegacySnapKeyring(mocks, { setSelectedAccounts }); + mockWithKeyringV2Unsafe(mocks, { + [MOCK_SNAP_ID]: { setSelectedAccounts }, + }); mocks.AccountTreeController.getSelectedAccountGroup.mockReturnValue( MOCK_GROUP_ID, ); mocks.AccountTreeController.getAccountGroupObject.mockReturnValue( buildGroup(MOCK_GROUP_ID, MOCK_ACCOUNTS), ); - expect(service).toBeDefined(); publishUnlock(rootMessenger); await flushMicrotasks(); @@ -804,14 +1240,17 @@ describe('SnapAccountService', () => { mocks.AccountTreeController.getAccountGroupObject, ).toHaveBeenCalledWith(MOCK_GROUP_ID); expect(setSelectedAccounts).toHaveBeenCalledWith(MOCK_ACCOUNTS); + expect(service).toBeDefined(); }); it('does nothing when no account group is selected', async () => { - const { service, rootMessenger, mocks } = setup(); - const setSelectedAccounts = jest.fn().mockResolvedValue(undefined); - mockLegacySnapKeyring(mocks, { setSelectedAccounts }); + const { service, rootMessenger, mocks } = await setup({ + runnableSnaps: [buildSnap(MOCK_SNAP_ID)], + }); + mockWithKeyringV2Unsafe(mocks, { + [MOCK_SNAP_ID]: { setSelectedAccounts: jest.fn() }, + }); mocks.AccountTreeController.getSelectedAccountGroup.mockReturnValue(''); - expect(service).toBeDefined(); publishUnlock(rootMessenger); await flushMicrotasks(); @@ -819,14 +1258,21 @@ describe('SnapAccountService', () => { expect( mocks.AccountTreeController.getAccountGroupObject, ).not.toHaveBeenCalled(); - expect(setSelectedAccounts).not.toHaveBeenCalled(); + expect( + mocks.KeyringController.withKeyringV2Unsafe, + ).not.toHaveBeenCalled(); + expect(service).toBeDefined(); }); - it('logs an error when forwarding to the Snap keyring fails', async () => { - const { service, rootMessenger, mocks } = setup(); + it('logs an error when forwarding to a v2 Snap keyring fails', async () => { + const { service, rootMessenger, mocks } = await setup({ + runnableSnaps: [buildSnap(MOCK_SNAP_ID)], + }); const error = new Error('forward boom'); const setSelectedAccounts = jest.fn().mockRejectedValue(error); - mockLegacySnapKeyring(mocks, { setSelectedAccounts }); + mockWithKeyringV2Unsafe(mocks, { + [MOCK_SNAP_ID]: { setSelectedAccounts }, + }); mocks.AccountTreeController.getSelectedAccountGroup.mockReturnValue( MOCK_GROUP_ID, ); @@ -836,17 +1282,16 @@ describe('SnapAccountService', () => { const consoleErrorSpy = jest .spyOn(console, 'error') .mockImplementation(() => undefined); - expect(service).toBeDefined(); publishUnlock(rootMessenger); await flushMicrotasks(); expect(setSelectedAccounts).toHaveBeenCalledWith(MOCK_ACCOUNTS); expect(consoleErrorSpy).toHaveBeenCalledWith( - 'Error forwarding selected accounts:', + `Error forwarding selected accounts to Snap "${MOCK_SNAP_ID}":`, error, ); - + expect(service).toBeDefined(); consoleErrorSpy.mockRestore(); }); }); @@ -863,13 +1308,16 @@ describe('SnapAccountService', () => { ]; it('forwards the accounts from the event payload when the affected group is the selected one', async () => { - const { service, rootMessenger, mocks } = setup(); + const { service, rootMessenger, mocks } = await setup({ + runnableSnaps: [buildSnap(MOCK_SNAP_ID)], + }); const setSelectedAccounts = jest.fn().mockResolvedValue(undefined); - mockLegacySnapKeyring(mocks, { setSelectedAccounts }); + mockWithKeyringV2Unsafe(mocks, { + [MOCK_SNAP_ID]: { setSelectedAccounts }, + }); mocks.AccountTreeController.getSelectedAccountGroup.mockReturnValue( MOCK_GROUP_ID, ); - expect(service).toBeDefined(); publishEvent(rootMessenger, buildGroup(MOCK_GROUP_ID, MOCK_ACCOUNTS)); await flushMicrotasks(); @@ -878,16 +1326,19 @@ describe('SnapAccountService', () => { mocks.AccountTreeController.getAccountGroupObject, ).not.toHaveBeenCalled(); expect(setSelectedAccounts).toHaveBeenCalledWith(MOCK_ACCOUNTS); + expect(service).toBeDefined(); }); it('does nothing when the affected group is not the selected one', async () => { - const { service, rootMessenger, mocks } = setup(); - const setSelectedAccounts = jest.fn().mockResolvedValue(undefined); - mockLegacySnapKeyring(mocks, { setSelectedAccounts }); + const { service, rootMessenger, mocks } = await setup({ + runnableSnaps: [buildSnap(MOCK_SNAP_ID)], + }); + mockWithKeyringV2Unsafe(mocks, { + [MOCK_SNAP_ID]: { setSelectedAccounts: jest.fn() }, + }); mocks.AccountTreeController.getSelectedAccountGroup.mockReturnValue( OTHER_GROUP_ID, ); - expect(service).toBeDefined(); publishEvent(rootMessenger, buildGroup(MOCK_GROUP_ID, MOCK_ACCOUNTS)); await flushMicrotasks(); @@ -895,15 +1346,20 @@ describe('SnapAccountService', () => { expect( mocks.AccountTreeController.getAccountGroupObject, ).not.toHaveBeenCalled(); - expect(setSelectedAccounts).not.toHaveBeenCalled(); + expect( + mocks.KeyringController.withKeyringV2Unsafe, + ).not.toHaveBeenCalled(); + expect(service).toBeDefined(); }); it('does nothing when no account group is selected', async () => { - const { service, rootMessenger, mocks } = setup(); - const setSelectedAccounts = jest.fn().mockResolvedValue(undefined); - mockLegacySnapKeyring(mocks, { setSelectedAccounts }); + const { service, rootMessenger, mocks } = await setup({ + runnableSnaps: [buildSnap(MOCK_SNAP_ID)], + }); + mockWithKeyringV2Unsafe(mocks, { + [MOCK_SNAP_ID]: { setSelectedAccounts: jest.fn() }, + }); mocks.AccountTreeController.getSelectedAccountGroup.mockReturnValue(''); - expect(service).toBeDefined(); publishEvent(rootMessenger, buildGroup(MOCK_GROUP_ID, MOCK_ACCOUNTS)); await flushMicrotasks(); @@ -911,7 +1367,10 @@ describe('SnapAccountService', () => { expect( mocks.AccountTreeController.getAccountGroupObject, ).not.toHaveBeenCalled(); - expect(setSelectedAccounts).not.toHaveBeenCalled(); + expect( + mocks.KeyringController.withKeyringV2Unsafe, + ).not.toHaveBeenCalled(); + expect(service).toBeDefined(); }); }); @@ -919,14 +1378,17 @@ describe('SnapAccountService', () => { const MOCK_GROUP_ID = 'keyring:01JABC/group-1' as AccountGroupId; const OTHER_GROUP_ID = 'keyring:01JABC/group-2' as AccountGroupId; - it('clears the selected accounts when the removed group is the selected one', async () => { - const { service, rootMessenger, mocks } = setup(); + it('clears the selected accounts on every tracked v2 Snap keyring when the removed group is the selected one', async () => { + const { service, rootMessenger, mocks } = await setup({ + runnableSnaps: [buildSnap(MOCK_SNAP_ID)], + }); const setSelectedAccounts = jest.fn().mockResolvedValue(undefined); - mockLegacySnapKeyring(mocks, { setSelectedAccounts }); + mockWithKeyringV2Unsafe(mocks, { + [MOCK_SNAP_ID]: { setSelectedAccounts }, + }); mocks.AccountTreeController.getSelectedAccountGroup.mockReturnValue( MOCK_GROUP_ID, ); - expect(service).toBeDefined(); publishAccountGroupRemoved(rootMessenger, MOCK_GROUP_ID); await flushMicrotasks(); @@ -935,16 +1397,19 @@ describe('SnapAccountService', () => { mocks.AccountTreeController.getAccountGroupObject, ).not.toHaveBeenCalled(); expect(setSelectedAccounts).toHaveBeenCalledWith([]); + expect(service).toBeDefined(); }); it('does nothing when the removed group is not the selected one', async () => { - const { service, rootMessenger, mocks } = setup(); - const setSelectedAccounts = jest.fn().mockResolvedValue(undefined); - mockLegacySnapKeyring(mocks, { setSelectedAccounts }); + const { service, rootMessenger, mocks } = await setup({ + runnableSnaps: [buildSnap(MOCK_SNAP_ID)], + }); + mockWithKeyringV2Unsafe(mocks, { + [MOCK_SNAP_ID]: { setSelectedAccounts: jest.fn() }, + }); mocks.AccountTreeController.getSelectedAccountGroup.mockReturnValue( OTHER_GROUP_ID, ); - expect(service).toBeDefined(); publishAccountGroupRemoved(rootMessenger, MOCK_GROUP_ID); await flushMicrotasks(); @@ -952,7 +1417,10 @@ describe('SnapAccountService', () => { expect( mocks.AccountTreeController.getAccountGroupObject, ).not.toHaveBeenCalled(); - expect(setSelectedAccounts).not.toHaveBeenCalled(); + expect( + mocks.KeyringController.withKeyringV2Unsafe, + ).not.toHaveBeenCalled(); + expect(service).toBeDefined(); }); }); }); diff --git a/packages/snap-account-service/src/SnapAccountService.ts b/packages/snap-account-service/src/SnapAccountService.ts index 7f73096178..645da69a4f 100644 --- a/packages/snap-account-service/src/SnapAccountService.ts +++ b/packages/snap-account-service/src/SnapAccountService.ts @@ -3,14 +3,22 @@ import type { SnapKeyring as LegacySnapKeyring, SnapMessage, } from '@metamask/eth-snap-keyring'; +import { SnapKeyring, SnapKeyringState } from '@metamask/eth-snap-keyring/v2'; +import { KeyringEvent } from '@metamask/keyring-api'; +import { Keyring, KeyringType } from '@metamask/keyring-api/v2'; import type { KeyringControllerGetStateAction, KeyringControllerStateChangeEvent, KeyringControllerUnlockEvent, KeyringControllerWithControllerAction, - KeyringEntry, + KeyringControllerWithKeyringV2Action, + KeyringControllerWithKeyringV2UnsafeAction, } from '@metamask/keyring-controller'; -import { KeyringTypes } from '@metamask/keyring-controller'; +import { + isKeyringNotFoundError, + KeyringTypes, +} from '@metamask/keyring-controller'; +import { SnapManageAccountsMethod } from '@metamask/keyring-snap-sdk'; import type { AccountId, BaseKeyring } from '@metamask/keyring-utils'; import type { Messenger } from '@metamask/messenger'; import type { @@ -31,9 +39,9 @@ import type { Json } from '@metamask/utils'; import { projectLogger as log } from './logger'; import type { SnapAccountServiceEnsureReadyAction, - SnapAccountServiceGetLegacySnapKeyringAction, SnapAccountServiceGetSnapsAction, SnapAccountServiceHandleKeyringSnapMessageAction, + SnapAccountServiceMigrateAction, } from './SnapAccountService-method-action-types'; import { SnapPlatformWatcher } from './SnapPlatformWatcher'; import type { SnapPlatformWatcherConfig } from './SnapPlatformWatcher'; @@ -61,8 +69,8 @@ export const serviceName = 'SnapAccountService'; const MESSENGER_EXPOSED_METHODS = [ 'ensureReady', 'getSnaps', - 'getLegacySnapKeyring', 'handleKeyringSnapMessage', + 'migrate', ] as const; /** @@ -71,8 +79,8 @@ const MESSENGER_EXPOSED_METHODS = [ export type SnapAccountServiceActions = | SnapAccountServiceEnsureReadyAction | SnapAccountServiceGetSnapsAction - | SnapAccountServiceGetLegacySnapKeyringAction - | SnapAccountServiceHandleKeyringSnapMessageAction; + | SnapAccountServiceHandleKeyringSnapMessageAction + | SnapAccountServiceMigrateAction; /** * Actions from other messengers that {@link SnapAccountService} calls. @@ -83,6 +91,8 @@ type AllowedActions = | SnapControllerGetRunnableSnapsAction | KeyringControllerGetStateAction | KeyringControllerWithControllerAction + | KeyringControllerWithKeyringV2Action + | KeyringControllerWithKeyringV2UnsafeAction | AccountTreeControllerGetAccountGroupObjectAction | AccountTreeControllerGetSelectedAccountGroupAction; @@ -147,6 +157,17 @@ function isLegacySnapKeyring(keyring: { return keyring.type === KeyringTypes.snap; } +/** + * Checks if a given keyring is a Snap keyring (v2). + * + * @param keyring - The keyring to check. + * @returns `true` if the keyring is a Snap keyring (v2), `false` otherwise. + */ +function isSnapKeyring(keyring: Keyring): keyring is SnapKeyring { + // Using `KeyringType.Snap` (used for v2). + return keyring.type === KeyringType.Snap; +} + /** * Service responsible for managing account management snaps. */ @@ -162,6 +183,10 @@ export class SnapAccountService { readonly #tracker: SnapTracker; + #migrated = false; + + #migratePromise: Promise | null = null; + /** * Constructs a new {@link SnapAccountService}. * @@ -285,7 +310,11 @@ export class SnapAccountService { /** * Ensures everything is ready to use Snap accounts for the given Snap. * 1. Validates that `snapId` is a tracked account-management Snap. - * 2. Waits for the Snap platform to be fully started. + * 2. Runs the legacy -> v2 Snap keyring migration (cached — no-op if + * already done). + * 3. Atomically creates the v2 keyring for this Snap if it doesn't exist + * yet. + * 4. Waits for the Snap platform to be fully started. * * Safe to call concurrently — each step is idempotent or mutex-protected. * @@ -296,56 +325,210 @@ export class SnapAccountService { if (!this.#tracker.canUse(snapId)) { throw new Error(`Unknown snap: "${snapId}"`); } + + // Migrate from the global v1 Snap keyring to the per-Snap v2 keyring + // before doing anything else. + await this.migrate(); + + // We still try to create the keyring for the Snap here, since we might + // want to use a new Snap that never had accounts before. + await this.#ensureKeyringIsReady(snapId); + // Before doing anything with our Snap, we need to make sure the platform // is ready to process requests. await this.#watcher.ensureCanUseSnapPlatform(); } /** - * Atomically gets-or-creates the legacy (v1) Snap keyring — the keyring - * associated with {@link KeyringTypes.snap}. + * Migrate the legacy Snap keyring to the new (per-snap) Snap keyring v2. + * Safe to call concurrently — the migration runs only once; all callers + * await the same promise. * - * @returns The existing or newly-created Snap keyring instance. + * @returns A promise that resolves when the migration is complete. */ - async getLegacySnapKeyring(): Promise { - type Result = { - snapKeyring: LegacySnapKeyring; - }; + async migrate(): Promise { + if (this.#migrated) { + return; + } + if (!this.#migratePromise) { + this.#migratePromise = this.#migrate(); + + try { + await this.#migratePromise; + + // Only mark it as migrated after the migration logic completes successfully. If + // it fails, we want future calls to retry the migration. + this.#migrated = true; + } finally { + this.#migratePromise = null; + } + } + await this.#migratePromise; + } - // `KeyringController:withController` forbids returning a direct keyring - // reference (it checks the result via `Object.is`), so we smuggle the - // instance out wrapped in an object and unwrap it after the call. - // NOTE: This violates the abstraction of `KeyringController:withController`, but this - // is how we currently interact with the legacy Snap keyring. Once we migrate it to - // the Snap keyring v2, we won't be using the same pattern. - const result = await this.#messenger.call( + /** + * Performs the actual migration logic. Should only be called once, and is not + * safe to call concurrently. + */ + async #migrate(): Promise { + await this.#messenger.call( 'KeyringController:withController', - async (controller): Promise => { - let snapKeyring: KeyringEntry['keyring'] | undefined; + async (controller) => { + const { keyrings } = controller; - const found = controller.keyrings.find(({ keyring }) => + const legacySnapKeyringEntry = keyrings.find(({ keyring }) => isLegacySnapKeyring(keyring), ); - if (found) { - snapKeyring = found.keyring; + if (!legacySnapKeyringEntry) { + log('No legacy Snap keyring found. Migration not required.'); + return; } - if (!snapKeyring) { - const { - keyring: newSnapKeyring, - metadata: { id }, - } = await controller.addNewKeyring(KeyringTypes.snap); - snapKeyring = newSnapKeyring; + log('Migration started...'); + + // The legacy Snap keyring has never been a true `EthKeyring` so we + // need to cast it to `unknown` first. + const legacySnapKeyring = + legacySnapKeyringEntry.keyring as unknown as LegacySnapKeyring; + + // Compute the account list for each Snap, grouped by snap ID. + const states = new Map(); + for (const internalAccount of legacySnapKeyring.listAccounts()) { + // Convert `InternalAccount` to `KeyringAccount` since the Snap + // keyring (v2) expects accounts in that format and will verify it + // with `superstruct` when adding the keyring. + const { metadata, ...account } = internalAccount; + + const snap = metadata?.snap; + if (snap) { + const snapId = snap.id as SnapId; + + let state = states.get(snapId); + if (!state) { + state = { snapId, accounts: {} }; + states.set(snapId, state); + } + state.accounts[account.id] = account; + } + } - log(`Legacy Snap keyring created. ("${id}")`); + // Create the new Snap keyring (v2) for each Snap and migrate the + // accounts over. + for (const state of states.values()) { + log(`Migrating accounts for Snap "${state.snapId}"...`); + await controller.addNewKeyring( + // IMPORTANT: The Snap keyring (v2) can also be used as a v1 + // keyring. So the builder associated with the v2 keyring type is + // able to build both v1 and v2 keyrings. + KeyringType.Snap, + state, + ); } - // The legacy Snap keyring is not compatible with `EthKeyring`, so we need to cast here. - return { snapKeyring } as unknown as Result; + // Remove the legacy Snap keyring after migration. + log('Removing legacy Snap keyring...'); + await controller.removeKeyring(legacySnapKeyringEntry.metadata.id); + + log('Migration completed!'); + }, + ); + } + + /** + * Ensures a Snap keyring is ready for the given Snap. If it doesn't exist yet, it will be created. + * Safe to call concurrently. + * + * @param snapId - The Snap ID to ensure the keyring is ready for. + */ + async #ensureKeyringIsReady(snapId: SnapId): Promise { + await this.#messenger.call( + 'KeyringController:withController', + async (controller) => { + const hasKeyring = controller.keyrings.some( + ({ keyringV2 }) => + keyringV2 && + isSnapKeyring(keyringV2) && + keyringV2.snapId === snapId, + ); + + if (!hasKeyring) { + log(`Creating v2 keyring for Snap "${snapId}"...`); + await controller.addNewKeyring(KeyringType.Snap, { + snapId, + accounts: {}, + }); + } + }, + ); + } + + /** + * Shared body for {@link SnapAccountService.#withKeyringV2} and + * {@link SnapAccountService.#withKeyringV2Unsafe}. Hides the per-Snap + * filter and the cast back to {@link SnapKeyring} (the messenger action's + * callback receives a generic `Keyring`; the selector's type predicate + * doesn't flow through the messenger's generics). + * + * @param action - The messenger action to invoke. + * @param snapId - The Snap ID to look up the keyring for. + * @param operation - The operation to run with the matching keyring. + * @returns The result of the operation. + */ + async #withKeyringV2Call( + action: + | 'KeyringController:withKeyringV2' + | 'KeyringController:withKeyringV2Unsafe', + snapId: SnapId, + operation: (keyring: SnapKeyring) => Promise, + ): Promise { + return this.#messenger.call( + action, + { + filter: (keyring): keyring is SnapKeyring => + isSnapKeyring(keyring) && keyring.snapId === snapId, }, + async ({ keyring }) => operation(keyring as SnapKeyring), + ) as Result; + } + + /** + * Runs an operation against the v2 Snap keyring for the given Snap, under + * the `KeyringController` mutex. Throws `KeyringNotFound` if no v2 Snap + * keyring exists for the given Snap. + * + * @param snapId - The Snap ID to look up the keyring for. + * @param operation - The operation to run with the matching keyring. + * @returns The result of the operation. + */ + async #withKeyringV2( + snapId: SnapId, + operation: (keyring: SnapKeyring) => Promise, + ): Promise { + return this.#withKeyringV2Call( + 'KeyringController:withKeyringV2', + snapId, + operation, ); + } - return (result as Result).snapKeyring; + /** + * Lock-free variant of {@link SnapAccountService.#withKeyringV2}. Only use + * for operations that do not mutate keyring or controller state — see + * `KeyringController.withKeyringV2Unsafe` for the contract. + * + * @param snapId - The Snap ID to look up the keyring for. + * @param operation - The operation to run with the matching keyring. + * @returns The result of the operation. + */ + async #withKeyringV2Unsafe( + snapId: SnapId, + operation: (keyring: SnapKeyring) => Promise, + ): Promise { + return this.#withKeyringV2Call( + 'KeyringController:withKeyringV2Unsafe', + snapId, + operation, + ); } /** @@ -359,8 +542,50 @@ export class SnapAccountService { snapId: SnapId, message: SnapMessage, ): Promise { - const snapKeyring = await this.getLegacySnapKeyring(); - return snapKeyring.handleKeyringSnapMessage(snapId, message); + // Handle specific methods first. + if (message.method === SnapManageAccountsMethod.GetSelectedAccounts) { + const groupId = this.#getSelectedAccountGroupId(); + const accounts = this.#getAccountGroup(groupId)?.accounts ?? []; + + return await this.#withKeyringV2Unsafe(snapId, async (keyring) => + accounts.filter((id) => keyring.hasAccount(id)), + ); + } + + const event = message.method as KeyringEvent; // We assume the Snap platform always sends a valid `KeyringEvent` here. + log( + `Forwarding message "${event}" from Snap "${snapId}" to its keyring...`, + ); + + // We can create a new keyring if the message is an AccountCreated event. + const isAccountCreatedMessage = event === KeyringEvent.AccountCreated; + + // Create the Snap keyring if it doesn't exist yet (in an atomic way). We cannot assume + // the keyring exists (e.g for the MMI Snap). + // NOTE: We only auto-create it for v1 account creation flows. + if (isAccountCreatedMessage) { + await this.#ensureKeyringIsReady(snapId); + } + + // This part of the flow relies on v1 flows, but v2 keyrings are compatible with those messages + // too. + try { + return await this.#withKeyringV2(snapId, async (keyring) => + keyring.handleKeyringSnapMessage(message), + ); + } catch (error) { + if (isKeyringNotFoundError(error)) { + log( + `No Snap keyring found for Snap "${snapId}". Cannot handle message with method "${event}".`, + ); + + throw new Error( + `Cannot delegate keyring Snap message, keyring does not exist yet for Snap "${snapId}".`, + ); + } + + throw error; + } } /** @@ -388,20 +613,47 @@ export class SnapAccountService { return; } - const forwardSelectedAccounts = async (): Promise => { - if (accounts.length) { - log( - `Forwarding selected accounts (from "${groupId}"): ${accounts.join(', ')}`, - ); - } else { - log(`Clearing selected accounts (from "${groupId}")`); - } + if (accounts.length) { + log( + `Forwarding selected accounts (from "${groupId}"): ${accounts.join(', ')}`, + ); + } else { + log(`Clearing selected accounts (from "${groupId}")`); + } - const snapKeyring = await this.getLegacySnapKeyring(); - await snapKeyring.setSelectedAccounts(accounts); + const forwardSelectedAccounts = async (): Promise => { + await Promise.all( + this.#tracker.getSnaps().map(async (snapId) => { + try { + // We can safely invoke this method without taking the controller lock + // because it should not mutate the keyring state. So we can use + // `withKeyringV2Unsafe` in this case. + await this.#withKeyringV2Unsafe(snapId, async (keyring) => { + // The group's accounts may belong to several Snaps; only + // forward the subset this Snap actually owns. An empty + // subset still gets forwarded to explicitly clear the + // Snap selected accounts. + await keyring.setSelectedAccounts( + accounts.filter((id) => keyring.hasAccount(id)), + ); + }); + } catch (error) { + // Tracked Snaps without a v2 keyring yet are expected — + // forwarding will resume on the next event once `ensureReady` + // has run. + if (!isKeyringNotFoundError(error)) { + console.error( + `Error forwarding selected accounts to Snap "${snapId}":`, + error, + ); + } + } + }), + ); }; - // There is nothing we can do if forwarding fails. This will auto-recover on the next relevant event. + // There is nothing we can do if forwarding fails. This will auto-recover on + // the next relevant event. forwardSelectedAccounts().catch((error) => { console.error('Error forwarding selected accounts:', error); }); diff --git a/yarn.lock b/yarn.lock index c4599a4918..b0061ebd97 100644 --- a/yarn.lock +++ b/yarn.lock @@ -5462,7 +5462,9 @@ __metadata: "@metamask/account-tree-controller": "npm:^7.4.0" "@metamask/auto-changelog": "npm:^6.1.0" "@metamask/eth-snap-keyring": "npm:^22.0.1" + "@metamask/keyring-api": "npm:^23.1.0" "@metamask/keyring-controller": "npm:^25.5.0" + "@metamask/keyring-snap-sdk": "npm:^9.0.1" "@metamask/keyring-utils": "npm:^3.2.1" "@metamask/messenger": "npm:^1.2.0" "@metamask/snaps-controllers": "npm:^19.0.0"