Add support for state methods to snaps-simulation

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

@@ -7,7 +7,7 @@
     "url": "https://github.com/MetaMask/snaps.git"
   },
   "source": {
-    "shasum": "REliam7FRJGwKCI4NaC2G3mBSD5iYR7DCEhrXLcBDqA=",
+    "shasum": "82KbG3cf0wtxooJpWzHeM1g4FhO8O7zSYCAAGNPshfM=",
     "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": "NItYOhAaWlS9q2r59/zlpoyVUyxojfsc5xMh65mLIwQ=",
+    "shasum": "5LsB950haZGnl0q5K7M4XgSh5J2e0p5O1Ptl/e6kpSQ=",
     "location": {
       "npm": {
         "filePath": "dist/bundle.js",

packages/examples/packages/manage-state/snap.manifest.json

@@ -7,7 +7,7 @@
     "url": "https://github.com/MetaMask/snaps.git"
   },
   "source": {
-    "shasum": "3jbTBm2Gtm5+qWdWgNR2sgwEGwWmKsGK7QPeXN9yOpE=",
+    "shasum": "fIXije73reQctVWFkOL9kdLWns7uDs7UWbPPL1J0f2o=",
     "location": {
       "npm": {
         "filePath": "dist/bundle.js",

packages/snaps-simulation/src/controllers.test.ts

@@ -5,10 +5,10 @@ import {
 } from '@metamask/permission-controller';
 
 import { getControllers } from './controllers';
-import type { MiddlewareHooks } from './simulation';
+import type { RestrictedMiddlewareHooks } from './simulation';
 import { getMockOptions } from './test-utils';
 
-const MOCK_HOOKS: MiddlewareHooks = {
+const MOCK_HOOKS: RestrictedMiddlewareHooks = {
   getIsLocked: jest.fn(),
   getMnemonic: jest.fn(),
   getSnapFile: jest.fn(),

packages/snaps-simulation/src/controllers.ts

@@ -29,7 +29,7 @@ import { getSafeJson } from '@metamask/utils';
 import { getPermissionSpecifications } from './methods';
 import { UNRESTRICTED_METHODS } from './methods/constants';
 import type { SimulationOptions } from './options';
-import type { MiddlewareHooks } from './simulation';
+import type { RestrictedMiddlewareHooks } from './simulation';
 import type { RunSagaFunction } from './store';
 
 export type RootControllerAllowedActions =
@@ -49,7 +49,7 @@ export type RootControllerMessenger = ControllerMessenger<
 
 export type GetControllersOptions = {
   controllerMessenger: ControllerMessenger<any, any>;
-  hooks: MiddlewareHooks;
+  hooks: RestrictedMiddlewareHooks;
   runSaga: RunSagaFunction;
   options: SimulationOptions;
 };

packages/snaps-simulation/src/methods/hooks/index.ts

@@ -1,5 +1,6 @@
 export * from './get-preferences';
+export * from './interface';
 export * from './notifications';
+export * from './permitted';
 export * from './request-user-approval';
 export * from './state';
-export * from './interface';

packages/snaps-simulation/src/methods/hooks/permitted/index.ts

@@ -0,0 +1 @@
+export * from './state';

packages/snaps-simulation/src/methods/hooks/permitted/state.test.ts

@@ -0,0 +1,119 @@
+import { createStore, getState, setState } from '../../../store';
+import { getMockOptions } from '../../../test-utils';
+import {
+  getPermittedClearSnapStateMethodImplementation,
+  getPermittedGetSnapStateMethodImplementation,
+  getPermittedUpdateSnapStateMethodImplementation,
+} from './state';
+
+describe('getPermittedGetSnapStateMethodImplementation', () => {
+  it('returns the implementation of the `getSnapState` hook', async () => {
+    const { store, runSaga } = createStore(getMockOptions());
+    const fn = getPermittedGetSnapStateMethodImplementation(runSaga);
+
+    expect(await fn(true)).toBeNull();
+
+    store.dispatch(
+      setState({
+        state: JSON.stringify({
+          foo: 'bar',
+        }),
+        encrypted: true,
+      }),
+    );
+
+    expect(await fn(true)).toStrictEqual({
+      foo: 'bar',
+    });
+  });
+
+  it('returns the implementation of the `getSnapState` hook for unencrypted state', async () => {
+    const { store, runSaga } = createStore(getMockOptions());
+    const fn = getPermittedGetSnapStateMethodImplementation(runSaga);
+
+    expect(await fn(false)).toBeNull();
+
+    store.dispatch(
+      setState({
+        state: JSON.stringify({
+          foo: 'bar',
+        }),
+        encrypted: false,
+      }),
+    );
+
+    expect(await fn(false)).toStrictEqual({
+      foo: 'bar',
+    });
+  });
+});
+
+describe('getPermittedUpdateSnapStateMethodImplementation', () => {
+  it('returns the implementation of the `updateSnapState` hook', async () => {
+    const { store, runSaga } = createStore(getMockOptions());
+    const fn = getPermittedUpdateSnapStateMethodImplementation(runSaga);
+
+    expect(getState(true)(store.getState())).toBeNull();
+
+    await fn({ foo: 'bar' }, true);
+
+    expect(getState(true)(store.getState())).toStrictEqual(
+      JSON.stringify({
+        foo: 'bar',
+      }),
+    );
+  });
+
+  it('returns the implementation of the `updateSnapState` hook for unencrypted state', async () => {
+    const { store, runSaga } = createStore(getMockOptions());
+    const fn = getPermittedUpdateSnapStateMethodImplementation(runSaga);
+
+    expect(getState(false)(store.getState())).toBeNull();
+
+    await fn({ foo: 'bar' }, false);
+
+    expect(getState(false)(store.getState())).toStrictEqual(
+      JSON.stringify({
+        foo: 'bar',
+      }),
+    );
+  });
+});
+
+describe('getPermittedClearSnapStateMethodImplementation', () => {
+  it('returns the implementation of the `clearSnapState` hook', async () => {
+    const { store, runSaga } = createStore(getMockOptions());
+    const fn = getPermittedClearSnapStateMethodImplementation(runSaga);
+
+    store.dispatch(
+      setState({
+        state: JSON.stringify({
+          foo: 'bar',
+        }),
+        encrypted: true,
+      }),
+    );
+
+    await fn(true);
+
+    expect(getState(true)(store.getState())).toBeNull();
+  });
+
+  it('returns the implementation of the `clearSnapState` hook for unencrypted state', async () => {
+    const { store, runSaga } = createStore(getMockOptions());
+    const fn = getPermittedClearSnapStateMethodImplementation(runSaga);
+
+    store.dispatch(
+      setState({
+        state: JSON.stringify({
+          foo: 'bar',
+        }),
+        encrypted: false,
+      }),
+    );
+
+    await fn(false);
+
+    expect(getState(false)(store.getState())).toBeNull();
+  });
+});

packages/snaps-simulation/src/methods/hooks/permitted/state.ts

@@ -0,0 +1,90 @@
+import { parseJson } from '@metamask/snaps-utils';
+import type { Json } from '@metamask/utils';
+import type { SagaIterator } from 'redux-saga';
+import { put, select } from 'redux-saga/effects';
+
+import type { RunSagaFunction } from '../../../store';
+import { clearState, getState, setState } from '../../../store';
+
+/**
+ * Get the Snap state from the store.
+ *
+ * @param encrypted - Whether to get the encrypted or unencrypted state.
+ * Defaults to encrypted state.
+ * @returns The state of the Snap.
+ * @yields Selects the state from the store.
+ */
+function* getSnapStateImplementation(encrypted: boolean): SagaIterator<string> {
+  const state = yield select(getState(encrypted));
+  // TODO: Use actual decryption implementation
+  return parseJson(state);
+}
+
+/**
+ * Get the implementation of the `getSnapState` hook.
+ *
+ * @param runSaga - The function to run a saga outside the usual Redux flow.
+ * @returns The implementation of the `getSnapState` hook.
+ */
+export function getPermittedGetSnapStateMethodImplementation(
+  runSaga: RunSagaFunction,
+) {
+  return async (...args: Parameters<typeof getSnapStateImplementation>) => {
+    return await runSaga(getSnapStateImplementation, ...args).toPromise();
+  };
+}
+
+/**
+ * Update the Snap state in the store.
+ *
+ * @param newState - The new state.
+ * @param encrypted - Whether to update the encrypted or unencrypted state.
+ * Defaults to encrypted state.
+ * @yields Puts the new state in the store.
+ */
+function* updateSnapStateImplementation(
+  newState: Record<string, Json>,
+  encrypted: boolean,
+): SagaIterator<void> {
+  // TODO: Use actual encryption implementation
+  yield put(setState({ state: JSON.stringify(newState), encrypted }));
+}
+
+/**
+ * Get the implementation of the `updateSnapState` hook.
+ *
+ * @param runSaga - The function to run a saga outside the usual Redux flow.
+ * @returns The implementation of the `updateSnapState` hook.
+ */
+export function getPermittedUpdateSnapStateMethodImplementation(
+  runSaga: RunSagaFunction,
+) {
+  return async (...args: Parameters<typeof updateSnapStateImplementation>) => {
+    return await runSaga(updateSnapStateImplementation, ...args).toPromise();
+  };
+}
+
+/**
+ * Clear the Snap state in the store.
+ *
+ * @param encrypted - Whether to clear the encrypted or unencrypted state.
+ * Defaults to encrypted state.
+ * @yields Puts the new state in the store.
+ */
+function* clearSnapStateImplementation(encrypted: boolean): SagaIterator<void> {
+  yield put(clearState({ encrypted }));
+}
+
+/**
+ * Get the implementation of the `clearSnapState` hook.
+ *
+ * @param runSaga - The function to run a saga outside the usual Redux flow.
+ * @returns The implementation of the `clearSnapState` hook.
+ */
+export function getPermittedClearSnapStateMethodImplementation(
+  runSaga: RunSagaFunction,
+) {
+  return async (...args: Parameters<typeof clearSnapStateImplementation>) => {
+    runSaga(clearSnapStateImplementation, ...args).result();
+  };
+}

packages/snaps-simulation/src/methods/specifications.test.ts

@@ -5,7 +5,7 @@ import {
 } from '@metamask/snaps-utils/test-utils';
 
 import { getControllers, registerSnap } from '../controllers';
-import type { MiddlewareHooks } from '../simulation';
+import type { RestrictedMiddlewareHooks } from '../simulation';
 import { getMockOptions } from '../test-utils/options';
 import {
   asyncResolve,
@@ -14,7 +14,7 @@ import {
   resolve,
 } from './specifications';
 
-const MOCK_HOOKS: MiddlewareHooks = {
+const MOCK_HOOKS: RestrictedMiddlewareHooks = {
   getMnemonic: jest.fn(),
   getSnapFile: jest.fn(),
   createInterface: jest.fn(),

packages/snaps-simulation/src/middleware/engine.test.ts

@@ -7,11 +7,18 @@ describe('createJsonRpcEngine', () => {
     const { store } = createStore(getMockOptions());
     const engine = createJsonRpcEngine({
       store,
-      hooks: {
+      restrictedHooks: {
         getMnemonic: jest.fn(),
-        getSnapFile: jest.fn().mockResolvedValue('foo'),
         getIsLocked: jest.fn(),
+        getClientCryptography: jest.fn(),
+      },
+      permittedHooks: {
+        getSnapFile: jest.fn().mockResolvedValue('foo'),
+        getSnapState: jest.fn(),
+        updateSnapState: jest.fn(),
+        clearSnapState: jest.fn(),
         getInterfaceState: jest.fn(),
+        getInterfaceContext: jest.fn(),
         createInterface: jest.fn(),
         updateInterface: jest.fn(),
         resolveInterface: jest.fn(),

packages/snaps-simulation/src/middleware/engine.ts

@@ -6,14 +6,18 @@ import { createSnapsMethodMiddleware } from '@metamask/snaps-rpc-methods';
 import type { Json } from '@metamask/utils';
 
 import { DEFAULT_JSON_RPC_ENDPOINT } from '../constants';
-import type { MiddlewareHooks } from '../simulation';
+import type {
+  PermittedMiddlewareHooks,
+  RestrictedMiddlewareHooks,
+} from '../simulation';
 import type { Store } from '../store';
 import { createInternalMethodsMiddleware } from './internal-methods';
 import { createMockMiddleware } from './mock';
 
 export type CreateJsonRpcEngineOptions = {
   store: Store;
-  hooks: MiddlewareHooks;
+  restrictedHooks: RestrictedMiddlewareHooks;
+  permittedHooks: PermittedMiddlewareHooks;
   permissionMiddleware: JsonRpcMiddleware<RestrictedMethodParameters, Json>;
   endpoint?: string;
 };
@@ -26,21 +30,23 @@ export type CreateJsonRpcEngineOptions = {
  *
  * @param options - The options to use when creating the engine.
  * @param options.store - The Redux store to use.
- * @param options.hooks - Any hooks used by the middleware handlers.
+ * @param options.restrictedHooks - Any hooks used by the middleware handlers.
+ * @param options.permittedHooks - Any hooks used by the middleware handlers.
  * @param options.permissionMiddleware - The permission middleware to use.
  * @param options.endpoint - The JSON-RPC endpoint to use for Ethereum requests.
  * @returns A JSON-RPC engine.
  */
 export function createJsonRpcEngine({
   store,
-  hooks,
+  restrictedHooks,
+  permittedHooks,
   permissionMiddleware,
   endpoint = DEFAULT_JSON_RPC_ENDPOINT,
 }: CreateJsonRpcEngineOptions) {
   const engine = new JsonRpcEngine();
   engine.push(createMockMiddleware(store));
-  engine.push(createInternalMethodsMiddleware(hooks));
-  engine.push(createSnapsMethodMiddleware(true, hooks));
+  engine.push(createInternalMethodsMiddleware(restrictedHooks));
+  engine.push(createSnapsMethodMiddleware(true, permittedHooks));
   engine.push(permissionMiddleware);
   engine.push(
     createFetchMiddleware({