Add property and hook for specifying custom cryptography functions

packages/examples/packages/bip32/package.json

@@ -43,7 +43,7 @@
     "test:watch": "jest --watch"
   },
   "dependencies": {
-    "@metamask/key-tree": "^9.1.2",
+    "@metamask/key-tree": "^10.0.1",
     "@metamask/snaps-sdk": "workspace:^",
     "@metamask/utils": "^10.0.0",
     "@noble/ed25519": "^1.6.0",

packages/examples/packages/bip32/snap.manifest.json

@@ -7,7 +7,7 @@
     "url": "https://github.com/MetaMask/snaps.git"
   },
   "source": {
-    "shasum": "ZgNEuQpFXjusvZX+0dsqN/jWaTFnk1T9mePMO2OxoQs=",
+    "shasum": "hRuh420QB8uksiS3rFwrvqNoQD5XTH/QyWkhFkmNBD8=",
     "location": {
       "npm": {
         "filePath": "dist/bundle.js",

packages/examples/packages/bip44/package.json

@@ -43,7 +43,7 @@
     "test:watch": "jest --watch"
   },
   "dependencies": {
-    "@metamask/key-tree": "^9.1.2",
+    "@metamask/key-tree": "^10.0.1",
     "@metamask/snaps-sdk": "workspace:^",
     "@metamask/utils": "^10.0.0",
     "@noble/bls12-381": "^1.2.0"

packages/examples/packages/bip44/snap.manifest.json

@@ -7,7 +7,7 @@
     "url": "https://github.com/MetaMask/snaps.git"
   },
   "source": {
-    "shasum": "dxKtnBcjmsoplDWa7k/dGTEVKyxl3slprHFoosSOPVI=",
+    "shasum": "kztNgPuBct9iJIGhWZs2i/yluGPJSQi0xl5+00opVGs=",
     "location": {
       "npm": {
         "filePath": "dist/bundle.js",

packages/examples/packages/browserify-plugin/snap.manifest.json

@@ -7,7 +7,7 @@
     "url": "https://github.com/MetaMask/snaps.git"
   },
   "source": {
-    "shasum": "ecGX3duI1nyJ8BOjkIPLze204JXMQKL8Eq1ir8Mm/dg=",
+    "shasum": "4yLB19XYAdGgHBPFlVOzCkb/JUZCjSajPRSQWs+a3uE=",
     "location": {
       "npm": {
         "filePath": "dist/bundle.js",

packages/examples/packages/browserify/snap.manifest.json

@@ -7,7 +7,7 @@
     "url": "https://github.com/MetaMask/snaps.git"
   },
   "source": {
-    "shasum": "KSkMBlnuET6wdxlrTCFlg6h1GDiCK8ShQoTbKPse0Ek=",
+    "shasum": "+0hxp1uhfCqe9KR+4RPDSPGHFTgRyGULKLn9XWwCmsY=",
     "location": {
       "npm": {
         "filePath": "dist/bundle.js",

packages/examples/packages/invoke-snap/packages/consumer-signer/package.json

@@ -43,7 +43,7 @@
     "test:watch": "jest --watch"
   },
   "dependencies": {
-    "@metamask/key-tree": "^9.1.2",
+    "@metamask/key-tree": "^10.0.1",
     "@metamask/snaps-sdk": "workspace:^",
     "@metamask/utils": "^10.0.0",
     "@noble/hashes": "^1.3.1"

packages/examples/packages/invoke-snap/packages/core-signer/package.json

@@ -43,7 +43,7 @@
     "test:watch": "jest --watch"
   },
   "dependencies": {
-    "@metamask/key-tree": "^9.1.2",
+    "@metamask/key-tree": "^10.0.1",
     "@metamask/snaps-sdk": "workspace:^",
     "@metamask/utils": "^10.0.0",
     "@noble/curves": "^1.1.0",

packages/examples/packages/invoke-snap/packages/core-signer/snap.manifest.json

@@ -7,7 +7,7 @@
     "url": "https://github.com/MetaMask/snaps.git"
   },
   "source": {
-    "shasum": "edIB0mUiM5lDzUI7tbS+VWnBxyB/ujEYRoQ/luywavA=",
+    "shasum": "g0lygIry0x1ULrACMgFTncUXfstO2l+7iM7/D65BXqY=",
     "location": {
       "npm": {
         "filePath": "dist/bundle.js",

packages/snaps-controllers/package.json

@@ -83,6 +83,7 @@
     "@metamask/base-controller": "^7.0.2",
     "@metamask/json-rpc-engine": "^10.0.1",
     "@metamask/json-rpc-middleware-stream": "^8.0.5",
+    "@metamask/key-tree": "^10.0.1",
     "@metamask/object-multiplex": "^2.0.0",
     "@metamask/permission-controller": "^11.0.3",
     "@metamask/phishing-controller": "^12.0.2",

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

@@ -55,6 +55,7 @@
 } from '@metamask/snaps-utils/test-utils';
 import type { SemVerRange, SemVerVersion, Json } from '@metamask/utils';
 import {
+  hexToBytes,
   assert,
   AssertionError,
   base64ToBytes,
@@ -1851,7 +1852,7 @@
   });
 
   // This isn't stable in CI unfortunately
   it.skip('throws if the Snap is terminated while executing', async () => {
     const { manifest, sourceCode, svgIcon } =
       await getMockSnapFilesWithUpdatedChecksum({
         sourceCode: `
@@ -8959,6 +8960,38 @@
 
       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-controllers/src/snaps/SnapController.ts

@@ -8,6 +8,7 @@ import type {
   ControllerStateChangeEvent,
 } from '@metamask/base-controller';
 import { BaseController } from '@metamask/base-controller';
+import type { CryptographicFunctions } from '@metamask/key-tree';
 import type {
   Caveat,
   GetEndowments,
@@ -711,7 +712,14 @@ type SnapControllerArgs = {
    * @returns The feature flags.
    */
   getFeatureFlags: () => DynamicFeatureFlags;
+
+  /**
+   * The cryptographic functions to use for the client. This may be an empty
+   * object to fall back to the default cryptographic functions.
+   */
+  clientCryptography?: CryptographicFunctions;
 };
+
 type AddSnapArgs = {
   id: SnapId;
   origin: string;
@@ -799,6 +807,8 @@ export class SnapController extends BaseController<
 
   #getFeatureFlags: () => DynamicFeatureFlags;
 
+  #clientCryptography: CryptographicFunctions | undefined;
+
   #detectSnapLocation: typeof detectSnapLocation;
 
   #snapsRuntimeData: Map<SnapId, SnapRuntimeData>;
@@ -832,6 +842,7 @@ export class SnapController extends BaseController<
     encryptor,
     getMnemonic,
     getFeatureFlags = () => ({}),
+    clientCryptography,
   }: SnapControllerArgs) {
     super({
       messenger,
@@ -887,6 +898,7 @@ export class SnapController extends BaseController<
     this.#encryptor = encryptor;
     this.#getMnemonic = getMnemonic;
     this.#getFeatureFlags = getFeatureFlags;
+    this.#clientCryptography = clientCryptography;
     this.#preinstalledSnaps = preinstalledSnaps;
     this._onUnhandledSnapError = this._onUnhandledSnapError.bind(this);
     this._onOutboundRequest = this._onOutboundRequest.bind(this);
@@ -1754,7 +1766,13 @@ export class SnapController extends BaseController<
 
     const salt = passedSalt ?? this.#encryptor.generateSalt();
     const mnemonicPhrase = await this.#getMnemonic();
-    const entropy = await getEncryptionEntropy({ snapId, mnemonicPhrase });
+
+    const entropy = await getEncryptionEntropy({
+      snapId,
+      mnemonicPhrase,
+      cryptographicFunctions: this.#clientCryptography,
+    });
+
     const encryptionKey = await this.#encryptor.keyFromPassword(
       entropy,
       salt,

packages/snaps-controllers/src/test-utils/controller.ts

@@ -573,13 +573,15 @@ export const getSnapControllerOptions = (
     state: undefined,
     fetchFunction: jest.fn(),
     getMnemonic: async () => Promise.resolve(TEST_SECRET_RECOVERY_PHRASE_BYTES),
+    clientCryptography: {},
     encryptor: getSnapControllerEncryptor(),
     ...opts,
   } as SnapControllerConstructorParams;
 
   options.state = {
     snaps: {},
     snapStates: {},
+    unencryptedSnapStates: {},
     ...options.state,
   };
   return options;

packages/snaps-rpc-methods/package.json

@@ -55,7 +55,7 @@
     "test:watch": "jest --watch"
   },
   "dependencies": {
-    "@metamask/key-tree": "^9.1.2",
+    "@metamask/key-tree": "^10.0.1",
     "@metamask/permission-controller": "^11.0.3",
     "@metamask/rpc-errors": "^7.0.1",
     "@metamask/snaps-sdk": "workspace:^",

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

@@ -11,6 +11,7 @@ describe('specificationBuilder', () => {
   const methodHooks = {
     getMnemonic: jest.fn(),
     getUnlockPromise: jest.fn(),
+    getClientCryptography: jest.fn(),
   };
 
   const specification = getBip32EntropyBuilder.specificationBuilder({
@@ -62,10 +63,15 @@ describe('getBip32EntropyImplementation', () => {
       const getMnemonic = jest
         .fn()
         .mockResolvedValue(TEST_SECRET_RECOVERY_PHRASE_BYTES);
+      const getClientCryptography = jest.fn().mockReturnValue({});
 
       expect(
-        // @ts-expect-error Missing other required properties.
-        await getBip32EntropyImplementation({ getUnlockPromise, getMnemonic })({
+        await getBip32EntropyImplementation({
+          getUnlockPromise,
+          getMnemonic,
+          getClientCryptography,
+          // @ts-expect-error Missing other required properties.
+        })({
           params: { path: ['m', "44'", "1'"], curve: 'secp256k1' },
         }),
       ).toMatchInlineSnapshot(`
@@ -75,6 +81,7 @@ describe('getBip32EntropyImplementation', () => {
           "depth": 2,
           "index": 2147483649,
           "masterFingerprint": 1404659567,
+          "network": "mainnet",
           "parentFingerprint": 1829122711,
           "privateKey": "0xc73cedb996e7294f032766853a8b7ba11ab4ce9755fc052f2f7b9000044c99af",
           "publicKey": "0x048e129862c1de5ca86468add43b001d32fd34b8113de716ecd63fa355b7f1165f0e76f5dc6095100f9fdaa76ddf28aa3f21406ac5fda7c71ffbedb45634fe2ceb",
@@ -87,10 +94,15 @@ describe('getBip32EntropyImplementation', () => {
       const getMnemonic = jest
         .fn()
         .mockResolvedValue(TEST_SECRET_RECOVERY_PHRASE_BYTES);
+      const getClientCryptography = jest.fn().mockReturnValue({});
 
       expect(
-        // @ts-expect-error Missing other required properties.
-        await getBip32EntropyImplementation({ getUnlockPromise, getMnemonic })({
+        await getBip32EntropyImplementation({
+          getUnlockPromise,
+          getMnemonic,
+          getClientCryptography,
+          // @ts-expect-error Missing other required properties.
+        })({
           params: {
             path: ['m', "44'", "1'", "0'", '0', '1'],
             curve: 'secp256k1',
@@ -103,6 +115,7 @@ describe('getBip32EntropyImplementation', () => {
           "depth": 5,
           "index": 1,
           "masterFingerprint": 1404659567,
+          "network": "mainnet",
           "parentFingerprint": 3495658567,
           "privateKey": "0x43a9353dfebf7209c3feb1843510299e2b0f4fa09151dccc3824df88451be37c",
           "publicKey": "0x0467f3cac111f47782b6c2d8d0984d51e22c128d24ec3eaca044509a386771d17206c740c7337c399d8ade8f52a60029340f288e11de82fffd3b69c5b863f6a515",
@@ -116,9 +129,15 @@ describe('getBip32EntropyImplementation', () => {
         .fn()
         .mockResolvedValue(TEST_SECRET_RECOVERY_PHRASE_BYTES);
 
+      const getClientCryptography = jest.fn().mockReturnValue({});
+
       expect(
-        // @ts-expect-error Missing other required properties.
-        await getBip32EntropyImplementation({ getUnlockPromise, getMnemonic })({
+        await getBip32EntropyImplementation({
+          getUnlockPromise,
+          getMnemonic,
+          getClientCryptography,
+          // @ts-expect-error Missing other required properties.
+        })({
           params: {
             path: ['m', "44'", "1'", "0'", "0'", "1'"],
             curve: 'ed25519',
@@ -131,6 +150,7 @@ describe('getBip32EntropyImplementation', () => {
           "depth": 5,
           "index": 2147483649,
           "masterFingerprint": 650419359,
+          "network": "mainnet",
           "parentFingerprint": 660188756,
           "privateKey": "0x5e6ebe8f5c33833e6c86f8769da173daa206b9dfd1956efcd2b115d82376bb5e",
           "publicKey": "0x0012affaf55babdfb59b76adcf00f69442f019974124639108470409d47e25e19f",
@@ -144,9 +164,15 @@ describe('getBip32EntropyImplementation', () => {
         .fn()
         .mockResolvedValue(TEST_SECRET_RECOVERY_PHRASE_BYTES);
 
+      const getClientCryptography = jest.fn().mockReturnValue({});
+
       expect(
-        // @ts-expect-error Missing other required properties.
-        await getBip32EntropyImplementation({ getUnlockPromise, getMnemonic })({
+        await getBip32EntropyImplementation({
+          getUnlockPromise,
+          getMnemonic,
+          getClientCryptography,
+          // @ts-expect-error Missing other required properties.
+        })({
           params: {
             path: ['m', "44'", "1'", "0'", "0'", "1'"],
             curve: 'ed25519Bip32',
@@ -159,11 +185,49 @@ describe('getBip32EntropyImplementation', () => {
           "depth": 5,
           "index": 2147483649,
           "masterFingerprint": 1587894111,
+          "network": "mainnet",
           "parentFingerprint": 3236688876,
           "privateKey": "0x88a59d7aa9fe82d8f98843ef474195178eb71956dee597252e7a5fbeebbc734e9b5bfdd17f82144a2bea78c8ab19bef26dc93f36e96eaa41453b65cb3daa1817",
           "publicKey": "0xd91d18b4540a2f30341e8463d5f9b25b14fae9a236dcbea338b668a318bb0867",
         }
       `);
     });
+
+    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);
+    });
   });
 });