[15.x] Backport encryption improvements

jest.config.js

@@ -41,7 +41,7 @@ module.exports = {
   // An object that configures minimum threshold enforcement for coverage results
   coverageThreshold: {
     global: {
-      branches: 79.41,
+      branches: 79.8,
       functions: 93.22,
       lines: 91.5,
       statements: 91.69,

package.json

@@ -44,7 +44,7 @@
   },
   "dependencies": {
     "@ethereumjs/tx": "^4.2.0",
-    "@metamask/browser-passworder": "^4.2.0",
+    "@metamask/browser-passworder": "^4.3.0",
     "@metamask/eth-hd-keyring": "^7.0.1",
     "@metamask/eth-sig-util": "^7.0.0",
     "@metamask/eth-simple-keyring": "^6.0.1",

src/KeyringController.test.ts

@@ -78,6 +78,18 @@ async function initializeKeyringController({
   return keyringController;
 }
 
+/**
+ * Delete the encryption key and salt from the `memStore` of the given keyring controller.
+ *
+ * @param keyringController - The keyring controller to delete the encryption key and salt from.
+ */
+function deleteEncryptionKeyAndSalt(keyringController: KeyringController) {
+  const keyringControllerState = keyringController.memStore.getState();
+  delete keyringControllerState.encryptionKey;
+  delete keyringControllerState.encryptionSalt;
+  keyringController.memStore.updateState(keyringControllerState);
+}
+
 describe('KeyringController', () => {
   afterEach(() => {
     sinon.restore();
@@ -199,83 +211,86 @@ describe('KeyringController', () => {
     });
 
     describe('when `cacheEncryptionKey` is enabled', () => {
-      it('should save an up to date encryption salt to the `memStore` when `password` is unset and `encryptionKey` is set', async () => {
-        const keyringController = await initializeKeyringController({
-          password: PASSWORD,
-          constructorOptions: {
-            cacheEncryptionKey: true,
-          },
+      describe('when `encryptionKey` is set', () => {
+        it('should save an up to date encryption salt to the `memStore`', async () => {
+          const keyringController = await initializeKeyringController({
+            password: PASSWORD,
+            constructorOptions: {
+              cacheEncryptionKey: true,
+            },
+          });
+          const vaultEncryptionKey = '🔑';
+          const vaultEncryptionSalt = '🧂';
+          const vault = JSON.stringify({ salt: vaultEncryptionSalt });
+          keyringController.store.updateState({ vault });
+
+          await keyringController.unlockKeyrings(
+            undefined,
+            vaultEncryptionKey,
+            vaultEncryptionSalt,
+          );
+
+          expect(keyringController.memStore.getState().encryptionKey).toBe(
+            vaultEncryptionKey,
+          );
+          expect(keyringController.memStore.getState().encryptionSalt).toBe(
+            vaultEncryptionSalt,
+          );
+
+          const response = await keyringController.persistAllKeyrings();
+
+          expect(response).toBe(true);
+          expect(keyringController.memStore.getState().encryptionKey).toBe(
+            vaultEncryptionKey,
+          );
+          expect(keyringController.memStore.getState().encryptionSalt).toBe(
+            vaultEncryptionSalt,
+          );
         });
-        delete keyringController.password;
-        const vaultEncryptionKey = '🔑';
-        const vaultEncryptionSalt = '🧂';
-        const vault = JSON.stringify({ salt: vaultEncryptionSalt });
-        keyringController.store.updateState({ vault });
-
-        await keyringController.unlockKeyrings(
-          undefined,
-          vaultEncryptionKey,
-          vaultEncryptionSalt,
-        );
-
-        expect(keyringController.memStore.getState().encryptionKey).toBe(
-          vaultEncryptionKey,
-        );
-        expect(keyringController.memStore.getState().encryptionSalt).toBe(
-          vaultEncryptionSalt,
-        );
-
-        const response = await keyringController.persistAllKeyrings();
-
-        expect(response).toBe(true);
-        expect(keyringController.memStore.getState().encryptionKey).toBe(
-          vaultEncryptionKey,
-        );
-        expect(keyringController.memStore.getState().encryptionSalt).toBe(
-          vaultEncryptionSalt,
-        );
       });
 
-      it('should save an up to date encryption salt to the `memStore` when `password` is set through `createNewVaultAndKeychain`', async () => {
-        const keyringController = await initializeKeyringController({
-          password: PASSWORD,
-          constructorOptions: {
-            cacheEncryptionKey: true,
-          },
+      describe('when `encryptionKey` is not set and `password` is set', () => {
+        it('should save an up to date encryption salt to the `memStore` when `password` is set through `createNewVaultAndKeychain`', async () => {
+          const keyringController = await initializeKeyringController({
+            password: PASSWORD,
+            constructorOptions: {
+              cacheEncryptionKey: true,
+            },
+          });
+          await keyringController.createNewVaultAndKeychain(PASSWORD);
+          deleteEncryptionKeyAndSalt(keyringController);
+
+          const response = await keyringController.persistAllKeyrings();
+
+          expect(response).toBe(true);
+          expect(keyringController.memStore.getState().encryptionKey).toBe(
+            MOCK_HARDCODED_KEY,
+          );
+          expect(keyringController.memStore.getState().encryptionSalt).toBe(
+            MOCK_ENCRYPTION_SALT,
+          );
         });
 
-        await keyringController.createNewVaultAndKeychain(PASSWORD);
-
-        const response = await keyringController.persistAllKeyrings();
-
-        expect(response).toBe(true);
-        expect(keyringController.memStore.getState().encryptionKey).toBe(
-          MOCK_HARDCODED_KEY,
-        );
-        expect(keyringController.memStore.getState().encryptionSalt).toBe(
-          MOCK_ENCRYPTION_SALT,
-        );
-      });
-
-      it('should save an up to date encryption salt to the `memStore` when `password` is set through `submitPassword`', async () => {
-        const keyringController = await initializeKeyringController({
-          password: PASSWORD,
-          constructorOptions: {
-            cacheEncryptionKey: true,
-          },
+        it('should save an up to date encryption salt to the `memStore` when `password` is set through `submitPassword`', async () => {
+          const keyringController = await initializeKeyringController({
+            password: PASSWORD,
+            constructorOptions: {
+              cacheEncryptionKey: true,
+            },
+          });
+          await keyringController.submitPassword(PASSWORD);
+          deleteEncryptionKeyAndSalt(keyringController);
+
+          const response = await keyringController.persistAllKeyrings();
+
+          expect(response).toBe(true);
+          expect(keyringController.memStore.getState().encryptionKey).toBe(
+            MOCK_HARDCODED_KEY,
+          );
+          expect(keyringController.memStore.getState().encryptionSalt).toBe(
+            MOCK_ENCRYPTION_SALT,
+          );
         });
-
-        await keyringController.submitPassword(PASSWORD);
-
-        const response = await keyringController.persistAllKeyrings();
-
-        expect(response).toBe(true);
-        expect(keyringController.memStore.getState().encryptionKey).toBe(
-          MOCK_HARDCODED_KEY,
-        );
-        expect(keyringController.memStore.getState().encryptionSalt).toBe(
-          MOCK_ENCRYPTION_SALT,
-        );
       });
     });
   });
@@ -836,29 +851,81 @@ describe('KeyringController', () => {
     });
 
     describe('with old vault format', () => {
-      [true, false].forEach((cacheEncryptionKey) => {
-        describe(`with cacheEncryptionKey = ${cacheEncryptionKey}`, () => {
-          it('should update the vault', async () => {
-            const mockEncryptor = new MockEncryptor();
-            const keyringController = await initializeKeyringController({
-              password: PASSWORD,
-              constructorOptions: {
-                cacheEncryptionKey: true,
-                encryptor: mockEncryptor,
-              },
-            });
-            const initialVault = keyringController.store.getState().vault;
-            const updatedVaultMock =
-              '{"vault": "updated_vault_detail", "salt": "salt"}';
-            sinon.stub(mockEncryptor, 'updateVault').resolves(updatedVaultMock);
-            sinon.stub(mockEncryptor, 'encrypt').resolves(updatedVaultMock);
-
-            await keyringController.unlockKeyrings(PASSWORD);
-            const updatedVault = keyringController.store.getState().vault;
-
-            expect(initialVault).not.toBe(updatedVault);
-            expect(updatedVault).toBe(updatedVaultMock);
+      describe(`with cacheEncryptionKey = true and encryptionKey is unset`, () => {
+        it('should update the vault', async () => {
+          const mockEncryptor = new MockEncryptor();
+          const keyringController = await initializeKeyringController({
+            password: PASSWORD,
+            constructorOptions: {
+              cacheEncryptionKey: true,
+              encryptor: mockEncryptor,
+            },
           });
+          deleteEncryptionKeyAndSalt(keyringController);
+          const initialVault = keyringController.store.getState().vault;
+          const mockEncryptionResult = {
+            data: '0x1234',
+            iv: 'an iv',
+          };
+          sinon.stub(mockEncryptor, 'isVaultUpdated').returns(false);
+          sinon
+            .stub(mockEncryptor, 'encryptWithKey')
+            .resolves(mockEncryptionResult);
+
+          await keyringController.unlockKeyrings(PASSWORD);
+          const updatedVault = keyringController.store.getState().vault;
+
+          expect(initialVault).not.toBe(updatedVault);
+          expect(updatedVault).toBe(
+            JSON.stringify({
+              ...mockEncryptionResult,
+              salt: MOCK_ENCRYPTION_SALT,
+            }),
+          );
+        });
+      });
+
+      describe(`with cacheEncryptionKey = true and encryptionKey is set`, () => {
+        it('should not update the vault', async () => {
+          const mockEncryptor = new MockEncryptor();
+          const keyringController = await initializeKeyringController({
+            password: PASSWORD,
+            constructorOptions: {
+              cacheEncryptionKey: true,
+              encryptor: mockEncryptor,
+            },
+          });
+          const initialVault = keyringController.store.getState().vault;
+          sinon.stub(mockEncryptor, 'isVaultUpdated').returns(false);
+
+          await keyringController.unlockKeyrings(PASSWORD);
+          const updatedVault = keyringController.store.getState().vault;
+
+          expect(initialVault).toBe(updatedVault);
+        });
+      });
+
+      describe(`with cacheEncryptionKey = false`, () => {
+        it('should update the vault', async () => {
+          const mockEncryptor = new MockEncryptor();
+          const keyringController = await initializeKeyringController({
+            password: PASSWORD,
+            constructorOptions: {
+              cacheEncryptionKey: false,
+              encryptor: mockEncryptor,
+            },
+          });
+          const initialVault = keyringController.store.getState().vault;
+          const updatedVaultMock =
+            '{"vault": "updated_vault_detail", "salt": "salt"}';
+          sinon.stub(mockEncryptor, 'isVaultUpdated').returns(false);
+          sinon.stub(mockEncryptor, 'encrypt').resolves(updatedVaultMock);
+
+          await keyringController.unlockKeyrings(PASSWORD);
+          const updatedVault = keyringController.store.getState().vault;
+
+          expect(initialVault).not.toBe(updatedVault);
+          expect(updatedVault).toBe(updatedVaultMock);
         });
       });
     });

src/KeyringController.ts

@@ -802,7 +802,15 @@ class KeyringController extends EventEmitter {
     if (this.#cacheEncryptionKey) {
       assertIsExportableKeyEncryptor(this.#encryptor);
 
-      if (this.password) {
+      if (encryptionKey) {
+        const key = await this.#encryptor.importKey(encryptionKey);
+        const vaultJSON = await this.#encryptor.encryptWithKey(
+          key,
+          serializedKeyrings,
+        );
+        vaultJSON.salt = encryptionSalt;
+        vault = JSON.stringify(vaultJSON);
+      } else if (this.password) {
         const { vault: newVault, exportedKeyString } =
           await this.#encryptor.encryptWithDetail(
             this.password,
@@ -811,14 +819,6 @@ class KeyringController extends EventEmitter {
 
         vault = newVault;
         newEncryptionKey = exportedKeyString;
-      } else if (encryptionKey) {
-        const key = await this.#encryptor.importKey(encryptionKey);
-        const vaultJSON = await this.#encryptor.encryptWithKey(
-          key,
-          serializedKeyrings,
-        );
-        vaultJSON.salt = encryptionSalt;
-        vault = JSON.stringify(vaultJSON);
       }
     } else {
       if (typeof this.password !== 'string') {
@@ -930,9 +930,9 @@ class KeyringController extends EventEmitter {
 
     if (
       this.password &&
-      this.#encryptor.updateVault &&
-      (await this.#encryptor.updateVault(encryptedVault, this.password)) !==
-        encryptedVault
+      (!this.#cacheEncryptionKey || !encryptionKey) &&
+      this.#encryptor.isVaultUpdated &&
+      !this.#encryptor.isVaultUpdated(encryptedVault)
     ) {
       // Re-encrypt the vault with safer method if one is available
       await this.persistAllKeyrings();

src/test/encryptor.mock.ts

@@ -81,6 +81,10 @@ export class MockEncryptor implements ExportableKeyEncryptor {
     return _vault;
   }
 
+  isVaultUpdated(_vault: string) {
+    return true;
+  }
+
   generateSalt() {
     return MOCK_ENCRYPTION_SALT;
   }

src/types.ts

@@ -2,6 +2,7 @@ import type {
   DetailedDecryptResult,
   DetailedEncryptionResult,
   EncryptionResult,
+  KeyDerivationOptions,
 } from '@metamask/browser-passworder';
 import type { Json, Keyring } from '@metamask/utils';
 
@@ -63,14 +64,17 @@ export type GenericEncryptor = {
    */
   decrypt: (password: string, encryptedString: string) => Promise<unknown>;
   /**
-   * Optional vault migration helper. Updates the provided vault, re-encrypting
-   * data with a safer algorithm if one is available.
+   * Optional vault migration helper. Checks if the provided vault is up to date
+   * with the desired encryption algorithm.
    *
-   * @param vault - The encrypted string to update.
-   * @param password - The password to decrypt the vault with.
+   * @param vault - The encrypted string to check.
+   * @param targetDerivationParams - The desired target derivation params.
    * @returns The updated encrypted string.
    */
-  updateVault?: (vault: string, password: string) => Promise<string>;
+  isVaultUpdated?: (
+    vault: string,
+    targetDerivationParams?: KeyDerivationOptions,
+  ) => boolean;
 };
 
 /**