feat(keyring-controller): persist vault when keyring state changes during unlock

packages/keyring-controller/CHANGELOG.md

@@ -10,6 +10,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 
 - Expose `KeyringController:signTransaction` method through `KeyringController` messenger ([#8408](https://github.com/MetaMask/core/pull/8408))
+- Persist vault when keyring state changes during unlock ([#8415](https://github.com/MetaMask/core/pull/8415))
+  - If a keyring's serialized state differs after deserialization (e.g. a migration ran, or metadata was missing), the vault is now re-persisted so the change is not lost on the next unlock.
 
 ### Changed
 

packages/keyring-controller/src/KeyringController.test.ts

@@ -2945,6 +2945,81 @@ describe('KeyringController', () => {
       );
     });
 
+    it('should update the vault if the keyring state changes during deserialization', async () => {
+      const oldState = { version: 1, foo: 'bar' };
+      const newState = { version: 2, foo: 'bar' };
+
+      // A keyring that migrates its own state in-place during deserialization: after
+      // deserialize(oldState), the original `data` object is mutated and serialize()
+      // returns the updated state.
+      class MigratingKeyring {
+        static type = 'Migrating Keyring';
+
+        type = 'Migrating Keyring';
+
+        #state: Record<string, unknown> = {};
+
+        async serialize(): Promise<Record<string, unknown>> {
+          return this.#state;
+        }
+
+        async deserialize(data: Record<string, unknown>): Promise<void> {
+          // Mutate in-place to simulate a keyring that upgrades its own format.
+          data.version = 2;
+          this.#state = data;
+        }
+
+        async getAccounts(): Promise<string[]> {
+          return [];
+        }
+
+        async addAccounts(): Promise<string[]> {
+          return [];
+        }
+      }
+
+      await withController(
+        {
+          skipVaultCreation: true,
+          state: {
+            vault: createVault([
+              // #updateVault requires at least one HD keyring to be present.
+              {
+                type: KeyringTypes.hd,
+                data: {},
+                metadata: { id: '0', name: '' },
+              },
+              {
+                type: MigratingKeyring.type,
+                data: oldState,
+                metadata: { id: '1', name: '' },
+              },
+            ]),
+          },
+          keyringBuilders: [
+            keyringBuilderFactory(MigratingKeyring as unknown as KeyringClass),
+          ],
+        },
+        async ({ controller, encryptor }) => {
+          const encryptWithKeySpy = jest.spyOn(encryptor, 'encryptWithKey');
+
+          await controller.submitPassword(password);
+
+          // Migration should have triggered a new vault update that we need to
+          // re-encrypt:
+          expect(encryptWithKeySpy).toHaveBeenCalledWith(
+            defaultCredentials,
+            expect.arrayContaining([
+              expect.objectContaining({
+                type: MigratingKeyring.type,
+                data: newState,
+              }),
+            ]),
+          );
+        },
+      );
+    });
+
     it('should unlock the wallet if the state has a duplicate account and the encryption parameters are outdated', async () => {
       stubKeyringClassWithAccount(MockKeyring, '0x123');
       stubKeyringClassWithAccount(HdKeyring, '0x123');

packages/keyring-controller/src/KeyringController.ts

@@ -32,7 +32,7 @@ import { Mutex } from 'async-mutex';
 import type { MutexInterface } from 'async-mutex';
 import Wallet, { thirdparty as importers } from 'ethereumjs-wallet';
 import type { Patch } from 'immer';
-import { cloneDeep, isEqual } from 'lodash';
+import { cloneDeep } from 'lodash';
 // When generating a ULID within the same millisecond, monotonicFactory provides some guarantees regarding sort order.
 import { ulid } from 'ulid';
 
@@ -1526,7 +1526,7 @@ export class KeyringController<
     encryptionKey: string,
     encryptionSalt?: string,
   ): Promise<void> {
-    const { newMetadata } = await this.#withRollback(async () => {
+    const { hasChanged } = await this.#withRollback(async () => {
       const result = await this.#unlockKeyrings({
         encryptionKey,
         encryptionSalt,
@@ -1539,7 +1539,7 @@ export class KeyringController<
       // if new metadata has been generated during login, we
       // can attempt to upgrade the vault.
       await this.#withRollback(async () => {
-        if (newMetadata) {
+        if (hasChanged) {
           await this.#updateVault();
         }
       });
@@ -1572,18 +1572,18 @@ export class KeyringController<
    * @returns Promise resolving when the operation completes.
    */
   async submitPassword(password: string): Promise<void> {
-    const { newMetadata } = await this.#withRollback(async () => {
+    const { hasChanged } = await this.#withRollback(async () => {
       const result = await this.#unlockKeyrings({ password });
       this.#setUnlocked();
       return result;
     });
 
     try {
       // If there are stronger encryption params available, or
-      // if new metadata has been generated during login, we
+      // if the keyring state has changed during deserialization, we
       // can attempt to upgrade the vault.
       await this.#withRollback(async () => {
-        if (newMetadata || this.#isNewEncryptionAvailable()) {
+        if (hasChanged || this.#isNewEncryptionAvailable()) {
           await this.#deriveAndSetEncryptionKey(password, {
             // If the vault is being upgraded, we want to ignore the metadata
             // that is already in the vault, so we can effectively
@@ -2179,24 +2179,24 @@ export class KeyringController<
     serializedKeyrings: SerializedKeyring[],
   ): Promise<{
     keyrings: { keyring: EthKeyring; metadata: KeyringMetadata }[];
-    newMetadata: boolean;
+    hasChanged: boolean;
   }> {
     await this.#clearKeyrings();
     const keyrings: { keyring: EthKeyring; metadata: KeyringMetadata }[] = [];
-    let newMetadata = false;
+    let hasChanged = false;
 
     for (const serializedKeyring of serializedKeyrings) {
       const result = await this.#restoreKeyring(serializedKeyring);
       if (result) {
         const { keyring, metadata } = result;
         keyrings.push({ keyring, metadata });
-        if (result.newMetadata) {
-          newMetadata = true;
+        if (result.hasChanged) {
+          hasChanged = true;
         }
       }
     }
 
-    return { keyrings, newMetadata };
+    return { keyrings, hasChanged };
   }
 
   /**
@@ -2217,7 +2217,7 @@ export class KeyringController<
         },
   ): Promise<{
     keyrings: { keyring: EthKeyring; metadata: KeyringMetadata }[];
-    newMetadata: boolean;
+    hasChanged: boolean;
   }> {
     return this.#withVaultLock(async () => {
       if (!this.state.vault) {
@@ -2255,7 +2255,7 @@ export class KeyringController<
         );
       }
 
-      const { keyrings, newMetadata } =
+      const { keyrings, hasChanged } =
         await this.#restoreSerializedKeyrings(vault);
 
       const updatedKeyrings = await this.#getUpdatedKeyrings();
@@ -2266,7 +2266,7 @@ export class KeyringController<
         state.encryptionSalt = this.#encryptionKey?.salt;
       });
 
-      return { keyrings, newMetadata };
+      return { keyrings, hasChanged };
     });
   }
 
@@ -2496,30 +2496,37 @@ export class KeyringController<
   async #restoreKeyring(
     serialized: SerializedKeyring,
   ): Promise<
-    | { keyring: EthKeyring; metadata: KeyringMetadata; newMetadata: boolean }
+    | { keyring: EthKeyring; metadata: KeyringMetadata; hasChanged: boolean }
     | undefined
   > {
     this.#assertControllerMutexIsLocked();
 
     try {
       const { type, data, metadata: serializedMetadata } = serialized;
-      let newMetadata = false;
-      let metadata = serializedMetadata;
+
+      const oldState = JSON.stringify(data);
       const keyring = await this.#createKeyring(type, data);
+      const newState = JSON.stringify(await keyring.serialize());
+      let hasChanged = oldState !== newState;
+
       await this.#assertNoDuplicateAccounts([keyring]);
-      // If metadata is missing, assume the data is from an installation before
-      // we had keyring metadata.
+
+      // If metadata is missing, assume the data is from an installation before we had
+      // keyring metadata.
+      let metadata = serializedMetadata;
       if (!metadata) {
-        newMetadata = true;
+        hasChanged = true;
         metadata = getDefaultKeyringMetadata();
       }
+
       // The keyring is added to the keyrings array only if it's successfully restored
       // and the metadata is successfully added to the controller
       this.#keyrings.push({
         keyring,
         metadata,
       });
-      return { keyring, metadata, newMetadata };
+
+      return { keyring, metadata, hasChanged };
     } catch (error) {
       console.error(error);
       this.#unsupportedKeyrings.push(serialized);
@@ -2603,7 +2610,7 @@ export class KeyringController<
       const newState = JSON.stringify(await this.#getSessionState());
 
       // State is committed only if the operation is successful and need to trigger a vault update.
-      if (!isEqual(oldState, newState)) {
+      if (oldState !== newState) {
         await this.#updateVault();
       }