Clear cache on lock

Author: Mrtenz

packages/snaps-controllers/coverage.json

@@ -1,6 +1,6 @@
 {
   "branches": 92.41,
   "functions": 96.56,
-  "lines": 97.99,
-  "statements": 97.71
+  "lines": 98,
+  "statements": 97.72
 }

packages/snaps-controllers/package.json

@@ -84,6 +84,7 @@
     "@metamask/json-rpc-engine": "^10.0.1",
     "@metamask/json-rpc-middleware-stream": "^8.0.5",
     "@metamask/key-tree": "^10.0.1",
+    "@metamask/keyring-controller": "^19.0.2",
     "@metamask/object-multiplex": "^2.0.0",
     "@metamask/permission-controller": "^11.0.4",
     "@metamask/phishing-controller": "^12.3.1",

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 {
@@ -2120,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(

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

@@ -9,6 +9,7 @@ import type {
 } from '@metamask/base-controller';
 import { BaseController } from '@metamask/base-controller';
 import type { CryptographicFunctions } from '@metamask/key-tree';
+import type { KeyringControllerLockEvent } from '@metamask/keyring-controller';
 import type {
   Caveat,
   GetEndowments,
@@ -607,7 +608,8 @@ export type AllowedActions =
 export type AllowedEvents =
   | ExecutionServiceEvents
   | SnapInstalled
-  | SnapUpdated;
+  | SnapUpdated
+  | KeyringControllerLockEvent;
 
 type SnapControllerMessenger = RestrictedControllerMessenger<
   typeof controllerName,
@@ -967,6 +969,11 @@ export class SnapController extends BaseController<
       },
     );
 
+    this.messagingSystem.subscribe(
+      'KeyringController:lock',
+      this.#handleLock.bind(this),
+    );
+
     this.#initializeStateMachine();
     this.#registerMessageHandlers();
 
@@ -1832,6 +1839,7 @@ export class SnapController extends BaseController<
       const useCache =
         this.#hasCachedEncryptionKey(snapId) ||
         this.#encryptor.isVaultUpdated(state);
+
       const { key } = await this.#getSnapEncryptionKey({
         snapId,
         salt,
@@ -4007,4 +4015,17 @@ export class SnapController extends BaseController<
       },
     });
   }
+
+  /**
+   * Handle the `KeyringController:lock` event.
+   *
+   * Currently this clears the cached encrypted state (if any) for all Snaps.
+   */
+  #handleLock() {
+    for (const runtime of this.#snapsRuntimeData.values()) {
+      runtime.encryptionKey = null;
+      runtime.encryptionSalt = null;
+      runtime.state = undefined;
+    }
+  }
 }

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

@@ -469,6 +469,7 @@ export const getSnapControllerMessenger = (
       'SnapController:snapUpdated',
       'SnapController:stateChange',
       'SnapController:snapRolledback',
+      'KeyringController:lock',
     ],
     allowedActions: [
       'ApprovalController:addRequest',