feat(snap-account-service): add migration logic + keyring v2 support

Author: ccharly

packages/snap-account-service/CHANGELOG.md

@@ -10,11 +10,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: "<id>"` 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))
@@ -32,6 +33,10 @@ 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`.
+  - 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.
 
 ### Changed
 

packages/snap-account-service/package.json

@@ -56,6 +56,7 @@
     "@metamask/account-api": "^1.0.4",
     "@metamask/account-tree-controller": "^7.3.0",
     "@metamask/eth-snap-keyring": "^22.0.1",
+    "@metamask/keyring-api": "^23.1.0",
     "@metamask/keyring-controller": "^25.5.0",
     "@metamask/messenger": "^1.2.0",
     "@metamask/snaps-controllers": "^19.0.0",

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.
  *
@@ -32,6 +36,18 @@ export type SnapAccountServiceEnsureReadyAction = {
   handler: SnapAccountService['ensureReady'];
 };
 
+/**
+ * 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 A promise that resolves when the migration is complete.
+ */
+export type SnapAccountServiceMigrateAction = {
+  type: `SnapAccountService:migrate`;
+  handler: SnapAccountService['migrate'];
+};
+
 /**
  * Atomically gets-or-creates the legacy (v1) Snap keyring — the keyring
  * associated with {@link KeyringTypes.snap}.
@@ -61,5 +77,6 @@ export type SnapAccountServiceHandleKeyringSnapMessageAction = {
 export type SnapAccountServiceMethodActions =
   | SnapAccountServiceGetSnapsAction
   | SnapAccountServiceEnsureReadyAction
+  | SnapAccountServiceMigrateAction
   | SnapAccountServiceGetLegacySnapKeyringAction
   | SnapAccountServiceHandleKeyringSnapMessageAction;

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

@@ -1,5 +1,7 @@
 import type { AccountGroupId } from '@metamask/account-api';
+import { SNAP_KEYRING_TYPE } from '@metamask/eth-snap-keyring';
 import type { SnapKeyring, SnapMessage } from '@metamask/eth-snap-keyring';
+import { KeyringType } from '@metamask/keyring-api/v2';
 import {
   KeyringControllerState,
   KeyringTypes,
@@ -436,6 +438,7 @@ function setup({
 }
 
 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', () => {
@@ -458,6 +461,93 @@ describe('SnapAccountService', () => {
     });
   });
 
+  describe('migrate', () => {
+    it('runs the migration only once when called concurrently', async () => {
+      const { service, mocks } = 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 } = 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 } = 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);
+    });
+  });
+
   describe('ensureReady', () => {
     it('throws when the Snap is not tracked', async () => {
       const { service } = setup();
@@ -489,6 +579,73 @@ describe('SnapAccountService', () => {
       expect(await service.ensureReady(MOCK_SNAP_ID)).toBeUndefined();
     });
 
+    it('runs the migration before checking platform readiness', async () => {
+      const { service, mocks } = setup({
+        runnableSnaps: [buildSnap(MOCK_SNAP_ID, true)],
+      });
+      mocks.KeyringController.withController.mockResolvedValue(undefined);
+
+      await service.init();
+      // `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 } = setup({
+        runnableSnaps: [buildSnap(MOCK_SNAP_ID, true)],
+      });
+      mocks.KeyringController.withController.mockImplementation(
+        async (operation) =>
+          operation({ keyrings: [], addNewKeyring, removeKeyring }),
+      );
+
+      await service.init();
+      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 } = setup({
+        runnableSnaps: [buildSnap(MOCK_SNAP_ID, true)],
+      });
+      // 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.init();
+      await service.ensureReady(MOCK_SNAP_ID);
+
+      expect(addNewKeyring).not.toHaveBeenCalled();
+    });
+
     it('waits for the Snap platform to become ready', async () => {
       const { service, rootMessenger } = setup({
         snapIsReady: false,
@@ -525,9 +682,9 @@ describe('SnapAccountService', () => {
         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);
 
@@ -584,7 +741,7 @@ describe('SnapAccountService', () => {
         return undefined;
       });
 
-      await Promise.resolve();
+      await flushMicrotasks();
       expect(ensureOnboardingComplete).toHaveBeenCalledTimes(1);
       expect(resolved).toBe(false);