feat: allow for multiple keys in snap_getState

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_getState` with array `key` parameter ([#4125](https://github.com/MetaMask/snaps/pull/4125))
+
 ## [3.0.0]
 
 ### Added

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

@@ -191,6 +191,51 @@ describe('onRpcRequest', () => {
         items: ['foo'],
       });
     });
+
+    it('returns state for multiple keys', async () => {
+      const { request } = await installSnap({
+        options: {
+          state: {
+            nested: { key: 'foo' },
+            items: ['bar'],
+          },
+        },
+      });
+
+      const response = await request({
+        method: 'getState',
+        params: {
+          key: ['nested.key', 'items'],
+        },
+      });
+
+      expect(response).toRespondWith({
+        'nested.key': 'foo',
+        items: ['bar'],
+      });
+    });
+
+    it('maps missing keys to `null` when an array of keys is provided', async () => {
+      const { request } = await installSnap({
+        options: {
+          state: {
+            items: ['foo'],
+          },
+        },
+      });
+
+      const response = await request({
+        method: 'getState',
+        params: {
+          key: ['items', 'missing'],
+        },
+      });
+
+      expect(response).toRespondWith({
+        items: ['foo'],
+        missing: null,
+      });
+    });
   });
 
   describe('clearState', () => {

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

@@ -1,6 +1,7 @@
 import {
   MethodNotFoundError,
   type OnRpcRequestHandler,
+  type GetStateParams,
 } from '@metamask/snaps-sdk';
 
 import type {
@@ -50,7 +51,7 @@ export const onRpcRequest: OnRpcRequestHandler = async ({ request }) => {
     }
 
     case 'getState': {
-      const params = request.params as BaseParams;
+      const params = request.params as GetStateParams;
       return await snap.request({
         method: 'snap_getState',
         params: {

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_getState`, returning a `Record<string, Json>` mapping each key to its resolved value ([#4125](https://github.com/MetaMask/snaps/pull/4125))
+
 ## [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.4,
       functions: 98.92,
       lines: 99.22,
-      statements: 98.95,
+      statements: 98.96,
     },
   },
 });

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

@@ -239,6 +239,97 @@ describe('snap_getState', () => {
       });
     });
 
+    it('returns the state for multiple keys', async () => {
+      const { implementation } = getStateHandler;
+
+      const getUnlockPromise = jest.fn().mockResolvedValue(undefined);
+      const hooks = { getUnlockPromise };
+
+      const messenger = getMessenger();
+
+      messenger.registerActionHandler(
+        'SnapController:getSnapState',
+        async () => ({ foo: 'bar', baz: 'qux' }),
+      );
+
+      const engine = new JsonRpcEngine();
+
+      engine.push(createOriginMiddleware(MOCK_SNAP_ID));
+      engine.push((request, response, next, end) => {
+        const result = implementation(
+          request as JsonRpcRequestWithOrigin<GetStateParameters>,
+          response,
+          next,
+          end,
+          hooks,
+          messenger,
+        );
+
+        result?.catch(end);
+      });
+
+      const response = await engine.handle({
+        jsonrpc: '2.0',
+        id: 1,
+        method: 'snap_getState',
+        params: {
+          key: ['foo', 'baz'],
+        },
+      });
+
+      expect(response).toStrictEqual({
+        jsonrpc: '2.0',
+        id: 1,
+        result: {
+          foo: 'bar',
+          baz: 'qux',
+        },
+      });
+    });
+
+    it('maps missing keys to `null` when an array of keys is provided', async () => {
+      const { implementation } = getStateHandler;
+
+      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<GetStateParameters>,
+          response,
+          next,
+          end,
+          hooks,
+          messenger,
+        );
+
+        result?.catch(end);
+      });
+
+      const response = await engine.handle({
+        jsonrpc: '2.0',
+        id: 1,
+        method: 'snap_getState',
+        params: {
+          key: ['foo', 'missing'],
+        },
+      });
+
+      expect(response).toStrictEqual({
+        jsonrpc: '2.0',
+        id: 1,
+        result: {
+          foo: 'bar',
+          missing: null,
+        },
+      });
+    });
+
     it('throws if the parameters are invalid', async () => {
       const { implementation } = getStateHandler;
 
@@ -322,4 +413,29 @@ describe('get', () => {
       'Invalid params: Key contains forbidden characters.',
     );
   });
+
+  it('returns a record of values when an array of keys is provided', () => {
+    const state = { a: { b: { c: 'value' } }, d: 'other' };
+    expect(get(state, ['a.b.c', 'd'])).toStrictEqual({
+      'a.b.c': 'value',
+      d: 'other',
+    });
+  });
+
+  it('returns an empty object when an empty array is provided', () => {
+    expect(get(object, [])).toStrictEqual({});
+  });
+
+  it('maps missing keys to `null` in array mode', () => {
+    expect(get(object, ['a.b.c', 'a.b.missing'])).toStrictEqual({
+      'a.b.c': 'value',
+      'a.b.missing': null,
+    });
+  });
+
+  it('throws if an array key contains a forbidden segment', () => {
+    expect(() => get(object, ['a.b.c', '__proto__.polluted'])).toThrow(
+      'Invalid params: Key contains forbidden characters.',
+    );
+  });
 });

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

@@ -5,7 +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 { GetStateParams, GetStateResult } from '@metamask/snaps-sdk';
+import {
+  selectiveUnion,
+  type GetStateParams,
+  type GetStateResult,
+} from '@metamask/snaps-sdk';
 import { type InferMatching } from '@metamask/snaps-utils';
 import {
   boolean,
@@ -23,7 +27,7 @@ import type {
   SnapControllerGetSnapStateAction,
 } from '../types';
 import type { MethodHooksObject } from '../utils';
-import { FORBIDDEN_KEYS, StateKeyStruct } from '../utils';
+import { FORBIDDEN_KEYS, StateKeysStruct, StateKeyStruct } from '../utils';
 
 const hookNames: MethodHooksObject<GetStateMethodHooks> = {
   getUnlockPromise: true,
@@ -82,7 +86,14 @@ export const getStateHandler = {
 >;
 
 const GetStateParametersStruct = object({
-  key: optional(StateKeyStruct),
+  key: optional(
+    selectiveUnion((value) => {
+      if (Array.isArray(value)) {
+        return StateKeysStruct;
+      }
+      return StateKeyStruct;
+    }),
+  ),
   encrypted: optional(boolean()),
 });
 
@@ -174,24 +185,39 @@ function getValidatedParams(params?: unknown) {
 /**
  * Get the value of a key in an object. The key may contain Lodash-style path
  * syntax, e.g., `a.b.c` (with the exception of array syntax). If the key does
- * not exist, `null` is returned.
+ * not exist, `null` is returned. If an array of keys is provided, the result
+ * is an object mapping each key to its value.
  *
  * This is a simplified version of Lodash's `get` function, but Lodash doesn't
  * seem to be maintained anymore, so we're using our own implementation.
  *
  * @param value - The object to get the key from.
- * @param key - The key to get.
+ * @param key - The key or keys to get.
  * @returns The value of the key in the object, or `null` if the key does not
- * exist.
+ * exist. If an array of keys is provided, returns a `Record` mapping each key
+ * to its value.
  */
 export function get(
   value: Record<string, Json> | null,
-  key?: string | undefined,
+  key?: string | string[] | undefined,
 ): Json {
   if (key === undefined) {
     return value;
   }
 
+  if (Array.isArray(key)) {
+    const result: Record<string, Json> = {};
+
+    // 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];
+      result[currentKey] = get(value, currentKey);
+    }
+
+    return result;
+  }
+
   const keys = key.split('.');
   let result: Json = 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 `GetStateParams` for `snap_getState` ([#4125](https://github.com/MetaMask/snaps/pull/4125))
+
 ## [12.0.1]
 
 ### Fixed

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

@@ -3,9 +3,10 @@ import type { Json } from '@metamask/utils';
 /**
  * An object containing the parameters for the `snap_getState` method.
  *
- * @property key - The key of the state to retrieve. If not provided, the entire
- * state is retrieved. This may contain Lodash-style path syntax, for example,
- * `a.b.c`, with the exception of array syntax.
+ * @property key - The key or keys of the state to retrieve. If not provided,
+ * the entire state is retrieved. 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 result is an object mapping each key to its 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
@@ -16,7 +17,7 @@ import type { Json } from '@metamask/utils';
  * while the client is locked.
  */
 export type GetStateParams = {
-  key?: string;
+  key?: string | string[];
   encrypted?: boolean;
 };
 

packages/test-snaps/CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Changed
+
+- Update `GetState` component to support comma-separated keys via `snap_getState` array `key` parameter ([#4125](https://github.com/MetaMask/snaps/pull/4125))
+
 ## [3.5.2]
 
 ### Fixed