feat(snap-account-service): add handleKeyringSnapMessage

Author: ccharly

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,9 +36,26 @@ export type SnapAccountServiceEnsureReadyAction = {
   handler: SnapAccountService['ensureReady'];
 };
 
+/**
+ * Handle a message from a Snap.
+ *
+ * Only `AccountCreated` triggers lazy keyring creation, since that is the
+ * single entry point for the v1 event-driven flow. All other events from
+ * unknown Snaps throw an error.
+ *
+ * @param snapId - ID of the Snap.
+ * @param message - Message sent by the Snap.
+ * @returns The execution result.
+ */
+export type SnapAccountServiceHandleKeyringSnapMessageAction = {
+  type: `SnapAccountService:handleKeyringSnapMessage`;
+  handler: SnapAccountService['handleKeyringSnapMessage'];
+};
+
 /**
  * Union of all SnapAccountService action types.
  */
 export type SnapAccountServiceMethodActions =
   | SnapAccountServiceGetSnapsAction
-  | SnapAccountServiceEnsureReadyAction;
+  | SnapAccountServiceEnsureReadyAction
+  | SnapAccountServiceHandleKeyringSnapMessageAction;

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

@@ -43,6 +43,14 @@ type Mocks = {
     withController: jest.MockedFunction<
       (callback: (controller: any) => Promise<unknown>) => Promise<unknown>
     >;
+    withKeyringV2Unsafe: jest.MockedFunction<
+      (
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        selector: any,
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        operation: (selected: { keyring: any }) => Promise<unknown>,
+      ) => Promise<unknown>
+    >;
   };
 };
 
@@ -76,6 +84,7 @@ function getMessenger(
       'SnapController:getRunnableSnaps',
       'KeyringController:getState',
       'KeyringController:withController',
+      'KeyringController:withKeyringV2Unsafe',
     ],
     events: [
       'SnapController:stateChange',
@@ -253,6 +262,7 @@ function setup({
             removeKeyring: jest.fn().mockResolvedValue(undefined),
           }),
         ),
+      withKeyringV2Unsafe: jest.fn(),
     },
   };
 
@@ -273,6 +283,11 @@ function setup({
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
     mocks.KeyringController.withController as any,
   );
+  rootMessenger.registerActionHandler(
+    'KeyringController:withKeyringV2Unsafe',
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    mocks.KeyringController.withKeyringV2Unsafe as any,
+  );
 
   const service = new SnapAccountService({ messenger, config });
 
@@ -690,4 +705,101 @@ describe('SnapAccountService', () => {
       expect(resolved).toBe(true);
     });
   });
+
+  describe('handleKeyringSnapMessage', () => {
+    it('returns an empty array for `getSelectedAccounts`', async () => {
+      const { service, mocks } = setup();
+
+      const result = await service.handleKeyringSnapMessage(MOCK_SNAP_ID, {
+        method: 'getSelectedAccounts',
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      } as any);
+
+      expect(result).toStrictEqual([]);
+      expect(mocks.KeyringController.withController).not.toHaveBeenCalled();
+      expect(
+        mocks.KeyringController.withKeyringV2Unsafe,
+      ).not.toHaveBeenCalled();
+    });
+
+    it('throws for non-AccountCreated events', async () => {
+      const { service } = setup();
+
+      await expect(
+        service.handleKeyringSnapMessage(MOCK_SNAP_ID, {
+          method: 'accountUpdated',
+          // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        } as any),
+      ).rejects.toThrow(
+        'Cannot delegate keyring Snap message, keyring does not exist yet.',
+      );
+    });
+
+    it('auto-creates a Snap keyring on AccountCreated when one does not exist, then delegates', async () => {
+      const addNewKeyring = jest.fn().mockResolvedValue(undefined);
+      const removeKeyring = jest.fn().mockResolvedValue(undefined);
+      const handleSnapMessage = jest.fn().mockResolvedValue('ok');
+      const { service, mocks } = setup();
+      mocks.KeyringController.withController.mockImplementation(
+        async (callback) =>
+          callback({ keyrings: [], addNewKeyring, removeKeyring }),
+      );
+      mocks.KeyringController.withKeyringV2Unsafe.mockImplementation(
+        async (_selector, operation) =>
+          operation({
+            keyring: { handleKeyringSnapMessage: handleSnapMessage },
+          }),
+      );
+
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      const message = { method: 'notify:accountCreated' } as any;
+      const result = await service.handleKeyringSnapMessage(
+        MOCK_SNAP_ID,
+        message,
+      );
+
+      expect(addNewKeyring).toHaveBeenCalledWith(KeyringType.Snap, {
+        snapId: MOCK_SNAP_ID,
+        accounts: {},
+      });
+      expect(handleSnapMessage).toHaveBeenCalledWith(message);
+      expect(result).toBe('ok');
+    });
+
+    it('does not re-create the Snap keyring on AccountCreated when one already exists', async () => {
+      const addNewKeyring = jest.fn().mockResolvedValue(undefined);
+      const removeKeyring = jest.fn().mockResolvedValue(undefined);
+      const handleSnapMessage = jest.fn().mockResolvedValue('ok');
+      const { service, mocks } = setup();
+      mocks.KeyringController.withController.mockImplementation(
+        async (callback) =>
+          callback({
+            keyrings: [
+              {
+                keyringV2: {
+                  type: KeyringType.Snap,
+                  snapId: MOCK_SNAP_ID,
+                },
+              },
+            ],
+            addNewKeyring,
+            removeKeyring,
+          }),
+      );
+      mocks.KeyringController.withKeyringV2Unsafe.mockImplementation(
+        async (_selector, operation) =>
+          operation({
+            keyring: { handleKeyringSnapMessage: handleSnapMessage },
+          }),
+      );
+
+      await service.handleKeyringSnapMessage(MOCK_SNAP_ID, {
+        method: 'notify:accountCreated',
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      } as any);
+
+      expect(addNewKeyring).not.toHaveBeenCalled();
+      expect(handleSnapMessage).toHaveBeenCalled();
+    });
+  });
 });

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

@@ -6,12 +6,15 @@ import type {
   SnapKeyring,
   SnapKeyringState,
 } from '@metamask/eth-snap-keyring/v2';
+import type { SnapMessage } from '@metamask/eth-snap-keyring';
+import { KeyringEvent } from '@metamask/keyring-api';
 import type { Keyring } from '@metamask/keyring-api/v2';
 import { KeyringType } from '@metamask/keyring-api/v2';
 import type {
   KeyringControllerGetStateAction,
   KeyringControllerStateChangeEvent,
   KeyringControllerWithControllerAction,
+  KeyringControllerWithKeyringV2UnsafeAction,
 } from '@metamask/keyring-controller';
 import type { Messenger } from '@metamask/messenger';
 import type {
@@ -24,13 +27,15 @@ import type {
   SnapControllerSnapUninstalledEvent,
   SnapControllerStateChangeEvent,
 } from '@metamask/snaps-controllers';
+import type { Json } from '@metamask/snaps-sdk';
 import { SnapId } from '@metamask/snaps-sdk';
 import type { TruncatedSnap } from '@metamask/snaps-utils';
 
 import { projectLogger as log } from './logger';
 import type {
   SnapAccountServiceEnsureReadyAction,
   SnapAccountServiceGetSnapsAction,
+  SnapAccountServiceHandleKeyringSnapMessageAction,
 } from './SnapAccountService-method-action-types';
 import { SnapPlatformWatcher } from './SnapPlatformWatcher';
 import type { SnapPlatformWatcherConfig } from './SnapPlatformWatcher';
@@ -67,14 +72,19 @@ export const serviceName = 'SnapAccountService';
  * All of the methods within {@link SnapAccountService} that are exposed via
  * the messenger.
  */
-const MESSENGER_EXPOSED_METHODS = ['ensureReady', 'getSnaps'] as const;
+const MESSENGER_EXPOSED_METHODS = [
+  'ensureReady',
+  'getSnaps',
+  'handleKeyringSnapMessage',
+] as const;
 
 /**
  * Actions that {@link SnapAccountService} exposes to other consumers.
  */
 export type SnapAccountServiceActions =
   | SnapAccountServiceEnsureReadyAction
-  | SnapAccountServiceGetSnapsAction;
+  | SnapAccountServiceGetSnapsAction
+  | SnapAccountServiceHandleKeyringSnapMessageAction;
 
 /**
  * Actions from other messengers that {@link SnapAccountService} calls.
@@ -83,7 +93,8 @@ type AllowedActions =
   | SnapControllerGetStateAction
   | SnapControllerGetRunnableSnapsAction
   | KeyringControllerGetStateAction
-  | KeyringControllerWithControllerAction;
+  | KeyringControllerWithControllerAction
+  | KeyringControllerWithKeyringV2UnsafeAction;
 
 /**
  * Events that {@link SnapAccountService} exposes to other consumers.
@@ -352,6 +363,89 @@ export class SnapAccountService {
     await this.#watcher.ensureCanUseSnapPlatform();
   }
 
+  /**
+   * Handle a message from a Snap.
+   *
+   * Only `AccountCreated` triggers lazy keyring creation, since that is the
+   * single entry point for the v1 event-driven flow. All other events from
+   * unknown Snaps throw an error.
+   *
+   * @param snapId - ID of the Snap.
+   * @param message - Message sent by the Snap.
+   * @returns The execution result.
+   */
+  async handleKeyringSnapMessage(
+    snapId: SnapId,
+    message: SnapMessage,
+  ): Promise<Json> {
+    // We assume the Snap platform always sends a valid `KeyringEvent` here.
+    const event = message.method as KeyringEvent;
+
+    if (message.method === 'getSelectedAccounts') {
+      return [];
+    }
+
+    const isSnapKeyringForThisSnap = (
+      keyring: Keyring,
+    ): keyring is SnapKeyring =>
+      isSnapKeyring(keyring) && keyring.snapId === snapId;
+
+    // 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.#messenger.call(
+        'KeyringController:withController',
+        async (controller) => {
+          const hasSnapKeyring = controller.keyrings.some(
+            ({ keyringV2 }) => keyringV2 && isSnapKeyringForThisSnap(keyringV2),
+          );
+
+          if (!hasSnapKeyring) {
+            log(`Auto-creating Snap keyring for "${snapId}"...`);
+
+            await controller.addNewKeyring(KeyringType.Snap, {
+              snapId,
+              accounts: {},
+            });
+
+            log(`Snap keyring for "${snapId}" is ready!`);
+          }
+        },
+      );
+    } else {
+      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.',
+      );
+    }
+
+    // This part of the flow relies on v1 flows, so we have to go with
+    // "unsafe" to avoid deadlocks. The keyring persistence will be done in a
+    // later stage of `handleKeyringSnapMessage` after the message is handled,
+    // so we don't have to worry about that here.
+    const result = await this.#messenger.call(
+      'KeyringController:withKeyringV2Unsafe',
+      { filter: (keyring) => isSnapKeyringForThisSnap(keyring) },
+      async ({ keyring }) => {
+        const snapKeyring = keyring as SnapKeyring;
+        return snapKeyring.handleKeyringSnapMessage(message);
+      },
+    );
+
+    log(
+      `Snap keyring for "${snapId}" handled keyring message "${event}" successfully!`,
+    );
+
+    return result as Json;
+  }
+
   /**
    * Handles a Snap being added (installed or enabled). If the Snap is an
    * account-management Snap, adds it to the internal set of tracked Snaps.

packages/snap-account-service/src/index.ts

@@ -9,6 +9,7 @@ export type {
 export type {
   SnapAccountServiceEnsureReadyAction,
   SnapAccountServiceGetSnapsAction,
+  SnapAccountServiceHandleKeyringSnapMessageAction,
 } from './SnapAccountService-method-action-types';
 export { SnapPlatformWatcher } from './SnapPlatformWatcher';
 export type { SnapPlatformWatcherConfig } from './SnapPlatformWatcher';