Cache snap state in memory

Author: Mrtenz

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

@@ -97,6 +97,7 @@ import {
   MOCK_WALLET_SNAP_PERMISSION,
   MockSnapsRegistry,
   sleep,
+  waitForStateChange,
 } from '../test-utils';
 import { delay } from '../utils';
 import { LEGACY_ENCRYPTION_KEY_DERIVATION_OPTIONS } from './constants';
@@ -8801,6 +8802,7 @@ describe('SnapController', () => {
       );
 
       const newState = { myVariable: 2 };
+      const promise = waitForStateChange(messenger);
 
       await messenger.call(
         'SnapController:updateSnapState',
@@ -8817,6 +8819,8 @@ describe('SnapController', () => {
         DEFAULT_ENCRYPTION_KEY_DERIVATION_OPTIONS,
       );
 
+      await promise;
+
       const result = await messenger.call(
         'SnapController:getSnapState',
         MOCK_SNAP_ID,
@@ -8831,7 +8835,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 +8861,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 +9081,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,13 +9149,17 @@ 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();

packages/snaps-controllers/src/snaps/SnapController.ts

@@ -65,6 +65,7 @@ import type {
   TruncatedSnapFields,
 } from '@metamask/snaps-utils';
 import {
+  withMutex,
   logWarning,
   getPlatformVersion,
   assertIsSnapManifest,
@@ -252,6 +253,16 @@ export interface SnapRuntimeData {
    * A boolean flag to determine whether the Snap is currently being stopped.
    */
   stopping: boolean;
+
+  /**
+   * Cached encrypted state of the Snap.
+   */
+  state: Record<string, Json> | null;
+
+  /**
+   * Cached unencrypted state of the Snap.
+   */
+  unencryptedState: Record<string, Json> | null;
 }
 
 export type SnapError = {
@@ -906,6 +917,7 @@ export class SnapController extends BaseController<
     this._onOutboundResponse = this._onOutboundResponse.bind(this);
     this.#rollbackSnapshots = new Map();
     this.#snapsRuntimeData = new Map();
+
     this.#pollForLastRequestStatus();
 
     /* eslint-disable @typescript-eslint/unbound-method */
@@ -1860,6 +1872,68 @@ export class SnapController extends BaseController<
     return JSON.stringify(encryptedState);
   }
 
+  /**
+   * Get the new Snap state to persist based on the given state and encryption
+   * flag.
+   *
+   * - If the state is null, return null.
+   * - If the state should not be encrypted, return the JSON stringified state.
+   * - Otherwise if the state should be encrypted, return the encrypted state.
+   *
+   * @param snapId - The Snap ID.
+   * @param state - The state to persist.
+   * @param encrypted - A flag to indicate whether to use encrypted storage or
+   * not.
+   * @returns The state to persist.
+   */
+  async #getStateToPersist(
+    snapId: SnapId,
+    state: Record<string, Json> | null,
+    encrypted: boolean,
+  ) {
+    if (state === null) {
+      return null;
+    }
+
+    if (encrypted) {
+      return await this.#encryptSnapState(snapId, state);
+    }
+
+    return JSON.stringify(state);
+  }
+
+  /**
+   * Persist the state of a Snap.
+   *
+   * @param snapId - The Snap ID.
+   * @param newSnapState - The new state of the Snap.
+   * @param encrypted - A flag to indicate whether to use encrypted storage or
+   * not.
+   */
+  #persistSnapState = withMutex(
+    async (
+      snapId: SnapId,
+      newSnapState: Record<string, Json> | null,
+      encrypted: boolean,
+    ) => {
+      const newState = await this.#getStateToPersist(
+        snapId,
+        newSnapState,
+        encrypted,
+      );
+
+      if (encrypted) {
+        return this.update((state) => {
+          state.snapStates[snapId] = newState;
+        });
+      }
+
+      return this.update((state) => {
+        state.unencryptedSnapStates[snapId] = newState;
+      });
+    },
+  );
+
   /**
    * Updates the own state of the snap with the given id.
    * This is distinct from the state MetaMask uses to manage snaps.
@@ -1873,17 +1947,19 @@ export class SnapController extends BaseController<
     newSnapState: Record<string, Json>,
     encrypted: boolean,
   ) {
-    if (encrypted) {
-      const encryptedState = await this.#encryptSnapState(snapId, newSnapState);
+    const runtime = this.#getRuntimeExpect(snapId);
 
-      this.update((state) => {
-        state.snapStates[snapId] = encryptedState;
-      });
+    if (encrypted) {
+      runtime.state = newSnapState;
     } else {
-      this.update((state) => {
-        state.unencryptedSnapStates[snapId] = JSON.stringify(newSnapState);
-      });
+      runtime.unencryptedState = newSnapState;
     }
+
+    // This is intentionally run asynchronously to avoid blocking the main
+    // thread.
+    this.#persistSnapState(snapId, newSnapState, encrypted).catch((error) => {
+      logError(error);
+    });
   }
 
   /**
@@ -1894,12 +1970,17 @@ export class SnapController extends BaseController<
    * @param encrypted - A flag to indicate whether to use encrypted storage or not.
    */
   clearSnapState(snapId: SnapId, encrypted: boolean) {
-    this.update((state) => {
-      if (encrypted) {
-        state.snapStates[snapId] = null;
-      } else {
-        state.unencryptedSnapStates[snapId] = null;
-      }
+    const runtime = this.#getRuntimeExpect(snapId);
+    if (encrypted) {
+      runtime.state = null;
+    } else {
+      runtime.unencryptedState = null;
+    }
+
+    // This is intentionally run asynchronously to avoid blocking the main
+    // thread.
+    this.#persistSnapState(snapId, null, encrypted).catch((error) => {
+      logError(error);
     });
   }
 
@@ -1912,6 +1993,13 @@ export class SnapController extends BaseController<
    * @returns The requested snap state or null if no state exists.
    */
   async getSnapState(snapId: SnapId, encrypted: boolean): Promise<Json> {
+    const runtime = this.#getRuntimeExpect(snapId);
+    const cachedState = encrypted ? runtime.state : runtime.unencryptedState;
+
+    if (cachedState !== undefined) {
+      return cachedState;
+    }
+
     const state = encrypted
       ? this.state.snapStates[snapId]
       : this.state.unencryptedSnapStates[snapId];
@@ -1921,11 +2009,17 @@ export class SnapController extends BaseController<
     }
 
     if (!encrypted) {
-      // For performance reasons, we do not validate that the state is JSON, since we control serialization.
-      return JSON.parse(state);
+      // For performance reasons, we do not validate that the state is JSON,
+      // since we control serialization.
+      const json = JSON.parse(state);
+      runtime.unencryptedState = json;
+
+      return json;
     }
 
     const decrypted = await this.#decryptSnapState(snapId, state);
+    runtime.state = decrypted;
+
     return decrypted;
   }
 
@@ -3706,6 +3800,8 @@ export class SnapController extends BaseController<
       pendingOutboundRequests: 0,
       interpreter,
       stopping: false,
+      state: undefined,
+      unencryptedState: undefined,
     });
   }
 

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

@@ -1,4 +1,8 @@
 import type { ApprovalRequest } from '@metamask/approval-controller';
+import type {
+  ControllerMessenger,
+  RestrictedControllerMessenger,
+} from '@metamask/base-controller';
 import {
   encryptWithKey,
   decryptWithKey,
@@ -48,16 +52,17 @@ import type {
   SnapInterfaceControllerEvents,
   StoredInterface,
 } from '../interface/SnapInterfaceController';
+import { SnapController } from '../snaps';
 import type {
   AllowedActions,
   AllowedEvents,
   PersistedSnapControllerState,
   SnapControllerActions,
   SnapControllerEvents,
+  SnapControllerStateChangeEvent,
   SnapsRegistryActions,
   SnapsRegistryEvents,
 } from '../snaps';
-import { SnapController } from '../snaps';
 import type { KeyDerivationOptions } from '../types';
 import { MOCK_CRONJOB_PERMISSION } from './cronjob';
 import { getNodeEES, getNodeEESMessenger } from './execution-environment';
@@ -830,3 +835,27 @@ export const getRestrictedSnapInsightsControllerMessenger = (
 
   return controllerMessenger;
 };
+
+/**
+ * Wait for the state change event to be emitted by the messenger.
+ *
+ * @param messenger - The messenger to listen to.
+ * @returns A promise that resolves when the state change event is emitted.
+ */
+export async function waitForStateChange(
+  messenger:
+    | ControllerMessenger<any, SnapControllerStateChangeEvent>
+    | RestrictedControllerMessenger<
+        'SnapController',
+        any,
+        SnapControllerStateChangeEvent,
+        any,
+        'SnapController:stateChange'
+      >,
+) {
+  return new Promise<void>((resolve) => {
+    messenger.subscribe('SnapController:stateChange', () => {
+      resolve();
+    });
+  });
+}

packages/snaps-utils/package.json

@@ -90,6 +90,7 @@
     "@metamask/utils": "^10.0.0",
     "@noble/hashes": "^1.3.1",
     "@scure/base": "^1.1.1",
+    "async-mutex": "^0.4.0",
     "chalk": "^4.1.2",
     "cron-parser": "^4.5.0",
     "fast-deep-equal": "^3.1.3",

packages/snaps-utils/src/index.ts

@@ -20,6 +20,7 @@ export * from './json-rpc';
 export * from './localization';
 export * from './logging';
 export * from './manifest';
+export * from './mutex';
 export * from './namespace';
 export * from './path';
 export * from './platform-version';

packages/snaps-utils/src/mutex.ts

@@ -0,0 +1,24 @@
+import { Mutex } from 'async-mutex';
+
+/**
+ * Run a function with an async mutex, ensuring that only one instance of the
+ * function can run at a time.
+ *
+ * @param fn - The function to run with a mutex.
+ * @returns The wrapped function.
+ * @template OriginalFunction - The original function type. This is inferred
+ * from the `fn` argument, and used to determine the return type of the
+ * wrapped function.
+ */
+export function withMutex<
+  OriginalFunction extends (...args: any[]) => Promise<Type>,
+  Type,
+>(
+  fn: OriginalFunction,
+): (...args: Parameters<OriginalFunction>) => Promise<Type> {
+  const mutex = new Mutex();
+
+  return async (...args: Parameters<OriginalFunction>) => {
+    return await mutex.runExclusive(async () => await fn(...args));
+  };
+}

yarn.lock

@@ -6280,6 +6280,7 @@ __metadata:
     "@wdio/spec-reporter": "npm:^8.19.0"
     "@wdio/static-server-service": "npm:^8.19.0"
     "@wdio/types": "npm:^8.19.0"
+    async-mutex: "npm:^0.4.0"
     chalk: "npm:^4.1.2"
     cron-parser: "npm:^4.5.0"
     deepmerge: "npm:^4.2.2"