feat: allow for multiple keys in snap_setState

packages/examples/packages/manage-state/CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- Add tests for `snap_setState` with array `key` parameter ([#4126](https://github.com/MetaMask/snaps/pull/4126))
+
 ## [3.0.0]
 
 ### Added

packages/examples/packages/manage-state/src/index.test.ts

@@ -101,6 +101,29 @@ describe('onRpcRequest', () => {
       });
     });
 
+    it('sets the state for multiple keys', async () => {
+      const { request } = await installSnap();
+
+      expect(
+        await request({
+          method: 'setState',
+          params: {
+            key: ['foo', 'baz'],
+            value: { foo: 'bar', baz: 'qux' },
+          },
+        }),
+      ).toRespondWith(null);
+
+      expect(
+        await request({
+          method: 'getState',
+          params: {
+            key: ['foo', 'baz'],
+          },
+        }),
+      ).toRespondWith({ foo: 'bar', baz: 'qux' });
+    });
+
     it('throws if the state is not an object and no key is specified', async () => {
       const { request } = await installSnap();
 

packages/snaps-rpc-methods/CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- Add support for array `key` parameter in `snap_setState`, setting each key to its corresponding value in the provided object ([#4126](https://github.com/MetaMask/snaps/pull/4126))
+
 ## [17.1.2]
 
 ### Fixed

packages/snaps-rpc-methods/jest.config.js

@@ -10,10 +10,10 @@ module.exports = deepmerge(baseConfig, {
   ],
   coverageThreshold: {
     global: {
-      branches: 97.38,
+      branches: 97.43,
       functions: 98.92,
       lines: 99.22,
-      statements: 98.95,
+      statements: 98.96,
     },
   },
 });

packages/snaps-rpc-methods/src/permitted/setState.test.ts

@@ -522,6 +522,148 @@ describe('snap_setState', () => {
       });
     });
 
+    it('sets state for multiple keys', async () => {
+      const { implementation } = setStateHandler;
+
+      const getUnlockPromise = jest.fn().mockResolvedValue(undefined);
+      const hooks = { getUnlockPromise };
+
+      const messenger = getMessenger();
+
+      const engine = new JsonRpcEngine();
+
+      engine.push(createOriginMiddleware(MOCK_SNAP_ID));
+      engine.push((request, response, next, end) => {
+        const result = implementation(
+          request as JsonRpcRequestWithOrigin<SetStateParameters>,
+          response as PendingJsonRpcResponse<SetStateResult>,
+          next,
+          end,
+          hooks,
+          messenger,
+        );
+
+        result?.catch(end);
+      });
+
+      const response = await engine.handle({
+        jsonrpc: '2.0',
+        id: 1,
+        method: 'snap_setState',
+        params: {
+          key: ['foo', 'baz'],
+          value: { foo: 'newFoo', baz: 'newBaz' },
+        },
+      });
+
+      expect(response).toStrictEqual({
+        jsonrpc: '2.0',
+        id: 1,
+        result: null,
+      });
+
+      expect(messenger.call).toHaveBeenCalledWith(
+        'SnapController:updateSnapState',
+        MOCK_SNAP_ID,
+        { foo: 'newFoo', baz: 'newBaz' },
+        true,
+      );
+    });
+
+    it('sets missing keys to `null` when key is an array and value omits them', async () => {
+      const { implementation } = setStateHandler;
+
+      const getUnlockPromise = jest.fn().mockResolvedValue(undefined);
+      const hooks = { getUnlockPromise };
+
+      const messenger = getMessenger();
+
+      const engine = new JsonRpcEngine();
+
+      engine.push(createOriginMiddleware(MOCK_SNAP_ID));
+      engine.push((request, response, next, end) => {
+        const result = implementation(
+          request as JsonRpcRequestWithOrigin<SetStateParameters>,
+          response as PendingJsonRpcResponse<SetStateResult>,
+          next,
+          end,
+          hooks,
+          messenger,
+        );
+
+        result?.catch(end);
+      });
+
+      const response = await engine.handle({
+        jsonrpc: '2.0',
+        id: 1,
+        method: 'snap_setState',
+        params: {
+          key: ['foo', 'missing'],
+          value: { foo: 'newFoo' },
+        },
+      });
+
+      expect(response).toStrictEqual({
+        jsonrpc: '2.0',
+        id: 1,
+        result: null,
+      });
+
+      expect(messenger.call).toHaveBeenCalledWith(
+        'SnapController:updateSnapState',
+        MOCK_SNAP_ID,
+        { foo: 'newFoo', missing: null },
+        true,
+      );
+    });
+
+    it('throws if key is an array and value is not an object', async () => {
+      const { implementation } = setStateHandler;
+
+      const getUnlockPromise = jest.fn().mockResolvedValue(undefined);
+      const hooks = { getUnlockPromise };
+
+      const messenger = getMessenger();
+
+      const engine = new JsonRpcEngine();
+
+      engine.push(createOriginMiddleware(MOCK_SNAP_ID));
+      engine.push((request, response, next, end) => {
+        const result = implementation(
+          request as JsonRpcRequestWithOrigin<SetStateParameters>,
+          response,
+          next,
+          end,
+          hooks,
+          messenger,
+        );
+
+        result?.catch(end);
+      });
+
+      const response = await engine.handle({
+        jsonrpc: '2.0',
+        id: 1,
+        method: 'snap_setState',
+        params: {
+          key: ['foo', 'baz'],
+          value: 'not-an-object',
+        },
+      });
+
+      expect(response).toStrictEqual({
+        jsonrpc: '2.0',
+        id: 1,
+        error: {
+          code: errorCodes.rpc.invalidParams,
+          message:
+            'Invalid params: Value must be an object if key is an array.',
+          stack: expect.any(String),
+        },
+      });
+    });
+
     it('throws if the new state is not JSON serialisable', async () => {
       const { implementation } = setStateHandler;
 

packages/snaps-rpc-methods/src/permitted/setState.ts

@@ -5,10 +5,11 @@ import type {
 import type { Messenger } from '@metamask/messenger';
 import type { PermissionControllerHasPermissionAction } from '@metamask/permission-controller';
 import { providerErrors, rpcErrors } from '@metamask/rpc-errors';
-import type {
-  SetStateParams,
-  SetStateResult,
-  SnapId,
+import {
+  selectiveUnion,
+  type SetStateParams,
+  type SetStateResult,
+  type SnapId,
 } from '@metamask/snaps-sdk';
 import type { JsonObject } from '@metamask/snaps-sdk/jsx';
 import { getJsonSizeUnsafe, type InferMatching } from '@metamask/snaps-utils';
@@ -34,7 +35,7 @@ import type {
   SnapControllerUpdateSnapStateAction,
 } from '../types';
 import type { MethodHooksObject } from '../utils';
-import { FORBIDDEN_KEYS, StateKeyStruct } from '../utils';
+import { FORBIDDEN_KEYS, StateKeysStruct, StateKeyStruct } from '../utils';
 
 const hookNames: MethodHooksObject<SetStateMethodHooks> = {
   getUnlockPromise: true,
@@ -136,7 +137,14 @@ function getMutex(snapId: SnapId) {
 }
 
 const SetStateParametersStruct = objectStruct({
-  key: optional(StateKeyStruct),
+  key: optional(
+    selectiveUnion((value) => {
+      if (Array.isArray(value)) {
+        return StateKeysStruct;
+      }
+      return StateKeyStruct;
+    }),
+  ),
   value: JsonStruct,
   encrypted: optional(boolean()),
 });
@@ -191,6 +199,14 @@ async function setStateImplementation(
       );
     }
 
+    if (Array.isArray(key) && !isObject(value)) {
+      return end(
+        rpcErrors.invalidParams(
+          'Invalid params: Value must be an object if key is an array.',
+        ),
+      );
+    }
+
     if (encrypted) {
       await getUnlockPromise(true);
     }
@@ -267,20 +283,24 @@ function getValidatedParams(params?: unknown) {
  * If the key is `undefined`, the value is expected to be an object. In this
  * case, the value is returned as the new state.
  *
- * If the key is not `undefined`, the value is set in the state at the key. If
- * the key does not exist, it is created (and any missing intermediate keys are
- * created as well).
+ * If the key is a string, the value is set in the state at the key. If the key
+ * does not exist, it is created (and any missing intermediate keys are created
+ * as well).
+ *
+ * If the key is an array of strings, the value is expected to be an object
+ * mapping each key to its new value. Each key is set in the state.
  *
  * @param snapId - The Snap ID.
- * @param key - The key to set.
- * @param value - The value to set the key to.
+ * @param key - The key or keys to set.
+ * @param value - The value to set the key to. If `key` is an array, this must
+ * be an object mapping each key to its new value.
  * @param encrypted - Whether the state is encrypted.
  * @param messenger - The messenger used to call controller actions.
  * @returns The new state of the Snap.
  */
 async function getNewState(
   snapId: SnapId,
-  key: string | undefined,
+  key: string | string[] | undefined,
   value: Json,
   encrypted: boolean,
   messenger: Messenger<string, SetStateMethodActions>,
@@ -295,6 +315,21 @@ async function getNewState(
     snapId,
     encrypted,
   );
+
+  if (Array.isArray(key)) {
+    assert(isObject(value));
+    let newState = state;
+
+    // Intentionally using a classic for loop here for performance reasons.
+    // eslint-disable-next-line @typescript-eslint/prefer-for-of
+    for (let i = 0; i < key.length; i++) {
+      const currentKey = key[i];
+      newState = set(newState, currentKey, value[currentKey] ?? null);
+    }
+
+    return newState;
+  }
+
   return set(state, key, value);
 }
 

packages/snaps-rpc-methods/src/utils.ts

@@ -9,7 +9,7 @@ import { SLIP10Node } from '@metamask/key-tree';
 import type { Messenger } from '@metamask/messenger';
 import { rpcErrors } from '@metamask/rpc-errors';
 import type { MagicValue } from '@metamask/snaps-utils';
-import { refine, string } from '@metamask/superstruct';
+import { array, refine, string } from '@metamask/superstruct';
 import {
   assertExhaustive,
   add0x,
@@ -44,7 +44,7 @@ export type MethodHooksObject<HooksType extends Record<string, unknown>> = {
  * @returns The derived indices as a {@link HardenedBIP32Node} array.
  */
 function getDerivationPathArray(hash: Uint8Array): HardenedBIP32Node[] {
-  const array: HardenedBIP32Node[] = [];
+  const nodeArray: HardenedBIP32Node[] = [];
   const view = createDataView(hash);
 
   for (let index = 0; index < 8; index++) {
@@ -55,10 +55,10 @@ function getDerivationPathArray(hash: Uint8Array): HardenedBIP32Node[] {
     // the result is a positive number.
     // eslint-disable-next-line no-bitwise
     const pathIndex = (uint32 | HARDENED_VALUE) >>> 0;
-    array.push(`bip32:${pathIndex - HARDENED_VALUE}'` as const);
+    nodeArray.push(`bip32:${pathIndex - HARDENED_VALUE}'` as const);
   }
 
-  return array;
+  return nodeArray;
 }
 
 type BaseDeriveEntropyOptions = {
@@ -308,6 +308,8 @@ export const StateKeyStruct = refine(string(), 'state key', (value) => {
   return true;
 });
 
+export const StateKeysStruct = array(StateKeyStruct);
+
 /**
  * Get a value using the entropy source hooks: getMnemonic or getMnemonicSeed.
  * This function calls the passed hook and handles any errors that occur,

packages/snaps-sdk/CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- Add support for array `key` parameter in `SetStateParams` for `snap_setState` ([#4126](https://github.com/MetaMask/snaps/pull/4126))
+
 ## [12.0.1]
 
 ### Fixed

packages/snaps-sdk/src/types/methods/set-state.ts

@@ -3,10 +3,12 @@ import type { Json } from '@metamask/utils';
 /**
  * The request parameters for the `snap_setState` method.
  *
- * @property key - The key of the state to update. If not provided, the entire
- * state is updated. This may contain Lodash-style path syntax, for example,
- * `a.b.c`, with the exception of array syntax.
- * @property value - The value to set the state to.
+ * @property key - The key or keys of the state to update. If not provided, the
+ * entire state is updated. This may contain Lodash-style path syntax, for
+ * example, `a.b.c`, with the exception of array syntax. If an array of keys is
+ * provided, the value must be an object mapping each key to its new value.
+ * @property value - The value to set the state to. If `key` is an array, this
+ * must be an object mapping each key to its new value.
  * @property encrypted - Whether to use the separate encrypted state, or the
  * unencrypted state. Defaults to the encrypted state. Encrypted state can only
  * be used if the client is unlocked, while unencrypted state can be used
@@ -17,7 +19,7 @@ import type { Json } from '@metamask/utils';
  * while the client is locked.
  */
 export type SetStateParams = {
-  key?: string;
+  key?: string | string[];
   value: Json;
   encrypted?: boolean;
 };