refactor(account-tree-controller): force use of password for :exportState

packages/account-tree-controller/CHANGELOG.md

@@ -9,8 +9,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
-- Add `:{import,export}State` ([#9663](https://github.com/MetaMask/core/pull/9663)), ([#9826](https://github.com/MetaMask/core/pull/9826)), ([#9863](https://github.com/MetaMask/core/pull/9863)), ([#9864](https://github.com/MetaMask/core/pull/9864))
+- **BREAKING:** Add `:{import,export}State` ([#9663](https://github.com/MetaMask/core/pull/9663)), ([#9826](https://github.com/MetaMask/core/pull/9826)), ([#9863](https://github.com/MetaMask/core/pull/9863)), ([#9864](https://github.com/MetaMask/core/pull/9864)), ([#9883](https://github.com/MetaMask/core/pull/9883))
+  - The following actions must be registered on the controller's messenger: `MultichainAccountService:createMultichainAccountWallet`, `KeyringController:with{Controller,KeyringV2,KeyringV2Unsafe}`.
   - This can be used to serialize/deserialize the entire account-tree state (metadata + secrets if needed).
+  - The `password` is required whenever secrets are requested.
   - The payload is versionned and hard-coded to version 1 for now.
 - Add `AccountTreeController:initialized` event, emitted at the end of `init()` when the account tree is fully built and ready to consume ([#9880](https://github.com/MetaMask/core/pull/9880))
 - Add `AccountTreeController:uninitialized` event, emitted at the end of `clearState()` when the account tree has been torn down ([#9880](https://github.com/MetaMask/core/pull/9880))

packages/account-tree-controller/src/AccountTreeController-method-action-types.ts

@@ -236,11 +236,14 @@ export type AccountTreeControllerSyncWithUserStorageAtLeastOnceAction = {
 /**
  * Produces a versioned snapshot of the current wallet and group state.
  *
- * When `options.includeSecrets` is `true` and the vault is unlocked,
- * mnemonic phrases and private keys are included in the snapshot.
+ * When `options.includeSecrets` is `true`, `options.password` is required
+ * and verified against the vault before any secret is read. Without
+ * `includeSecrets`, only metadata (names, pinned, hidden) is exported and
+ * no password is needed.
  *
  * @param options - Export options.
  * @returns A promise resolving to an `AccountTreeSnapshot`.
+ * @throws If the vault is locked or the password is incorrect.
  */
 export type AccountTreeControllerExportStateAction = {
   type: `AccountTreeController:exportState`;

packages/account-tree-controller/src/AccountTreeController.test.ts

@@ -322,6 +322,7 @@ function setup({
     KeyringController: {
       keyrings: KeyringObject[];
       getState: jest.Mock;
+      verifyPassword: jest.Mock;
       withController: jest.Mock;
     };
     AccountsController: {
@@ -346,6 +347,7 @@ function setup({
     KeyringController: {
       keyrings,
       getState: jest.fn(),
+      verifyPassword: jest.fn().mockResolvedValue(undefined),
       withController: jest.fn(),
     },
     AccountsController: {
@@ -449,6 +451,11 @@ function setup({
       mocks.KeyringController.getState,
     );
 
+    messenger.registerActionHandler(
+      'KeyringController:verifyPassword',
+      mocks.KeyringController.verifyPassword,
+    );
+
     // Default: call the callback with no existing keyrings so private-key
     // imports are a no-op unless the test overrides this handler.
     mocks.KeyringController.withController.mockImplementation(
@@ -6682,8 +6689,64 @@ describe('AccountTreeController', () => {
       ).rejects.toThrow('Cannot export account tree when vault is locked');
 
       await expect(
-        controller.exportState({ includeSecrets: true }),
+        controller.exportState({
+          includeSecrets: true,
+          password: 'test-password',
+        }),
       ).rejects.toThrow('Cannot export account tree when vault is locked');
     });
+
+    it('throws when exporting with a wrong password', async () => {
+      const { controller, mocks } = setup({
+        accounts: [MOCK_HD_ACCOUNT_1],
+        keyrings: [MOCK_HD_KEYRING_1],
+      });
+
+      controller.init();
+
+      mocks.KeyringController.verifyPassword.mockRejectedValue(
+        new Error('Invalid password'),
+      );
+
+      await expect(
+        controller.exportState({
+          includeSecrets: true,
+          password: 'wrong-password',
+        }),
+      ).rejects.toThrow('Invalid password');
+    });
+
+    it('verifies password before exporting', async () => {
+      const { controller, messenger, mocks } = setup({
+        accounts: [MOCK_HD_ACCOUNT_1],
+        keyrings: [MOCK_HD_KEYRING_1],
+      });
+
+      controller.init();
+
+      messenger.registerActionHandler(
+        'KeyringController:withKeyringV2Unsafe',
+        async (
+          _selector: unknown,
+          callback: (ctx: { keyring: unknown }) => unknown,
+        ) =>
+          callback({
+            keyring: {
+              toEntropySourceId: async () => MOCK_HD_KEYRING_1.metadata.id,
+              // Must be even-length for encodeMnemonic (uses Uint16Array internally).
+              mnemonic: new Uint8Array([1, 2, 3, 4]),
+            },
+          }),
+      );
+
+      await controller.exportState({
+        includeSecrets: true,
+        password: 'correct-password',
+      });
+
+      expect(mocks.KeyringController.verifyPassword).toHaveBeenCalledWith(
+        'correct-password',
+      );
+    });
   });
 });

packages/account-tree-controller/src/AccountTreeController.ts

@@ -1881,11 +1881,14 @@ export class AccountTreeController extends BaseController<
   /**
    * Produces a versioned snapshot of the current wallet and group state.
    *
-   * When `options.includeSecrets` is `true` and the vault is unlocked,
-   * mnemonic phrases and private keys are included in the snapshot.
+   * When `options.includeSecrets` is `true`, `options.password` is required
+   * and verified against the vault before any secret is read. Without
+   * `includeSecrets`, only metadata (names, pinned, hidden) is exported and
+   * no password is needed.
    *
    * @param options - Export options.
    * @returns A promise resolving to an `AccountTreeSnapshot`.
+   * @throws If the vault is locked or the password is incorrect.
    */
   async exportState(
     options?: ExportStateOptions,

packages/account-tree-controller/src/state/export.test.ts

@@ -85,6 +85,7 @@ function setup({
   mocks: {
     KeyringController: {
       getState: jest.Mock;
+      verifyPassword: jest.Mock;
       withKeyringV2Unsafe: jest.Mock;
       withKeyringV2: jest.Mock;
     };
@@ -96,6 +97,7 @@ function setup({
   const mocks = {
     KeyringController: {
       getState: jest.fn().mockReturnValue({ isUnlocked, keyrings: [] }),
+      verifyPassword: jest.fn().mockResolvedValue(undefined),
       withKeyringV2Unsafe: jest.fn(),
       withKeyringV2: jest.fn(),
     },
@@ -109,6 +111,8 @@ function setup({
       switch (action) {
         case 'KeyringController:getState':
           return mocks.KeyringController.getState();
+        case 'KeyringController:verifyPassword':
+          return mocks.KeyringController.verifyPassword(...args);
         case 'KeyringController:withKeyringV2Unsafe':
           return mocks.KeyringController.withKeyringV2Unsafe(...args);
         case 'KeyringController:withKeyringV2':
@@ -287,7 +291,10 @@ describe('exportState', () => {
         new Uint8Array([1, 2, 3, 4]),
       );
 
-      const snapshot = await exportState(context, { includeSecrets: true });
+      const snapshot = await exportState(context, {
+        includeSecrets: true,
+        password: 'test-password',
+      });
       const wallet = snapshot.serialize().wallets[0] as { value?: string };
 
       expect(Array.isArray(wallet.value)).toBe(true);
@@ -302,7 +309,10 @@ describe('exportState', () => {
       );
 
       await expect(
-        exportState(context, { includeSecrets: true }),
+        exportState(context, {
+          includeSecrets: true,
+          password: 'test-password',
+        }),
       ).rejects.toThrow('Failed to export mnemonic');
     });
 
@@ -469,7 +479,10 @@ describe('exportState', () => {
         encoding: AccountWalletPrivateKeyEncoding.Hexadecimal,
       });
 
-      const snapshot = await exportState(context, { includeSecrets: true });
+      const snapshot = await exportState(context, {
+        includeSecrets: true,
+        password: 'test-password',
+      });
       const group = snapshot.serialize().wallets[0]?.groups[0] as {
         value?: { privateKey: number[]; encoding: string; type: string };
       };
@@ -499,7 +512,10 @@ describe('exportState', () => {
       );
 
       await expect(
-        exportState(context, { includeSecrets: true }),
+        exportState(context, {
+          includeSecrets: true,
+          password: 'test-password',
+        }),
       ).rejects.toThrow('does not support exportAccount');
     });
 
@@ -515,7 +531,10 @@ describe('exportState', () => {
         makePrivateKeyExportHandler(undefined);
 
       await expect(
-        exportState(context, { includeSecrets: true }),
+        exportState(context, {
+          includeSecrets: true,
+          password: 'test-password',
+        }),
       ).rejects.toThrow('Failed to export private key');
     });
 

packages/account-tree-controller/src/state/export.ts

@@ -277,12 +277,26 @@ export async function exportState(
 ): Promise<AccountTreeSnapshot> {
   const state = context.getState();
 
-  const includeSecrets = options.includeSecrets ?? false;
   const { isUnlocked } = context.messenger.call('KeyringController:getState');
   if (!isUnlocked) {
     throw new Error('Cannot export account tree when vault is locked');
   }
 
+  // Use `options` here to let the compiler infer the type of `includeSecrets` based on the
+  // discriminated union.
+  if (options.includeSecrets) {
+    // We verify the password here to force consumers to have it in their flow
+    // before calling exportState. The password is never stored, so the only
+    // way to supply it is to ask the user, ensuring they are prompted upstream
+    // rather than having the export silently succeed without their interaction.
+    await context.messenger.call(
+      'KeyringController:verifyPassword',
+      options.password,
+    );
+  }
+
+  const includeSecrets = options.includeSecrets ?? false;
+
   const idMap = new IdMap();
   const entries: AccountTreeWalletEntry[] = [];
   let privateKeyWallet: AccountWalletPrivateKeyPayload | undefined;

packages/account-tree-controller/src/state/payload.ts

@@ -209,10 +209,18 @@ export function toGroupPayloadId(
 }
 
 /** Options accepted by {@link AccountTreeController.exportState}. */
-export type ExportStateOptions = {
-  /** When `true`, secrets (mnemonic / private keys) are included in the export. */
-  includeSecrets?: boolean;
-};
+export type ExportStateOptions =
+  | {
+      /** When `true`, secrets (mnemonic / private keys) are included in the snapshot. */
+      includeSecrets: true;
+      /** Password verified against the vault before any secret is read. */
+      password: string;
+    }
+  | {
+      /** When `false` or omitted, only metadata is exported — no password needed. */
+      includeSecrets?: false;
+      password?: never;
+    };
 
 const AccountWalletPayloadIdStruct = define<AccountWalletPayloadId>(
   'AccountWalletPayloadId',

packages/account-tree-controller/src/types.ts

@@ -16,6 +16,7 @@ import type {
 import type { TraceCallback } from '@metamask/controller-utils';
 import type {
   KeyringControllerGetStateAction,
+  KeyringControllerVerifyPasswordAction,
   KeyringControllerWithControllerAction,
   KeyringControllerWithKeyringV2Action,
   KeyringControllerWithKeyringV2UnsafeAction,
@@ -92,6 +93,7 @@ export type AllowedActions =
   | AccountsControllerListMultichainAccountsAction
   | AccountsControllerSetSelectedAccountAction
   | KeyringControllerGetStateAction
+  | KeyringControllerVerifyPasswordAction
   | SnapControllerGetSnapAction
   | UserStorageController.UserStorageControllerGetStateAction
   | UserStorageController.UserStorageControllerPerformGetStorageAction

packages/account-tree-controller/tests/mockMessenger.ts

@@ -65,6 +65,7 @@ export function getAccountTreeControllerMessenger(
       'MultichainAccountService:createMultichainAccountGroups',
       'MultichainAccountService:createMultichainAccountWallet',
       'KeyringController:getState',
+      'KeyringController:verifyPassword',
       'KeyringController:withController',
       'KeyringController:withKeyringV2',
       'KeyringController:withKeyringV2Unsafe',