Cache snap state in memory (#2980)

To improve performance when using state, especially when using the new
state methods (#2916), we now cache state in memory, and write updates
to disk asynchronously. The first time the state is fetched, it's cached
in the Snap's runtime data, which is used in subsequent calls.

Author: Mrtenz

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

@@ -47,7 +47,7 @@
     "@metamask/snaps-sdk": "workspace:^",
     "@metamask/utils": "^10.0.0",
     "@noble/curves": "^1.1.0",
-    "async-mutex": "^0.4.0"
+    "async-mutex": "^0.5.0"
   },
   "devDependencies": {
     "@jest/globals": "^29.5.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": "9+79ZuJLehTDLvMMK/dR0C29/5Q/GRdvTq8EaxTwQkU=",
+    "shasum": "5YpYX3b3wdRQEjPd2lUeNsNK7FwiflxMLCCPYtDeLnQ=",
     "location": {
       "npm": {
         "filePath": "dist/bundle.js",

packages/snaps-controllers/coverage.json

@@ -1,6 +1,6 @@
 {
-  "branches": 93.06,
-  "functions": 96.54,
-  "lines": 98.02,
-  "statements": 97.74
+  "branches": 92.96,
+  "functions": 96.56,
+  "lines": 98.05,
+  "statements": 97.77
 }

packages/snaps-controllers/package.json

@@ -95,6 +95,7 @@
     "@metamask/snaps-utils": "workspace:^",
     "@metamask/utils": "^10.0.0",
     "@xstate/fsm": "^2.0.0",
+    "async-mutex": "^0.5.0",
     "browserify-zlib": "^0.2.0",
     "concat-stream": "^2.0.0",
     "fast-deep-equal": "^3.1.3",

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

@@ -52,6 +52,7 @@ import {
   MOCK_SNAP_NAME,
   DEFAULT_SOURCE_PATH,
   DEFAULT_ICON_PATH,
+  TEST_SECRET_RECOVERY_PHRASE_BYTES,
 } from '@metamask/snaps-utils/test-utils';
 import type { SemVerRange, SemVerVersion, Json } from '@metamask/utils';
 import {
@@ -60,6 +61,7 @@ import {
   AssertionError,
   base64ToBytes,
   stringToBytes,
+  createDeferredPromise,
 } from '@metamask/utils';
 import { File } from 'buffer';
 import { webcrypto } from 'crypto';
@@ -78,6 +80,7 @@ import {
   getNodeEESMessenger,
   getPersistedSnapsState,
   getSnapController,
+  getSnapControllerEncryptor,
   getSnapControllerMessenger,
   getSnapControllerOptions,
   getSnapControllerWithEES,
@@ -97,6 +100,7 @@ import {
   MOCK_WALLET_SNAP_PERMISSION,
   MockSnapsRegistry,
   sleep,
+  waitForStateChange,
 } from '../test-utils';
 import { delay } from '../utils';
 import { LEGACY_ENCRYPTION_KEY_DERIVATION_OPTIONS } from './constants';
@@ -2117,6 +2121,59 @@ describe('SnapController', () => {
     await service.terminateAllSnaps();
   });
 
+  it('clears encrypted state of Snaps when the client is locked', async () => {
+    const rootMessenger = getControllerMessenger();
+    const messenger = getSnapControllerMessenger(rootMessenger);
+
+    const state = { myVariable: 1 };
+
+    const mockEncryptedState = await encrypt(
+      ENCRYPTION_KEY,
+      state,
+      undefined,
+      undefined,
+      DEFAULT_ENCRYPTION_KEY_DERIVATION_OPTIONS,
+    );
+
+    const getMnemonic = jest
+      .fn()
+      .mockReturnValue(TEST_SECRET_RECOVERY_PHRASE_BYTES);
+
+    const snapController = getSnapController(
+      getSnapControllerOptions({
+        messenger,
+        state: {
+          snaps: {
+            [MOCK_SNAP_ID]: getPersistedSnapObject(),
+          },
+          snapStates: {
+            [MOCK_SNAP_ID]: mockEncryptedState,
+          },
+        },
+        getMnemonic,
+      }),
+    );
+
+    expect(
+      await messenger.call('SnapController:getSnapState', MOCK_SNAP_ID, true),
+    ).toStrictEqual(state);
+    expect(getMnemonic).toHaveBeenCalledTimes(1);
+
+    rootMessenger.publish('KeyringController:lock');
+
+    expect(
+      await messenger.call('SnapController:getSnapState', MOCK_SNAP_ID, true),
+    ).toStrictEqual(state);
+
+    // We assume `getMnemonic` is called again because the controller needs to
+    // decrypt the state again. This is not an ideal way to test this, but it
+    // is the easiest to test this without exposing the internal state of the
+    // `SnapController`.
+    expect(getMnemonic).toHaveBeenCalledTimes(2);
+
+    snapController.destroy();
+  });
+
   describe('handleRequest', () => {
     it.each(
       Object.keys(handlerEndowments).filter(
@@ -8801,6 +8858,7 @@ describe('SnapController', () => {
       );
 
       const newState = { myVariable: 2 };
+      const promise = waitForStateChange(messenger);
 
       await messenger.call(
         'SnapController:updateSnapState',
@@ -8817,6 +8875,8 @@ describe('SnapController', () => {
         DEFAULT_ENCRYPTION_KEY_DERIVATION_OPTIONS,
       );
 
+      await promise;
+
       const result = await messenger.call(
         'SnapController:getSnapState',
         MOCK_SNAP_ID,
@@ -8831,7 +8891,7 @@ describe('SnapController', () => {
       snapController.destroy();
     });
 
-    it('different snaps use different encryption keys', async () => {
+    it('uses different encryption keys for different snaps', async () => {
       const messenger = getSnapControllerMessenger();
 
       const state = { foo: 'bar' };
@@ -8857,13 +8917,17 @@ describe('SnapController', () => {
         true,
       );
 
+      const promise = waitForStateChange(messenger);
+
       await messenger.call(
         'SnapController:updateSnapState',
         MOCK_LOCAL_SNAP_ID,
         state,
         true,
       );
 
+      await promise;
+
       const encryptedState1 = await encrypt(
         ENCRYPTION_KEY,
         state,
@@ -9073,13 +9137,17 @@ describe('SnapController', () => {
         undefined,
         DEFAULT_ENCRYPTION_KEY_DERIVATION_OPTIONS,
       );
+
+      const promise = waitForStateChange(messenger);
       await messenger.call(
         'SnapController:updateSnapState',
         MOCK_SNAP_ID,
         state,
         true,
       );
 
+      await promise;
+
       expect(updateSnapStateSpy).toHaveBeenCalledTimes(1);
       expect(snapController.state.snapStates[MOCK_SNAP_ID]).toStrictEqual(
         mockEncryptedState,
@@ -9137,17 +9205,126 @@ describe('SnapController', () => {
       );
 
       const state = { foo: 'bar' };
+
+      const promise = waitForStateChange(messenger);
       await messenger.call(
         'SnapController:updateSnapState',
         MOCK_SNAP_ID,
         state,
         true,
       );
 
+      await promise;
+
       expect(pbkdf2Sha512).toHaveBeenCalledTimes(1);
 
       snapController.destroy();
     });
+
+    it('queues multiple state updates', async () => {
+      const messenger = getSnapControllerMessenger();
+
+      jest.useFakeTimers();
+
+      const encryptor = getSnapControllerEncryptor();
+      const { promise, resolve } = createDeferredPromise();
+      const encryptWithKey = jest
+        .fn<
+          ReturnType<typeof encryptor.encryptWithKey>,
+          Parameters<typeof encryptor.encryptWithKey>
+        >()
+        .mockImplementation(async (...args) => {
+          resolve();
+          await sleep(1);
+          return await encryptor.encryptWithKey(...args);
+        });
+
+      const snapController = getSnapController(
+        getSnapControllerOptions({
+          messenger,
+          state: {
+            snaps: getPersistedSnapsState(),
+          },
+          encryptor: {
+            ...getSnapControllerEncryptor(),
+            // @ts-expect-error - Missing required properties.
+            encryptWithKey,
+          },
+        }),
+      );
+
+      const firstStateChange = waitForStateChange(messenger);
+      await messenger.call(
+        'SnapController:updateSnapState',
+        MOCK_SNAP_ID,
+        { foo: 'bar' },
+        true,
+      );
+
+      await messenger.call(
+        'SnapController:updateSnapState',
+        MOCK_SNAP_ID,
+        { bar: 'baz' },
+        true,
+      );
+
+      // We await this promise to ensure the timer is queued.
+      await promise;
+      jest.advanceTimersByTime(1);
+
+      // After this point the second update should be queued.
+      await firstStateChange;
+      const secondStateChange = waitForStateChange(messenger);
+
+      expect(encryptWithKey).toHaveBeenCalledTimes(1);
+
+      // This is a bit hacky, but we can't simply advance the timer by 1ms
+      // because the second timer is not running yet.
+      jest.useRealTimers();
+      await secondStateChange;
+
+      expect(encryptWithKey).toHaveBeenCalledTimes(2);
+
+      expect(
+        await messenger.call('SnapController:getSnapState', MOCK_SNAP_ID, true),
+      ).toStrictEqual({ bar: 'baz' });
+
+      snapController.destroy();
+    });
+
+    it('logs an error message if the state fails to persist', async () => {
+      const messenger = getSnapControllerMessenger();
+
+      const errorValue = new Error('Failed to persist state.');
+      const snapController = getSnapController(
+        getSnapControllerOptions({
+          messenger,
+          state: {
+            snaps: getPersistedSnapsState(),
+          },
+          // @ts-expect-error - Missing required properties.
+          encryptor: {
+            ...getSnapControllerEncryptor(),
+            encryptWithKey: jest.fn().mockRejectedValue(errorValue),
+          },
+        }),
+      );
+
+      const { promise, resolve } = createDeferredPromise();
+      const error = jest.spyOn(console, 'error').mockImplementation(resolve);
+
+      await messenger.call(
+        'SnapController:updateSnapState',
+        MOCK_SNAP_ID,
+        { foo: 'bar' },
+        true,
+      );
+
+      await promise;
+      expect(error).toHaveBeenCalledWith(errorValue);
+
+      snapController.destroy();
+    });
   });
 
   describe('SnapController:clearSnapState', () => {
@@ -9206,6 +9383,41 @@ describe('SnapController', () => {
 
       snapController.destroy();
     });
+
+    it('logs an error message if the state fails to persist', async () => {
+      const messenger = getSnapControllerMessenger();
+
+      const errorValue = new Error('Failed to persist state.');
+      const snapController = getSnapController(
+        getSnapControllerOptions({
+          messenger,
+          state: {
+            snaps: getPersistedSnapsState(),
+          },
+          // @ts-expect-error - Missing required properties.
+          encryptor: {
+            ...getSnapControllerEncryptor(),
+            encryptWithKey: jest.fn().mockRejectedValue(errorValue),
+          },
+        }),
+      );
+
+      const { promise, resolve } = createDeferredPromise();
+      const error = jest.spyOn(console, 'error').mockImplementation(resolve);
+
+      // @ts-expect-error - Property `update` is protected.
+      // eslint-disable-next-line jest/prefer-spy-on
+      snapController.update = jest.fn().mockImplementation(() => {
+        throw errorValue;
+      });
+
+      await messenger.call('SnapController:clearSnapState', MOCK_SNAP_ID, true);
+
+      await promise;
+      expect(error).toHaveBeenCalledWith(errorValue);
+
+      snapController.destroy();
+    });
   });
 
   describe('SnapController:updateBlockedSnaps', () => {