Validate keys

MetaMask · Dec 10, 2024 · f6e08f2 · f6e08f2
1 parent ace7f38
commit f6e08f2
Show file tree

Hide file tree

Showing 4 changed files with 75 additions and 9 deletions.
diff --git a/packages/snaps-rpc-methods/src/permitted/getState.ts b/packages/snaps-rpc-methods/src/permitted/getState.ts
@@ -8,14 +8,14 @@ import {
   create,
   object,
   optional,
-  string,
   StructError,
 } from '@metamask/superstruct';
 import type { PendingJsonRpcResponse, JsonRpcRequest } from '@metamask/utils';
 import { hasProperty, isPlainObject, type Json } from '@metamask/utils';
 
 import { manageStateBuilder } from '../restricted/manageState';
 import type { MethodHooksObject } from '../utils';
+import { StateKeyStruct } from '../utils';
 
 const hookNames: MethodHooksObject<GetStateHooks> = {
   hasPermission: true,
@@ -64,7 +64,7 @@ export type GetStateHooks = {
 };
 
 const GetStateParametersStruct = object({
-  key: optional(string()),
+  key: optional(StateKeyStruct),
   encrypted: optional(boolean()),
 });
 

diff --git a/packages/snaps-rpc-methods/src/permitted/setState.ts b/packages/snaps-rpc-methods/src/permitted/setState.ts
@@ -2,20 +2,21 @@ import type { JsonRpcEngineEndCallback } from '@metamask/json-rpc-engine';
 import type { PermittedHandlerExport } from '@metamask/permission-controller';
 import { providerErrors, rpcErrors } from '@metamask/rpc-errors';
 import type { SetStateParams, SetStateResult } from '@metamask/snaps-sdk';
+import type { JsonObject } from '@metamask/snaps-sdk/jsx';
 import { type InferMatching } from '@metamask/snaps-utils';
 import {
   boolean,
   create,
   object as objectStruct,
   optional,
-  string,
   StructError,
 } from '@metamask/superstruct';
 import type { PendingJsonRpcResponse, JsonRpcRequest } from '@metamask/utils';
-import { JsonStruct, isPlainObject, type Json } from '@metamask/utils';
+import { assert, JsonStruct, isPlainObject, type Json } from '@metamask/utils';
 
 import { manageStateBuilder } from '../restricted/manageState';
 import type { MethodHooksObject } from '../utils';
+import { StateKeyStruct } from '../utils';
 
 const hookNames: MethodHooksObject<SetStateHooks> = {
   hasPermission: true,
@@ -80,7 +81,7 @@ export type SetStateHooks = {
 };
 
 const SetStateParametersStruct = objectStruct({
-  key: optional(string()),
+  key: optional(StateKeyStruct),
   value: JsonStruct,
   encrypted: optional(boolean()),
 });
@@ -193,8 +194,9 @@ export function set(
   object: Record<string, Json> | null,
   key: string | undefined,
   value: Json,
-): Json {
-  if (key === undefined || key === '') {
+): JsonObject {
+  if (key === undefined) {
+    assert(isPlainObject(value));
     return value;
   }
 
@@ -223,5 +225,5 @@ export function set(
   }
 
   // This should never be reached.
-  return null;
+  return {};
 }
diff --git a/packages/snaps-rpc-methods/src/utils.test.ts b/packages/snaps-rpc-methods/src/utils.test.ts
@@ -1,8 +1,15 @@
 import { SIP_6_MAGIC_VALUE } from '@metamask/snaps-utils';
 import { TEST_SECRET_RECOVERY_PHRASE_BYTES } from '@metamask/snaps-utils/test-utils';
+import { create, is } from '@metamask/superstruct';
 
 import { ENTROPY_VECTORS } from './__fixtures__';
-import { deriveEntropy, getNode, getPathPrefix } from './utils';
+import {
+  deriveEntropy,
+  getNode,
+  getPathPrefix,
+  isValidStateKey,
+  StateKeyStruct,
+} from './utils';
 
 describe('deriveEntropy', () => {
   it.each(ENTROPY_VECTORS)(
@@ -14,6 +21,7 @@ describe('deriveEntropy', () => {
           salt,
           mnemonicPhrase: TEST_SECRET_RECOVERY_PHRASE_BYTES,
           magic: SIP_6_MAGIC_VALUE,
+          cryptographicFunctions: {},
         }),
       ).toStrictEqual(entropy);
     },
@@ -47,6 +55,7 @@ describe('getNode', () => {
       curve: 'secp256k1',
       path: ['m', "44'", "1'"],
       secretRecoveryPhrase: TEST_SECRET_RECOVERY_PHRASE_BYTES,
+      cryptographicFunctions: {},
     });
 
     expect(node).toMatchInlineSnapshot(`
@@ -69,6 +78,7 @@ describe('getNode', () => {
       curve: 'ed25519',
       path: ['m', "44'", "1'"],
       secretRecoveryPhrase: TEST_SECRET_RECOVERY_PHRASE_BYTES,
+      cryptographicFunctions: {},
     });
 
     expect(node).toMatchInlineSnapshot(`
@@ -86,3 +96,34 @@ describe('getNode', () => {
     `);
   });
 });
+
+describe('isValidStateKey', () => {
+  it.each(['foo', 'foo.bar', 'foo.bar.baz'])(
+    'returns `true` for "%s"',
+    (key) => {
+      expect(isValidStateKey(key)).toBe(true);
+    },
+  );
+
+  it.each(['', '.', '..', 'foo.', 'foo..bar', 'foo.bar.', 'foo.bar..baz'])(
+    'returns `false` for "%s"',
+    (key) => {
+      expect(isValidStateKey(key)).toBe(false);
+    },
+  );
+});
+
+describe('StateKeyStruct', () => {
+  it.each(['foo', 'foo.bar', 'foo.bar.baz'])('accepts "%s"', (key) => {
+    expect(is(key, StateKeyStruct)).toBe(true);
+  });
+
+  it.each(['', '.', '..', 'foo.', 'foo..bar', 'foo.bar.', 'foo.bar..baz'])(
+    'does not accept "%s"',
+    (key) => {
+      expect(() => create(key, StateKeyStruct)).toThrow(
+        'Invalid state key. Each part of the key must be non-empty.',
+      );
+    },
+  );
+});
diff --git a/packages/snaps-rpc-methods/src/utils.ts b/packages/snaps-rpc-methods/src/utils.ts
@@ -7,6 +7,7 @@ import type {
 } from '@metamask/key-tree';
 import { SLIP10Node } from '@metamask/key-tree';
 import type { MagicValue } from '@metamask/snaps-utils';
+import { refine, string } from '@metamask/superstruct';
 import type { Hex } from '@metamask/utils';
 import {
   assertExhaustive,
@@ -238,3 +239,25 @@ export async function getNode({
     cryptographicFunctions,
   );
 }
+
+/**
+ * Validate the key of a state object.
+ *
+ * @param key - The key to validate.
+ * @returns `true` if the key is valid, `false` otherwise.
+ */
+export function isValidStateKey(key: string | undefined) {
+  if (key === undefined) {
+    return true;
+  }
+
+  return key.split('.').every((part) => part.length > 0);
+}
+
+export const StateKeyStruct = refine(string(), 'state key', (value) => {
+  if (!isValidStateKey(value)) {
+    return 'Invalid state key. Each part of the key must be non-empty.';
+  }
+
+  return true;
+});