feat(multichain-account-service): use withKeyringV2 for Snap account providers

Author: ccharly

packages/multichain-account-service/CHANGELOG.md

@@ -9,6 +9,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.

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<SelectedKeyring, CallbackResult = void>(
-    selector: KeyringSelector,
-    operation: ({
-      keyring,
-      metadata,
-    }: {
-      keyring: SelectedKeyring;
-      metadata: KeyringMetadata;
-    }) => Promise<CallbackResult>,
-  ): Promise<CallbackResult> {
-    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`.
    *

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();
     });
 

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<SnapAccountProviderConfig>;
   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(),
+    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);
@@ -866,9 +884,8 @@ describe('SnapAccountProvider', () => {
       ).toHaveBeenCalledWith(extraSnapAccount2.id);
 
       // Should remove from keyring and recreate the missing account
-      expect(keyring.removeAccount).toHaveBeenCalledWith(
-        mockAccounts[1].address,
-      );
+      // (the keyring has extraSnapAccount2 with the same address as mockAccounts[1])
+      expect(keyring.removeAccount).toHaveBeenCalledWith(extraSnapAccount2.id);
       expect(createAccountsSpy).toHaveBeenCalledWith({
         entropySource: mockAccounts[1].options.entropy.id,
         groupIndex: mockAccounts[1].options.entropy.groupIndex,
@@ -949,6 +966,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();