Cache snap state in memory

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 @@
   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 @@
   AssertionError,
   base64ToBytes,
   stringToBytes,
+  createDeferredPromise,
 } from '@metamask/utils';
 import { File } from 'buffer';
 import { webcrypto } from 'crypto';
@@ -78,6 +80,7 @@
   getNodeEESMessenger,
   getPersistedSnapsState,
   getSnapController,
+  getSnapControllerEncryptor,
   getSnapControllerMessenger,
   getSnapControllerOptions,
   getSnapControllerWithEES,
@@ -97,6 +100,7 @@
   MOCK_WALLET_SNAP_PERMISSION,
   MockSnapsRegistry,
   sleep,
+  waitForStateChange,
 } from '../test-utils';
 import { delay } from '../utils';
 import { LEGACY_ENCRYPTION_KEY_DERIVATION_OPTIONS } from './constants';
@@ -1838,7 +1842,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: `
@@ -2117,6 +2121,59 @@
     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 @@
       );
 
       const newState = { myVariable: 2 };
+      const promise = waitForStateChange(messenger);
 
       await messenger.call(
         'SnapController:updateSnapState',
@@ -8817,6 +8875,8 @@
         DEFAULT_ENCRYPTION_KEY_DERIVATION_OPTIONS,
       );
 
+      await promise;
+
       const result = await messenger.call(
         'SnapController:getSnapState',
         MOCK_SNAP_ID,
@@ -8831,7 +8891,7 @@
       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 @@
         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 @@
         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 @@
       );
 
       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 @@
 
       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', () => {