Add tests to make sure custom client cryptography is used

Author: Mrtenz

packages/snaps-controllers/src/snaps/SnapController.test.tsx

@@ -55,6 +55,7 @@ import {
 } from '@metamask/snaps-utils/test-utils';
 import type { SemVerRange, SemVerVersion, Json } from '@metamask/utils';
 import {
+  hexToBytes,
   assert,
   AssertionError,
   base64ToBytes,
@@ -8959,6 +8960,38 @@ describe('SnapController', () => {
 
       snapController.destroy();
     });
+
+    it('uses custom client cryptography functions', async () => {
+      const messenger = getSnapControllerMessenger();
+
+      const pbkdf2Sha512 = jest
+        .fn()
+        .mockResolvedValue(hexToBytes(ENCRYPTION_KEY));
+
+      const snapController = getSnapController(
+        getSnapControllerOptions({
+          messenger,
+          state: {
+            snaps: getPersistedSnapsState(),
+          },
+          clientCryptography: {
+            pbkdf2Sha512,
+          },
+        }),
+      );
+
+      const state = { foo: 'bar' };
+      await messenger.call(
+        'SnapController:updateSnapState',
+        MOCK_SNAP_ID,
+        state,
+        true,
+      );
+
+      expect(pbkdf2Sha512).toHaveBeenCalledTimes(1);
+
+      snapController.destroy();
+    });
   });
 
   describe('SnapController:clearSnapState', () => {

packages/snaps-rpc-methods/src/restricted/getBip32Entropy.test.ts

@@ -192,5 +192,42 @@ describe('getBip32EntropyImplementation', () => {
         }
       `);
     });
+
+    it('uses custom client cryptography functions', async () => {
+      const getUnlockPromise = jest.fn().mockResolvedValue(undefined);
+      const getMnemonic = jest
+        .fn()
+        .mockResolvedValue(TEST_SECRET_RECOVERY_PHRASE_BYTES);
+
+      const pbkdf2Sha512 = jest.fn().mockResolvedValue(new Uint8Array(64));
+      const getClientCryptography = jest.fn().mockReturnValue({
+        pbkdf2Sha512,
+      });
+
+      expect(
+        await getBip32EntropyImplementation({
+          getUnlockPromise,
+          getMnemonic,
+          getClientCryptography,
+          // @ts-expect-error Missing other required properties.
+        })({
+          params: { path: ['m', "44'", "1'"], curve: 'secp256k1' },
+        }),
+      ).toMatchInlineSnapshot(`
+        {
+          "chainCode": "0x8472428420c7fd8ef7280545bb6d2bde1d7c6b490556ccd59895f242716388d1",
+          "curve": "secp256k1",
+          "depth": 2,
+          "index": 2147483649,
+          "masterFingerprint": 3276136937,
+          "network": "mainnet",
+          "parentFingerprint": 1981505209,
+          "privateKey": "0x71d945aba22cd337ff26a107073ae2606dee5dbf7ecfe5c25870b8eaf62b9f1b",
+          "publicKey": "0x0491c4b234ca9b394f40d90f09092e04fd3bca2aa68c57e1311b25acfd972c5a6fc7ffd19e7812127473aa2bd827917b6ec7b57bec73cf022fc1f1fa0593f48770",
+        }
+      `);
+
+      expect(pbkdf2Sha512).toHaveBeenCalledTimes(1);
+    });
   });
 });

packages/snaps-rpc-methods/src/restricted/getBip32PublicKey.test.ts

@@ -154,5 +154,36 @@ describe('getBip32PublicKeyImplementation', () => {
         `"0x022de17487a660993177ce2a85bb73b6cd9ad436184d57bdf5a93f5db430bea914"`,
       );
     });
+
+    it('uses custom client cryptography functions', async () => {
+      const getUnlockPromise = jest.fn().mockResolvedValue(undefined);
+      const getMnemonic = jest
+        .fn()
+        .mockResolvedValue(TEST_SECRET_RECOVERY_PHRASE_BYTES);
+
+      const pbkdf2Sha512 = jest.fn().mockResolvedValue(new Uint8Array(64));
+      const getClientCryptography = jest.fn().mockReturnValue({
+        pbkdf2Sha512,
+      });
+
+      expect(
+        await getBip32PublicKeyImplementation({
+          getUnlockPromise,
+          getMnemonic,
+          getClientCryptography,
+          // @ts-expect-error Missing other required properties.
+        })({
+          params: {
+            path: ['m', "44'", "1'", '1', '2', '3'],
+            curve: 'secp256k1',
+            compressed: true,
+          },
+        }),
+      ).toMatchInlineSnapshot(
+        `"0x03102d63c39b6dda3f9aa06b247c50653cd9d01a91efce00ccc8735e9714058a01"`,
+      );
+
+      expect(pbkdf2Sha512).toHaveBeenCalledTimes(1);
+    });
   });
 });

packages/snaps-rpc-methods/src/restricted/getBip44Entropy.test.ts

@@ -89,5 +89,43 @@ describe('getBip44EntropyImplementation', () => {
         }
       `);
     });
+
+    it('uses custom client cryptography functions', async () => {
+      const getUnlockPromise = jest.fn().mockResolvedValue(undefined);
+      const getMnemonic = jest
+        .fn()
+        .mockResolvedValue(TEST_SECRET_RECOVERY_PHRASE_BYTES);
+
+      const pbkdf2Sha512 = jest.fn().mockResolvedValue(new Uint8Array(64));
+      const getClientCryptography = jest.fn().mockReturnValue({
+        pbkdf2Sha512,
+      });
+
+      expect(
+        await getBip44EntropyImplementation({
+          getUnlockPromise,
+          getMnemonic,
+          getClientCryptography,
+          // @ts-expect-error Missing other required properties.
+        })({
+          params: { coinType: 1 },
+        }),
+      ).toMatchInlineSnapshot(`
+        {
+          "chainCode": "0x8472428420c7fd8ef7280545bb6d2bde1d7c6b490556ccd59895f242716388d1",
+          "coin_type": 1,
+          "depth": 2,
+          "index": 2147483649,
+          "masterFingerprint": 3276136937,
+          "network": "mainnet",
+          "parentFingerprint": 1981505209,
+          "path": "m / bip32:44' / bip32:1'",
+          "privateKey": "0x71d945aba22cd337ff26a107073ae2606dee5dbf7ecfe5c25870b8eaf62b9f1b",
+          "publicKey": "0x0491c4b234ca9b394f40d90f09092e04fd3bca2aa68c57e1311b25acfd972c5a6fc7ffd19e7812127473aa2bd827917b6ec7b57bec73cf022fc1f1fa0593f48770",
+        }
+      `);
+
+      expect(pbkdf2Sha512).toHaveBeenCalledTimes(1);
+    });
   });
 });

packages/snaps-rpc-methods/src/restricted/getEntropy.test.ts

@@ -72,4 +72,43 @@ describe('getEntropyImplementation', () => {
       '0x6d8e92de419401c7da3cedd5f60ce5635b26059c2a4a8003877fec83653a4921',
     );
   });
+
+  it('uses custom client cryptography functions', async () => {
+    const getUnlockPromise = jest.fn().mockResolvedValue(undefined);
+    const getMnemonic = jest
+      .fn()
+      .mockResolvedValue(TEST_SECRET_RECOVERY_PHRASE_BYTES);
+
+    const pbkdf2Sha512 = jest.fn().mockResolvedValue(new Uint8Array(64));
+    const getClientCryptography = jest.fn().mockReturnValue({
+      pbkdf2Sha512,
+    });
+
+    const methodHooks = {
+      getMnemonic,
+      getUnlockPromise,
+      getClientCryptography,
+    };
+
+    const implementation = getEntropyBuilder.specificationBuilder({
+      methodHooks,
+    }).methodImplementation;
+
+    const result = await implementation({
+      method: 'snap_getEntropy',
+      params: {
+        version: 1,
+        salt: 'foo',
+      },
+      context: {
+        origin: MOCK_SNAP_ID,
+      },
+    });
+
+    expect(result).toBe(
+      '0x9bea47f2180fd874147f2f455a5ccc779826cfeff005605190cf0c568b3de7b5',
+    );
+
+    expect(pbkdf2Sha512).toHaveBeenCalledTimes(1);
+  });
 });