From 16a30dc03ccbf5576d37615f70f09eb2207c8a7f Mon Sep 17 00:00:00 2001 From: Maarten Zuidhoorn Date: Fri, 29 Nov 2024 14:46:02 +0100 Subject: [PATCH 01/31] Add `snap_getState`, `snap_setState`, `snap_clearState` methods --- .../packages/manage-state/src/index.test.ts | 36 +-- .../packages/manage-state/src/index.ts | 42 +++- .../packages/manage-state/src/types.ts | 15 +- .../src/permitted/clearState.ts | 121 ++++++++++ .../src/permitted/getState.test.ts | 39 +++ .../src/permitted/getState.ts | 185 +++++++++++++++ .../src/permitted/handlers.ts | 6 + .../snaps-rpc-methods/src/permitted/index.ts | 6 +- .../src/permitted/setState.test.ts | 91 +++++++ .../src/permitted/setState.ts | 222 ++++++++++++++++++ .../src/types/methods/clear-state.ts | 16 ++ .../snaps-sdk/src/types/methods/get-state.ts | 22 ++ packages/snaps-sdk/src/types/methods/index.ts | 7 +- .../snaps-sdk/src/types/methods/methods.ts | 6 + .../snaps-sdk/src/types/methods/set-state.ts | 24 ++ packages/test-snaps/src/api.ts | 7 +- .../test-snaps/src/features/snaps/index.ts | 1 + .../manage-state/components/ClearData.tsx | 2 +- .../manage-state/components/SendData.tsx | 6 +- .../snaps/manage-state/hooks/useSnapState.ts | 2 +- .../src/features/snaps/state/State.tsx | 54 +++++ .../snaps/state/components/ClearState.tsx | 44 ++++ .../snaps/state/components/GetState.tsx | 66 ++++++ .../snaps/state/components/SetState.tsx | 90 +++++++ .../features/snaps/state/components/index.ts | 3 + .../src/features/snaps/state/constants.ts | 5 + .../src/features/snaps/state/hooks/index.ts | 1 + .../snaps/state/hooks/useSnapState.ts | 32 +++ .../src/features/snaps/state/index.ts | 1 + 29 files changed, 1117 insertions(+), 35 deletions(-) create mode 100644 packages/snaps-rpc-methods/src/permitted/clearState.ts create mode 100644 packages/snaps-rpc-methods/src/permitted/getState.test.ts create mode 100644 packages/snaps-rpc-methods/src/permitted/getState.ts create mode 100644 packages/snaps-rpc-methods/src/permitted/setState.test.ts create mode 100644 packages/snaps-rpc-methods/src/permitted/setState.ts create mode 100644 packages/snaps-sdk/src/types/methods/clear-state.ts create mode 100644 packages/snaps-sdk/src/types/methods/get-state.ts create mode 100644 packages/snaps-sdk/src/types/methods/set-state.ts create mode 100644 packages/test-snaps/src/features/snaps/state/State.tsx create mode 100644 packages/test-snaps/src/features/snaps/state/components/ClearState.tsx create mode 100644 packages/test-snaps/src/features/snaps/state/components/GetState.tsx create mode 100644 packages/test-snaps/src/features/snaps/state/components/SetState.tsx create mode 100644 packages/test-snaps/src/features/snaps/state/components/index.ts create mode 100644 packages/test-snaps/src/features/snaps/state/constants.ts create mode 100644 packages/test-snaps/src/features/snaps/state/hooks/index.ts create mode 100644 packages/test-snaps/src/features/snaps/state/hooks/useSnapState.ts create mode 100644 packages/test-snaps/src/features/snaps/state/index.ts diff --git a/packages/examples/packages/manage-state/src/index.test.ts b/packages/examples/packages/manage-state/src/index.test.ts index 2f86d99f65..7f0f8ed4e5 100644 --- a/packages/examples/packages/manage-state/src/index.test.ts +++ b/packages/examples/packages/manage-state/src/index.test.ts @@ -20,13 +20,13 @@ describe('onRpcRequest', () => { }); }); - describe('setState', () => { + describe('legacy_legacy_setState', () => { it('sets the state to the params', async () => { const { request } = await installSnap(); expect( await request({ - method: 'setState', + method: 'legacy_legacy_setState', params: { items: ['foo'], }, @@ -35,7 +35,7 @@ describe('onRpcRequest', () => { expect( await request({ - method: 'getState', + method: 'legacy_getState', }), ).toRespondWith({ items: ['foo'], @@ -47,7 +47,7 @@ describe('onRpcRequest', () => { expect( await request({ - method: 'setState', + method: 'legacy_legacy_setState', params: { items: ['foo'], encrypted: false, @@ -57,7 +57,7 @@ describe('onRpcRequest', () => { expect( await request({ - method: 'getState', + method: 'legacy_getState', }), ).toRespondWith({ items: [], @@ -65,7 +65,7 @@ describe('onRpcRequest', () => { expect( await request({ - method: 'getState', + method: 'legacy_getState', params: { encrypted: false, }, @@ -76,12 +76,12 @@ describe('onRpcRequest', () => { }); }); - describe('getState', () => { + describe('legacy_getState', () => { it('returns the state if no state has been set', async () => { const { request } = await installSnap(); const response = await request({ - method: 'getState', + method: 'legacy_getState', }); expect(response).toRespondWith({ @@ -93,14 +93,14 @@ describe('onRpcRequest', () => { const { request } = await installSnap(); await request({ - method: 'setState', + method: 'legacy_setState', params: { items: ['foo'], }, }); const response = await request({ - method: 'getState', + method: 'legacy_getState', }); expect(response).toRespondWith({ @@ -118,7 +118,7 @@ describe('onRpcRequest', () => { }); const response = await request({ - method: 'getState', + method: 'legacy_getState', params: { encrypted: false, }, @@ -130,12 +130,12 @@ describe('onRpcRequest', () => { }); }); - describe('clearState', () => { + describe('legacy_clearState', () => { it('clears the state', async () => { const { request } = await installSnap(); await request({ - method: 'setState', + method: 'legacy_setState', params: { items: ['foo'], }, @@ -143,13 +143,13 @@ describe('onRpcRequest', () => { expect( await request({ - method: 'clearState', + method: 'legacy_clearState', }), ).toRespondWith(true); expect( await request({ - method: 'getState', + method: 'legacy_getState', }), ).toRespondWith({ items: [], @@ -160,7 +160,7 @@ describe('onRpcRequest', () => { const { request } = await installSnap(); await request({ - method: 'setState', + method: 'legacy_setState', params: { items: ['foo'], encrypted: false, @@ -169,7 +169,7 @@ describe('onRpcRequest', () => { expect( await request({ - method: 'clearState', + method: 'legacy_clearState', params: { encrypted: false, }, @@ -178,7 +178,7 @@ describe('onRpcRequest', () => { expect( await request({ - method: 'getState', + method: 'legacy_getState', params: { encrypted: false, }, diff --git a/packages/examples/packages/manage-state/src/index.ts b/packages/examples/packages/manage-state/src/index.ts index 32c5aa32ef..b8d92a98a2 100644 --- a/packages/examples/packages/manage-state/src/index.ts +++ b/packages/examples/packages/manage-state/src/index.ts @@ -3,7 +3,7 @@ import { type OnRpcRequestHandler, } from '@metamask/snaps-sdk'; -import type { BaseParams, SetStateParams } from './types'; +import type { BaseParams, LegacySetStateParams, SetStateParams } from './types'; import { clearState, getState, setState } from './utils'; /** @@ -34,21 +34,55 @@ export const onRpcRequest: OnRpcRequestHandler = async ({ request }) => { switch (request.method) { case 'setState': { const params = request.params as SetStateParams; + return await snap.request({ + method: 'snap_setState', + params: { + key: params?.key, + value: params?.value, + encrypted: params?.encrypted, + }, + }); + } + case 'getState': { + const params = request.params as BaseParams; + return await snap.request({ + method: 'snap_getState', + params: { + key: params?.key, + encrypted: params?.encrypted, + }, + }); + } + + case 'clearState': { + const params = request.params as BaseParams; + return await snap.request({ + method: 'snap_clearState', + params: { + encrypted: params?.encrypted, + }, + }); + } + + case 'legacy_setState': { + const params = request.params as LegacySetStateParams; if (params.items) { await setState({ items: params.items }, params.encrypted); } + return true; } - case 'getState': { + case 'legacy_getState': { const params = request.params as BaseParams; - return await getState(params?.encrypted); + return await getState(params.encrypted); } - case 'clearState': { + case 'legacy_clearState': { const params = request.params as BaseParams; await clearState(params?.encrypted); + return true; } diff --git a/packages/examples/packages/manage-state/src/types.ts b/packages/examples/packages/manage-state/src/types.ts index 4f696071b7..c433b4b867 100644 --- a/packages/examples/packages/manage-state/src/types.ts +++ b/packages/examples/packages/manage-state/src/types.ts @@ -1,10 +1,17 @@ +import type { Json } from '@metamask/utils'; + import type { State } from './utils'; -export type BaseParams = { encrypted?: boolean }; +export type BaseParams = { key?: string; encrypted?: boolean }; /** * The parameters for the `setState` JSON-RPC method. - * - * The current state will be merged with the new state. */ -export type SetStateParams = BaseParams & Partial; +export type SetStateParams = BaseParams & { + value: Json; +}; + +/** + * The parameters for the `legacy_setState` JSON-RPC method. + */ +export type LegacySetStateParams = BaseParams & Partial; diff --git a/packages/snaps-rpc-methods/src/permitted/clearState.ts b/packages/snaps-rpc-methods/src/permitted/clearState.ts new file mode 100644 index 0000000000..176b6b8d12 --- /dev/null +++ b/packages/snaps-rpc-methods/src/permitted/clearState.ts @@ -0,0 +1,121 @@ +import type { JsonRpcEngineEndCallback } from '@metamask/json-rpc-engine'; +import type { PermittedHandlerExport } from '@metamask/permission-controller'; +import { providerErrors, rpcErrors } from '@metamask/rpc-errors'; +import type { ClearStateParams, ClearStateResult } from '@metamask/snaps-sdk'; +import { type InferMatching } from '@metamask/snaps-utils'; +import { + boolean, + create, + object, + optional, + StructError, +} from '@metamask/superstruct'; +import type { PendingJsonRpcResponse, JsonRpcRequest } from '@metamask/utils'; + +import type { MethodHooksObject } from '../utils'; + +const hookNames: MethodHooksObject = { + clearSnapState: true, + hasPermission: true, +}; + +/** + * `snap_clearState` clears the state of the Snap. + */ +export const clearStateHandler: PermittedHandlerExport< + ClearStateHooks, + ClearStateParameters, + ClearStateResult +> = { + methodNames: ['snap_clearState'], + implementation: clearStateImplementation, + hookNames, +}; + +export type ClearStateHooks = { + /** + * A function that clears the state of the requesting Snap. + */ + clearSnapState: (snapId: string, encrypted: boolean) => void; + + /** + * Check if the requesting origin has a given permission. + * + * @param permissionName - The name of the permission to check. + * @returns Whether the origin has the permission. + */ + hasPermission: (permissionName: string) => boolean; +}; + +const ClearStateParametersStruct = object({ + encrypted: optional(boolean()), +}); + +export type ClearStateParameters = InferMatching< + typeof ClearStateParametersStruct, + ClearStateParams +>; + +/** + * The `snap_clearState` method implementation. + * + * @param request - The JSON-RPC request object. + * @param response - The JSON-RPC response object. + * @param _next - The `json-rpc-engine` "next" callback. Not used by this + * function. + * @param end - The `json-rpc-engine` "end" callback. + * @param hooks - The RPC method hooks. + * @param hooks.clearSnapState - A function that clears the state of the + * requesting Snap. + * @param hooks.hasPermission - Check whether a given origin has a given + * permission. + * @returns Nothing. + */ +async function clearStateImplementation( + request: JsonRpcRequest, + response: PendingJsonRpcResponse, + _next: unknown, + end: JsonRpcEngineEndCallback, + { clearSnapState, hasPermission }: ClearStateHooks, +): Promise { + const { params } = request; + + if (!hasPermission('snap_manageState')) { + return end(providerErrors.unauthorized()); + } + + try { + const validatedParams = getValidatedParams(params); + const { encrypted = true } = validatedParams; + + // We expect the MM middleware stack to always add the origin to requests + const { origin } = request as JsonRpcRequest & { origin: string }; + clearSnapState(origin, encrypted); + + response.result = null; + } catch (error) { + return end(error); + } + + return end(); +} + +/** + * Validate the parameters of the `snap_clearState` method. + * + * @param params - The parameters to validate. + * @returns The validated parameters. + */ +function getValidatedParams(params?: unknown) { + try { + return create(params, ClearStateParametersStruct); + } catch (error) { + if (error instanceof StructError) { + throw rpcErrors.invalidParams({ + message: `Invalid params: ${error.message}.`, + }); + } + + throw error; + } +} diff --git a/packages/snaps-rpc-methods/src/permitted/getState.test.ts b/packages/snaps-rpc-methods/src/permitted/getState.test.ts new file mode 100644 index 0000000000..8fccc8cf40 --- /dev/null +++ b/packages/snaps-rpc-methods/src/permitted/getState.test.ts @@ -0,0 +1,39 @@ +import { get } from './getState'; + +describe('get', () => { + const object = { + a: { + b: { + c: 'value', + }, + }, + }; + + it('returns the value of the key', () => { + expect(get(object, 'a.b.c')).toBe('value'); + }); + + it('returns the object if the key is empty', () => { + expect(get(object, '')).toBe(object); + }); + + it('returns `null` if the object is `null`', () => { + expect(get(null, '')).toBeNull(); + }); + + it('returns `null` if the key does not exist', () => { + expect(get(object, 'a.b.d')).toBeNull(); + }); + + it('returns `null` if the parent key is not an object', () => { + expect(get(object, 'a.b.c.d')).toBeNull(); + }); + + it('returns `null` if the key is a prototype pollution attempt', () => { + expect(get(object, '__proto__.polluted')).toBeNull(); + }); + + it('returns `null` if the key is a constructor pollution attempt', () => { + expect(get(object, 'constructor.polluted')).toBeNull(); + }); +}); diff --git a/packages/snaps-rpc-methods/src/permitted/getState.ts b/packages/snaps-rpc-methods/src/permitted/getState.ts new file mode 100644 index 0000000000..9376a3215d --- /dev/null +++ b/packages/snaps-rpc-methods/src/permitted/getState.ts @@ -0,0 +1,185 @@ +import type { JsonRpcEngineEndCallback } from '@metamask/json-rpc-engine'; +import type { PermittedHandlerExport } from '@metamask/permission-controller'; +import { providerErrors, rpcErrors } from '@metamask/rpc-errors'; +import type { GetStateParams, GetStateResult } from '@metamask/snaps-sdk'; +import { type InferMatching } from '@metamask/snaps-utils'; +import { + boolean, + create, + object, + optional, + string, + StructError, +} from '@metamask/superstruct'; +import type { PendingJsonRpcResponse, JsonRpcRequest } from '@metamask/utils'; +import { hasProperty, isPlainObject, type Json } from '@metamask/utils'; + +import type { MethodHooksObject } from '../utils'; + +const hookNames: MethodHooksObject = { + hasPermission: true, + getSnapState: true, + getUnlockPromise: true, +}; + +/** + * `snap_getState` gets the state of the Snap. + */ +export const getStateHandler: PermittedHandlerExport< + GetStateHooks, + GetStateParameters, + GetStateResult +> = { + methodNames: ['snap_getState'], + implementation: getStateImplementation, + hookNames, +}; + +export type GetStateHooks = { + /** + * Check if the requesting origin has a given permission. + * + * @param permissionName - The name of the permission to check. + * @returns Whether the origin has the permission. + */ + hasPermission: (permissionName: string) => boolean; + + /** + * Get the state of the requesting Snap. + * + * @returns The current state of the Snap. + */ + getSnapState: ( + snapId: string, + encrypted: boolean, + ) => Promise>; + + /** + * Wait for the extension to be unlocked. + * + * @returns A promise that resolves once the extension is unlocked. + */ + getUnlockPromise: (shouldShowUnlockRequest: boolean) => Promise; +}; + +const GetStateParametersStruct = object({ + key: optional(string()), + encrypted: optional(boolean()), +}); + +export type GetStateParameters = InferMatching< + typeof GetStateParametersStruct, + GetStateParams +>; + +/** + * The `snap_getState` method implementation. + * + * @param request - The JSON-RPC request object. + * @param response - The JSON-RPC response object. + * @param _next - The `json-rpc-engine` "next" callback. Not used by this + * function. + * @param end - The `json-rpc-engine` "end" callback. + * @param hooks - The RPC method hooks. + * @param hooks.hasPermission - Check whether a given origin has a given + * permission. + * @param hooks.getSnapState - Get the state of the requesting Snap. + * @param hooks.getUnlockPromise - Wait for the extension to be unlocked. + * @returns Nothing. + */ +async function getStateImplementation( + request: JsonRpcRequest, + response: PendingJsonRpcResponse, + _next: unknown, + end: JsonRpcEngineEndCallback, + { hasPermission, getSnapState, getUnlockPromise }: GetStateHooks, +): Promise { + const { params } = request; + + if (!hasPermission('snap_manageState')) { + return end(providerErrors.unauthorized()); + } + + try { + const validatedParams = getValidatedParams(params); + const { key, encrypted = true } = validatedParams; + + if (encrypted) { + await getUnlockPromise(true); + } + + // We expect the MM middleware stack to always add the origin to requests + const { origin } = request as JsonRpcRequest & { origin: string }; + const state = await getSnapState(origin, encrypted); + + response.result = get(state, key); + } catch (error) { + return end(error); + } + + return end(); +} + +/** + * Validate the parameters of the `snap_getState` method. + * + * @param params - The parameters to validate. + * @returns The validated parameters. + */ +function getValidatedParams(params?: unknown) { + try { + return create(params, GetStateParametersStruct); + } catch (error) { + if (error instanceof StructError) { + throw rpcErrors.invalidParams({ + message: `Invalid params: ${error.message}.`, + }); + } + + throw error; + } +} + +/** + * 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. + * + * 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. + * @returns The value of the key in the object, or `null` if the key does not + * exist. + */ +export function get( + value: Record | null, + key?: string | undefined, +): Json { + if (key === undefined || key === '') { + return value; + } + + const keys = key.split('.'); + let result: Json = value; + + for (const currentKey of keys) { + if (['__proto__', 'constructor'].includes(currentKey)) { + return null; + } + + if (isPlainObject(result)) { + if (!hasProperty(result, currentKey)) { + return null; + } + + result = result[currentKey]; + continue; + } + + return null; + } + + return result; +} diff --git a/packages/snaps-rpc-methods/src/permitted/handlers.ts b/packages/snaps-rpc-methods/src/permitted/handlers.ts index 83730b1a15..5e0c6f201d 100644 --- a/packages/snaps-rpc-methods/src/permitted/handlers.ts +++ b/packages/snaps-rpc-methods/src/permitted/handlers.ts @@ -1,3 +1,4 @@ +import { clearStateHandler } from './clearState'; import { createInterfaceHandler } from './createInterface'; import { providerRequestHandler } from './experimentalProviderRequest'; import { getAllSnapsHandler } from './getAllSnaps'; @@ -7,10 +8,12 @@ import { getFileHandler } from './getFile'; import { getInterfaceContextHandler } from './getInterfaceContext'; import { getInterfaceStateHandler } from './getInterfaceState'; import { getSnapsHandler } from './getSnaps'; +import { getStateHandler } from './getState'; import { invokeKeyringHandler } from './invokeKeyring'; import { invokeSnapSugarHandler } from './invokeSnapSugar'; import { requestSnapsHandler } from './requestSnaps'; import { resolveInterfaceHandler } from './resolveInterface'; +import { setStateHandler } from './setState'; import { updateInterfaceHandler } from './updateInterface'; /* eslint-disable @typescript-eslint/naming-convention */ @@ -20,8 +23,10 @@ export const methodHandlers = { wallet_requestSnaps: requestSnapsHandler, wallet_invokeSnap: invokeSnapSugarHandler, wallet_invokeKeyring: invokeKeyringHandler, + snap_clearState: clearStateHandler, snap_getClientStatus: getClientStatusHandler, snap_getFile: getFileHandler, + snap_getState: getStateHandler, snap_createInterface: createInterfaceHandler, snap_updateInterface: updateInterfaceHandler, snap_getInterfaceState: getInterfaceStateHandler, @@ -29,6 +34,7 @@ export const methodHandlers = { snap_resolveInterface: resolveInterfaceHandler, snap_getCurrencyRate: getCurrencyRateHandler, snap_experimentalProviderRequest: providerRequestHandler, + snap_setState: setStateHandler, }; /* eslint-enable @typescript-eslint/naming-convention */ diff --git a/packages/snaps-rpc-methods/src/permitted/index.ts b/packages/snaps-rpc-methods/src/permitted/index.ts index 5aa676fce3..cadbf563fc 100644 --- a/packages/snaps-rpc-methods/src/permitted/index.ts +++ b/packages/snaps-rpc-methods/src/permitted/index.ts @@ -5,20 +5,24 @@ import type { GetClientStatusHooks } from './getClientStatus'; import type { GetCurrencyRateMethodHooks } from './getCurrencyRate'; import type { GetInterfaceStateMethodHooks } from './getInterfaceState'; import type { GetSnapsHooks } from './getSnaps'; +import type { GetStateHooks } from './getState'; import type { RequestSnapsHooks } from './requestSnaps'; import type { ResolveInterfaceMethodHooks } from './resolveInterface'; +import type { SetStateHooks } from './setState'; import type { UpdateInterfaceMethodHooks } from './updateInterface'; export type PermittedRpcMethodHooks = GetAllSnapsHooks & GetClientStatusHooks & GetSnapsHooks & + GetStateHooks & RequestSnapsHooks & CreateInterfaceMethodHooks & UpdateInterfaceMethodHooks & GetInterfaceStateMethodHooks & ResolveInterfaceMethodHooks & GetCurrencyRateMethodHooks & - ProviderRequestMethodHooks; + ProviderRequestMethodHooks & + SetStateHooks; export * from './handlers'; export * from './middleware'; diff --git a/packages/snaps-rpc-methods/src/permitted/setState.test.ts b/packages/snaps-rpc-methods/src/permitted/setState.test.ts new file mode 100644 index 0000000000..e509a06f25 --- /dev/null +++ b/packages/snaps-rpc-methods/src/permitted/setState.test.ts @@ -0,0 +1,91 @@ +import { set } from './setState'; + +describe('set', () => { + it('sets the state in an empty object', () => { + const object = {}; + + expect(set(object, 'key', 'value')).toStrictEqual({ + key: 'value', + }); + }); + + it('sets the state if the current state is `null`', () => { + const object = null; + + expect(set(object, 'key', 'value')).toStrictEqual({ + key: 'value', + }); + }); + + it('sets the state in an empty object with a nested key', () => { + const object = {}; + + expect(set(object, 'nested.key', 'newValue')).toStrictEqual({ + nested: { + key: 'newValue', + }, + }); + }); + + it('sets the state in an existing object', () => { + const object = { + key: 'oldValue', + }; + + expect(set(object, 'key', 'newValue')).toStrictEqual({ + key: 'newValue', + }); + }); + + it('sets the state in an existing object with a nested key', () => { + const object = { + nested: { + key: 'oldValue', + }, + }; + + expect(set(object, 'nested.key', 'newValue')).toStrictEqual({ + nested: { + key: 'newValue', + }, + }); + }); + + it('sets the state in an existing object with a nested key that does not exist', () => { + const object = { + nested: {}, + }; + + expect(set(object, 'nested.key', 'newValue')).toStrictEqual({ + nested: { + key: 'newValue', + }, + }); + }); + + it('overwrites the state in an existing object', () => { + const object = { + nested: { + key: 'oldValue', + }, + }; + + expect(set(object, undefined, { foo: 'bar' })).toStrictEqual({ + foo: 'bar', + }); + }); + + it('overwrites the nested state in an existing object', () => { + const object = { + nested: { + key: 'oldValue', + }, + }; + + expect(set(object, 'nested', { foo: 'bar' })).toStrictEqual({ + nested: { + foo: 'bar', + }, + }); + }); +}); diff --git a/packages/snaps-rpc-methods/src/permitted/setState.ts b/packages/snaps-rpc-methods/src/permitted/setState.ts new file mode 100644 index 0000000000..73e6cff2fa --- /dev/null +++ b/packages/snaps-rpc-methods/src/permitted/setState.ts @@ -0,0 +1,222 @@ +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 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 type { MethodHooksObject } from '../utils'; + +const hookNames: MethodHooksObject = { + hasPermission: true, + getSnapState: true, + getUnlockPromise: true, + updateSnapState: true, +}; + +/** + * `snap_setState` sets the state of the Snap. + */ +export const setStateHandler: PermittedHandlerExport< + SetStateHooks, + SetStateParameters, + SetStateResult +> = { + methodNames: ['snap_setState'], + implementation: setStateImplementation, + hookNames, +}; + +export type SetStateHooks = { + /** + * Check if the requesting origin has a given permission. + * + * @param permissionName - The name of the permission to check. + * @returns Whether the origin has the permission. + */ + hasPermission: (permissionName: string) => boolean; + + /** + * Get the state of the requesting Snap. + * + * @param snapId - The ID of the Snap. + * @param encrypted - Whether the state is encrypted. + * @returns The current state of the Snap. + */ + getSnapState: ( + snapId: string, + encrypted: boolean, + ) => Promise>; + + /** + * Wait for the extension to be unlocked. + * + * @returns A promise that resolves once the extension is unlocked. + */ + getUnlockPromise: (shouldShowUnlockRequest: boolean) => Promise; + + /** + * Update the state of the requesting Snap. + * + * @param snapId - The ID of the Snap. + * @param newState - The new state of the Snap. + * @param encrypted - Whether the state should be encrypted. + */ + updateSnapState: ( + snapId: string, + newState: Record, + encrypted: boolean, + ) => Promise; +}; + +const SetStateParametersStruct = objectStruct({ + key: optional(string()), + value: JsonStruct, + encrypted: optional(boolean()), +}); + +export type SetStateParameters = InferMatching< + typeof SetStateParametersStruct, + SetStateParams +>; + +/** + * The `snap_setState` method implementation. + * + * @param request - The JSON-RPC request object. + * @param response - The JSON-RPC response object. + * @param _next - The `json-rpc-engine` "next" callback. Not used by this + * function. + * @param end - The `json-rpc-engine` "end" callback. + * @param hooks - The RPC method hooks. + * @param hooks.hasPermission - Check whether a given origin has a given + * permission. + * @param hooks.getSnapState - Get the state of the requesting Snap. + * @param hooks.getUnlockPromise - Wait for the extension to be unlocked. + * @param hooks.updateSnapState - Update the state of the requesting Snap. + * @returns Nothing. + */ +async function setStateImplementation( + request: JsonRpcRequest, + response: PendingJsonRpcResponse, + _next: unknown, + end: JsonRpcEngineEndCallback, + { + hasPermission, + getSnapState, + getUnlockPromise, + updateSnapState, + }: SetStateHooks, +): Promise { + const { params } = request; + + if (!hasPermission('snap_manageState')) { + return end(providerErrors.unauthorized()); + } + + try { + const validatedParams = getValidatedParams(params); + const { key, value, encrypted = true } = validatedParams; + + if (encrypted) { + await getUnlockPromise(true); + } + + // We expect the MM middleware stack to always add the origin to requests + const { origin } = request as JsonRpcRequest & { origin: string }; + const state = await getSnapState(origin, encrypted); + + const newState = set(state, key, value); + if (!isPlainObject(newState)) { + return end( + rpcErrors.invalidParams('Invalid state: Root must be an object.'), + ); + } + + await updateSnapState(origin, newState, encrypted); + response.result = null; + } catch (error) { + return end(error); + } + + return end(); +} + +/** + * Validate the parameters of the `snap_setState` method. + * + * @param params - The parameters to validate. + * @returns The validated parameters. + */ +function getValidatedParams(params?: unknown) { + try { + return create(params, SetStateParametersStruct); + } catch (error) { + if (error instanceof StructError) { + throw rpcErrors.invalidParams({ + message: `Invalid params: ${error.message}.`, + }); + } + + throw error; + } +} + +/** + * Set 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, it is created (and any missing intermediate keys are created as + * well). + * + * This is a simplified version of Lodash's `set` function, but Lodash doesn't + * seem to be maintained anymore, so we're using our own implementation. + * + * @param object - The object to get the key from. + * @param key - The key to set. + * @param value - The value to set the key to. + * @returns The new object with the key set to the value. + */ +export function set( + // eslint-disable-next-line @typescript-eslint/default-param-last + object: Record | null, + key: string | undefined, + value: Json, +): Json { + if (key === undefined || key === '') { + return value; + } + + const keys = key.split('.'); + const requiredObject = object ?? {}; + let currentObject: Record = requiredObject; + + for (let i = 0; i < keys.length; i++) { + const currentKey = keys[i]; + if (['__proto__', 'constructor'].includes(currentKey)) { + return {}; + } + + if (i === keys.length - 1) { + currentObject[currentKey] = value; + return requiredObject; + } + + currentObject[currentKey] = isPlainObject(currentObject[currentKey]) + ? currentObject[currentKey] + : {}; + + currentObject = currentObject[currentKey] as Record; + } + + // This should never be reached. + return null; +} diff --git a/packages/snaps-sdk/src/types/methods/clear-state.ts b/packages/snaps-sdk/src/types/methods/clear-state.ts new file mode 100644 index 0000000000..0db3edc778 --- /dev/null +++ b/packages/snaps-sdk/src/types/methods/clear-state.ts @@ -0,0 +1,16 @@ +/** + * The request parameters for the `snap_clearState` method. + * + * @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 extension is unlocked, while unencrypted state can be used + * whether the extension is locked or unlocked. + */ +export type ClearStateParams = { + encrypted?: boolean; +}; + +/** + * The result returned by the `snap_clearState` method. + */ +export type ClearStateResult = null; diff --git a/packages/snaps-sdk/src/types/methods/get-state.ts b/packages/snaps-sdk/src/types/methods/get-state.ts new file mode 100644 index 0000000000..1692e7ea52 --- /dev/null +++ b/packages/snaps-sdk/src/types/methods/get-state.ts @@ -0,0 +1,22 @@ +import type { Json } from '@metamask/utils'; + +/** + * The request 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, e.g., + * `a.b.c`, with the exception of array syntax. + * @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 extension is unlocked, while unencrypted state can be used + * whether the extension is locked or unlocked. + */ +export type GetStateParams = { + key?: string; + encrypted?: boolean; +}; + +/** + * The result returned by the `snap_getState` method. + */ +export type GetStateResult = Json; diff --git a/packages/snaps-sdk/src/types/methods/index.ts b/packages/snaps-sdk/src/types/methods/index.ts index 02d604df85..70bd559285 100644 --- a/packages/snaps-sdk/src/types/methods/index.ts +++ b/packages/snaps-sdk/src/types/methods/index.ts @@ -1,9 +1,11 @@ +export * from './clear-state'; export * from './create-interface'; export * from './dialog'; export * from './get-bip32-entropy'; export * from './get-bip32-public-key'; export * from './get-bip44-entropy'; export * from './get-client-status'; +export * from './get-currency-rate'; export * from './get-entropy'; export * from './get-file'; export * from './get-interface-context'; @@ -11,14 +13,15 @@ export * from './get-interface-state'; export * from './get-locale'; export * from './get-preferences'; export * from './get-snaps'; +export * from './get-state'; export * from './invoke-keyring'; export * from './invoke-snap'; export * from './manage-accounts'; export * from './manage-state'; export * from './methods'; export * from './notify'; +export * from './provider-request'; export * from './request-snaps'; export * from './update-interface'; export * from './resolve-interface'; -export * from './get-currency-rate'; -export * from './provider-request'; +export * from './set-state'; diff --git a/packages/snaps-sdk/src/types/methods/methods.ts b/packages/snaps-sdk/src/types/methods/methods.ts index cf671a6254..21386d0fd8 100644 --- a/packages/snaps-sdk/src/types/methods/methods.ts +++ b/packages/snaps-sdk/src/types/methods/methods.ts @@ -1,4 +1,5 @@ import type { Method } from '../../internals'; +import type { ClearStateParams, ClearStateResult } from './clear-state'; import type { CreateInterfaceParams, CreateInterfaceResult, @@ -40,6 +41,7 @@ import type { GetPreferencesResult, } from './get-preferences'; import type { GetSnapsParams, GetSnapsResult } from './get-snaps'; +import type { GetStateParams, GetStateResult } from './get-state'; import type { InvokeKeyringParams, InvokeKeyringResult, @@ -56,6 +58,7 @@ import type { ResolveInterfaceParams, ResolveInterfaceResult, } from './resolve-interface'; +import type { SetStateParams, SetStateResult } from './set-state'; import type { UpdateInterfaceParams, UpdateInterfaceResult, @@ -67,6 +70,7 @@ import type { */ export type SnapMethods = { /* eslint-disable @typescript-eslint/naming-convention */ + snap_clearState: [ClearStateParams, ClearStateResult]; snap_dialog: [DialogParams, DialogResult]; snap_getBip32Entropy: [GetBip32EntropyParams, GetBip32EntropyResult]; snap_getBip32PublicKey: [GetBip32PublicKeyParams, GetBip32PublicKeyResult]; @@ -77,6 +81,7 @@ export type SnapMethods = { snap_getFile: [GetFileParams, GetFileResult]; snap_getLocale: [GetLocaleParams, GetLocaleResult]; snap_getPreferences: [GetPreferencesParams, GetPreferencesResult]; + snap_getState: [GetStateParams, GetStateResult]; snap_manageAccounts: [ManageAccountsParams, ManageAccountsResult]; snap_manageState: [ManageStateParams, ManageStateResult]; snap_notify: [NotifyParams, NotifyResult]; @@ -88,6 +93,7 @@ export type SnapMethods = { GetInterfaceContextResult, ]; snap_resolveInterface: [ResolveInterfaceParams, ResolveInterfaceResult]; + snap_setState: [SetStateParams, SetStateResult]; wallet_getSnaps: [GetSnapsParams, GetSnapsResult]; wallet_invokeKeyring: [InvokeKeyringParams, InvokeKeyringResult]; wallet_invokeSnap: [InvokeSnapParams, InvokeSnapResult]; diff --git a/packages/snaps-sdk/src/types/methods/set-state.ts b/packages/snaps-sdk/src/types/methods/set-state.ts new file mode 100644 index 0000000000..d5c30df46c --- /dev/null +++ b/packages/snaps-sdk/src/types/methods/set-state.ts @@ -0,0 +1,24 @@ +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, e.g., + * `a.b.c`, with the exception of array syntax. + * @property value - The value to set the state to. + * @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 extension is unlocked, while unencrypted state can be used + * whether the extension is locked or unlocked. + */ +export type SetStateParams = { + key?: string; + value: Json; + encrypted?: boolean; +}; + +/** + * The result returned by the `snap_setState` method. + */ +export type SetStateResult = null; diff --git a/packages/test-snaps/src/api.ts b/packages/test-snaps/src/api.ts index 913a98de23..4c4f458d6e 100644 --- a/packages/test-snaps/src/api.ts +++ b/packages/test-snaps/src/api.ts @@ -1,5 +1,7 @@ -import type { MetaMaskInpageProvider } from '@metamask/providers'; -import type { RequestArguments } from '@metamask/providers/dist/BaseProvider'; +import type { + MetaMaskInpageProvider, + RequestArguments, +} from '@metamask/providers'; import { logError } from '@metamask/snaps-utils'; import type { JsonRpcError, JsonRpcParams } from '@metamask/utils'; import type { BaseQueryFn } from '@reduxjs/toolkit/query/react'; @@ -118,6 +120,7 @@ export const baseApi = createApi({ }), invalidatesTags: [Tag.InstalledSnaps], }), + installSnaps: build.mutation({ query: (params) => ({ method: 'wallet_requestSnaps', diff --git a/packages/test-snaps/src/features/snaps/index.ts b/packages/test-snaps/src/features/snaps/index.ts index c80978a40d..d563710530 100644 --- a/packages/test-snaps/src/features/snaps/index.ts +++ b/packages/test-snaps/src/features/snaps/index.ts @@ -27,3 +27,4 @@ export * from './signature-insights'; export * from './updates'; export * from './wasm'; export * from './send-flow'; +export * from './state'; diff --git a/packages/test-snaps/src/features/snaps/manage-state/components/ClearData.tsx b/packages/test-snaps/src/features/snaps/manage-state/components/ClearData.tsx index d250e03022..32ea16a493 100644 --- a/packages/test-snaps/src/features/snaps/manage-state/components/ClearData.tsx +++ b/packages/test-snaps/src/features/snaps/manage-state/components/ClearData.tsx @@ -15,7 +15,7 @@ export const ClearData: FunctionComponent<{ encrypted: boolean }> = ({ const handleClick = () => { invokeSnap({ snapId: getSnapId(MANAGE_STATE_SNAP_ID, MANAGE_STATE_PORT), - method: 'clearState', + method: 'legacy_clearState', params: { encrypted }, tags: [encrypted ? Tag.TestState : Tag.UnencryptedTestState], }).catch(logError); diff --git a/packages/test-snaps/src/features/snaps/manage-state/components/SendData.tsx b/packages/test-snaps/src/features/snaps/manage-state/components/SendData.tsx index 2903b4ce7f..cb3949c817 100644 --- a/packages/test-snaps/src/features/snaps/manage-state/components/SendData.tsx +++ b/packages/test-snaps/src/features/snaps/manage-state/components/SendData.tsx @@ -22,11 +22,13 @@ export const SendData: FunctionComponent<{ encrypted: boolean }> = ({ const handleSubmit = (event: FormEvent) => { event.preventDefault(); + const items = snapState?.items ?? []; + invokeSnap({ snapId: getSnapId(MANAGE_STATE_SNAP_ID, MANAGE_STATE_PORT), - method: 'setState', + method: 'legacy_setState', params: { - items: [...snapState.items, value], + items: [...items, value], encrypted, }, tags: [encrypted ? Tag.TestState : Tag.UnencryptedTestState], diff --git a/packages/test-snaps/src/features/snaps/manage-state/hooks/useSnapState.ts b/packages/test-snaps/src/features/snaps/manage-state/hooks/useSnapState.ts index af53efa433..d0139f9bc1 100644 --- a/packages/test-snaps/src/features/snaps/manage-state/hooks/useSnapState.ts +++ b/packages/test-snaps/src/features/snaps/manage-state/hooks/useSnapState.ts @@ -19,7 +19,7 @@ export function useSnapState(encrypted: boolean) { const { data: state } = useInvokeQuery<{ data: State }>( { snapId, - method: 'getState', + method: 'legacy_getState', params: { encrypted }, tags: [encrypted ? Tag.TestState : Tag.UnencryptedTestState], }, diff --git a/packages/test-snaps/src/features/snaps/state/State.tsx b/packages/test-snaps/src/features/snaps/state/State.tsx new file mode 100644 index 0000000000..65b7f40daa --- /dev/null +++ b/packages/test-snaps/src/features/snaps/state/State.tsx @@ -0,0 +1,54 @@ +import type { FunctionComponent } from 'react'; + +import { Result, Snap } from '../../../components'; +import { ClearState, GetState, SetState } from './components'; +import { + MANAGE_STATE_SNAP_ID, + MANAGE_STATE_PORT, + MANAGE_STATE_VERSION, +} from './constants'; +import { useSnapState } from './hooks'; + +export const State: FunctionComponent = () => { + const encryptedState = useSnapState(true); + const unencryptedState = useSnapState(false); + + return ( + + + 
+          {JSON.stringify(encryptedState, null, 2)}
+
+ + + + + + + + 
+          {JSON.stringify(unencryptedState, null, 2)}
+
+ + + + + + ); +}; diff --git a/packages/test-snaps/src/features/snaps/state/components/ClearState.tsx b/packages/test-snaps/src/features/snaps/state/components/ClearState.tsx new file mode 100644 index 0000000000..4923db7fae --- /dev/null +++ b/packages/test-snaps/src/features/snaps/state/components/ClearState.tsx @@ -0,0 +1,44 @@ +import { logError } from '@metamask/snaps-utils'; +import type { FunctionComponent } from 'react'; +import { Button } from 'react-bootstrap'; + +import { Tag, useInvokeMutation } from '../../../../api'; +import { Result } from '../../../../components'; +import { getSnapId } from '../../../../utils'; +import { MANAGE_STATE_PORT, MANAGE_STATE_SNAP_ID } from '../constants'; + +export const ClearState: FunctionComponent<{ encrypted: boolean }> = ({ + encrypted, +}) => { + const [invokeSnap, { isLoading, data, error }] = useInvokeMutation(); + + const handleClick = () => { + invokeSnap({ + snapId: getSnapId(MANAGE_STATE_SNAP_ID, MANAGE_STATE_PORT), + method: 'clearState', + params: { encrypted }, + tags: [encrypted ? Tag.TestState : Tag.UnencryptedTestState], + }).catch(logError); + }; + + return ( + <> + + + + {JSON.stringify(data, null, 2)} + {JSON.stringify(error, null, 2)} + + + + ); +}; diff --git a/packages/test-snaps/src/features/snaps/state/components/GetState.tsx b/packages/test-snaps/src/features/snaps/state/components/GetState.tsx new file mode 100644 index 0000000000..7d4b925b35 --- /dev/null +++ b/packages/test-snaps/src/features/snaps/state/components/GetState.tsx @@ -0,0 +1,66 @@ +import { logError } from '@metamask/snaps-utils'; +import type { ChangeEvent, FormEvent, FunctionComponent } from 'react'; +import { useState } from 'react'; +import { Button, Form } from 'react-bootstrap'; + +import { Tag, useInvokeMutation } from '../../../../api'; +import { Result } from '../../../../components'; +import { getSnapId } from '../../../../utils'; +import { MANAGE_STATE_PORT, MANAGE_STATE_SNAP_ID } from '../constants'; + +export const GetState: FunctionComponent<{ encrypted: boolean }> = ({ + encrypted, +}) => { + const [key, setKey] = useState(''); + const [invokeSnap, { isLoading, data, error }] = useInvokeMutation(); + + const handleChange = (event: ChangeEvent) => { + setKey(event.target.value); + }; + + const handleSubmit = (event: FormEvent) => { + event.preventDefault(); + invokeSnap({ + snapId: getSnapId(MANAGE_STATE_SNAP_ID, MANAGE_STATE_PORT), + method: 'getState', + params: { + key, + encrypted, + }, + tags: [encrypted ? Tag.TestState : Tag.UnencryptedTestState], + }).catch(logError); + }; + + return ( + <> +
+ + Key + + + + +
+ + + + {JSON.stringify(data, null, 2)} + {JSON.stringify(error, null, 2)} + + + + ); +}; diff --git a/packages/test-snaps/src/features/snaps/state/components/SetState.tsx b/packages/test-snaps/src/features/snaps/state/components/SetState.tsx new file mode 100644 index 0000000000..03bdc9744a --- /dev/null +++ b/packages/test-snaps/src/features/snaps/state/components/SetState.tsx @@ -0,0 +1,90 @@ +import { logError } from '@metamask/snaps-utils'; +import type { ChangeEvent, FormEvent, FunctionComponent } from 'react'; +import { useState } from 'react'; +import { Button, Form } from 'react-bootstrap'; + +import { Tag, useInvokeMutation } from '../../../../api'; +import { Result } from '../../../../components'; +import { getSnapId } from '../../../../utils'; +import { MANAGE_STATE_PORT, MANAGE_STATE_SNAP_ID } from '../constants'; + +export const SetState: FunctionComponent<{ encrypted: boolean }> = ({ + encrypted, +}) => { + const [key, setKey] = useState(''); + const [value, setValue] = useState(''); + const [invokeSnap, { isLoading, data, error }] = useInvokeMutation(); + + const handleChangeKey = (event: ChangeEvent) => { + setKey(event.target.value); + }; + + const handleChangeValue = (event: ChangeEvent) => { + setValue(event.target.value); + }; + + const handleSubmit = (event: FormEvent) => { + event.preventDefault(); + invokeSnap({ + snapId: getSnapId(MANAGE_STATE_SNAP_ID, MANAGE_STATE_PORT), + method: 'setState', + params: { + key, + value: JSON.parse(value), + encrypted, + }, + tags: [encrypted ? Tag.TestState : Tag.UnencryptedTestState], + }).catch(logError); + }; + + return ( + <> +
+ + Key + + + + + Value + + + + +
+ + + + {JSON.stringify(data, null, 2)} + {JSON.stringify(error, null, 2)} + + + + ); +}; diff --git a/packages/test-snaps/src/features/snaps/state/components/index.ts b/packages/test-snaps/src/features/snaps/state/components/index.ts new file mode 100644 index 0000000000..3f13fdefea --- /dev/null +++ b/packages/test-snaps/src/features/snaps/state/components/index.ts @@ -0,0 +1,3 @@ +export * from './ClearState'; +export * from './GetState'; +export * from './SetState'; diff --git a/packages/test-snaps/src/features/snaps/state/constants.ts b/packages/test-snaps/src/features/snaps/state/constants.ts new file mode 100644 index 0000000000..cd2622fdbc --- /dev/null +++ b/packages/test-snaps/src/features/snaps/state/constants.ts @@ -0,0 +1,5 @@ +import packageJson from '@metamask/manage-state-example-snap/package.json'; + +export const MANAGE_STATE_SNAP_ID = 'npm:@metamask/manage-state-example-snap'; +export const MANAGE_STATE_PORT = 8014; +export const MANAGE_STATE_VERSION = packageJson.version; diff --git a/packages/test-snaps/src/features/snaps/state/hooks/index.ts b/packages/test-snaps/src/features/snaps/state/hooks/index.ts new file mode 100644 index 0000000000..3a6999aa10 --- /dev/null +++ b/packages/test-snaps/src/features/snaps/state/hooks/index.ts @@ -0,0 +1 @@ +export * from './useSnapState'; diff --git a/packages/test-snaps/src/features/snaps/state/hooks/useSnapState.ts b/packages/test-snaps/src/features/snaps/state/hooks/useSnapState.ts new file mode 100644 index 0000000000..af53efa433 --- /dev/null +++ b/packages/test-snaps/src/features/snaps/state/hooks/useSnapState.ts @@ -0,0 +1,32 @@ +import { Tag, useInvokeQuery } from '../../../../api'; +import { getSnapId, useInstalled } from '../../../../utils'; +import { MANAGE_STATE_PORT, MANAGE_STATE_SNAP_ID } from '../constants'; + +export type State = { + items: string[]; +}; + +/** + * Hook to retrieve the state of the snap. + * + * @param encrypted - A flag to indicate whether to use encrypted storage or not. + * @returns The state of the snap. + */ +export function useSnapState(encrypted: boolean) { + const snapId = getSnapId(MANAGE_STATE_SNAP_ID, MANAGE_STATE_PORT); + const isInstalled = useInstalled(snapId); + + const { data: state } = useInvokeQuery<{ data: State }>( + { + snapId, + method: 'getState', + params: { encrypted }, + tags: [encrypted ? Tag.TestState : Tag.UnencryptedTestState], + }, + { + skip: !isInstalled, + }, + ); + + return state; +} diff --git a/packages/test-snaps/src/features/snaps/state/index.ts b/packages/test-snaps/src/features/snaps/state/index.ts new file mode 100644 index 0000000000..e11bf86ccf --- /dev/null +++ b/packages/test-snaps/src/features/snaps/state/index.ts @@ -0,0 +1 @@ +export * from './State'; From 6000f185c15d7213e9ee7e51ece0b3d55f2301a8 Mon Sep 17 00:00:00 2001 From: Maarten Zuidhoorn Date: Fri, 29 Nov 2024 14:55:23 +0100 Subject: [PATCH 02/31] Add missing dependency to manage state example --- .../packages/manage-state/package.json | 3 ++- yarn.lock | 18 ++++++++++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/packages/examples/packages/manage-state/package.json b/packages/examples/packages/manage-state/package.json index bca541d171..6a7065d795 100644 --- a/packages/examples/packages/manage-state/package.json +++ b/packages/examples/packages/manage-state/package.json @@ -43,7 +43,8 @@ "test:watch": "jest --watch" }, "dependencies": { - "@metamask/snaps-sdk": "workspace:^" + "@metamask/snaps-sdk": "workspace:^", + "@metamask/utils": "^11.0.0" }, "devDependencies": { "@jest/globals": "^29.5.0", diff --git a/yarn.lock b/yarn.lock index 5e73b80771..17aa0af215 100644 --- a/yarn.lock +++ b/yarn.lock @@ -5144,6 +5144,7 @@ __metadata: "@metamask/snaps-cli": "workspace:^" "@metamask/snaps-jest": "workspace:^" "@metamask/snaps-sdk": "workspace:^" + "@metamask/utils": "npm:^11.0.0" "@swc/core": "npm:1.3.78" "@swc/jest": "npm:^0.2.26" "@typescript-eslint/eslint-plugin": "npm:^5.42.1" @@ -6475,6 +6476,23 @@ __metadata: languageName: node linkType: hard +"@metamask/utils@npm:^11.0.0": + version: 11.0.0 + resolution: "@metamask/utils@npm:11.0.0" + dependencies: + "@ethereumjs/tx": "npm:^4.2.0" + "@metamask/superstruct": "npm:^3.1.0" + "@noble/hashes": "npm:^1.3.1" + "@scure/base": "npm:^1.1.3" + "@types/debug": "npm:^4.1.7" + debug: "npm:^4.3.4" + pony-cause: "npm:^2.1.10" + semver: "npm:^7.5.4" + uuid: "npm:^9.0.1" + checksum: 10/a836fcbbc1e88af9c514f643bf5c6b85a9c9f433a6f43184e3c8e8d464390c6eafd788737904564c86b4ac81c4d3e0f070408ce6ae5b71f6abbd59a94f6137de + languageName: node + linkType: hard + "@metamask/utils@npm:^9.0.0, @metamask/utils@npm:^9.1.0": version: 9.2.1 resolution: "@metamask/utils@npm:9.2.1" From 4a336c298dcaf6d5e60147ada7440bc700db57db Mon Sep 17 00:00:00 2001 From: Maarten Zuidhoorn Date: Fri, 29 Nov 2024 15:01:15 +0100 Subject: [PATCH 03/31] Fix example tests --- packages/examples/packages/manage-state/src/index.test.ts | 6 +++--- packages/examples/packages/manage-state/src/index.ts | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/packages/examples/packages/manage-state/src/index.test.ts b/packages/examples/packages/manage-state/src/index.test.ts index 7f0f8ed4e5..f2a61df6f6 100644 --- a/packages/examples/packages/manage-state/src/index.test.ts +++ b/packages/examples/packages/manage-state/src/index.test.ts @@ -20,13 +20,13 @@ describe('onRpcRequest', () => { }); }); - describe('legacy_legacy_setState', () => { + describe('legacy_setState', () => { it('sets the state to the params', async () => { const { request } = await installSnap(); expect( await request({ - method: 'legacy_legacy_setState', + method: 'legacy_setState', params: { items: ['foo'], }, @@ -47,7 +47,7 @@ describe('onRpcRequest', () => { expect( await request({ - method: 'legacy_legacy_setState', + method: 'legacy_setState', params: { items: ['foo'], encrypted: false, diff --git a/packages/examples/packages/manage-state/src/index.ts b/packages/examples/packages/manage-state/src/index.ts index b8d92a98a2..12e1cb24c4 100644 --- a/packages/examples/packages/manage-state/src/index.ts +++ b/packages/examples/packages/manage-state/src/index.ts @@ -76,7 +76,7 @@ export const onRpcRequest: OnRpcRequestHandler = async ({ request }) => { case 'legacy_getState': { const params = request.params as BaseParams; - return await getState(params.encrypted); + return await getState(params?.encrypted); } case 'legacy_clearState': { From d98f4279d24a9de5974cf2265cb8bfea073114a9 Mon Sep 17 00:00:00 2001 From: Maarten Zuidhoorn Date: Tue, 10 Dec 2024 10:01:52 +0100 Subject: [PATCH 04/31] Throw for invalid keys --- .../snaps-rpc-methods/src/permitted/getState.test.ts | 10 +++++++--- packages/snaps-rpc-methods/src/permitted/getState.ts | 4 +++- .../snaps-rpc-methods/src/permitted/setState.test.ts | 12 ++++++++++++ packages/snaps-rpc-methods/src/permitted/setState.ts | 4 +++- 4 files changed, 25 insertions(+), 5 deletions(-) diff --git a/packages/snaps-rpc-methods/src/permitted/getState.test.ts b/packages/snaps-rpc-methods/src/permitted/getState.test.ts index 8fccc8cf40..67a5df5a8d 100644 --- a/packages/snaps-rpc-methods/src/permitted/getState.test.ts +++ b/packages/snaps-rpc-methods/src/permitted/getState.test.ts @@ -29,11 +29,15 @@ describe('get', () => { expect(get(object, 'a.b.c.d')).toBeNull(); }); - it('returns `null` if the key is a prototype pollution attempt', () => { - expect(get(object, '__proto__.polluted')).toBeNull(); + it('throws an error if the key is a prototype pollution attempt', () => { + expect(() => get(object, '__proto__.polluted')).toThrow( + 'Invalid params: Key contains forbidden characters.', + ); }); it('returns `null` if the key is a constructor pollution attempt', () => { - expect(get(object, 'constructor.polluted')).toBeNull(); + expect(() => get(object, 'constructor.polluted')).toThrow( + 'Invalid params: Key contains forbidden characters.', + ); }); }); diff --git a/packages/snaps-rpc-methods/src/permitted/getState.ts b/packages/snaps-rpc-methods/src/permitted/getState.ts index 9376a3215d..f4efe3a976 100644 --- a/packages/snaps-rpc-methods/src/permitted/getState.ts +++ b/packages/snaps-rpc-methods/src/permitted/getState.ts @@ -166,7 +166,9 @@ export function get( for (const currentKey of keys) { if (['__proto__', 'constructor'].includes(currentKey)) { - return null; + throw rpcErrors.invalidParams( + 'Invalid params: Key contains forbidden characters.', + ); } if (isPlainObject(result)) { diff --git a/packages/snaps-rpc-methods/src/permitted/setState.test.ts b/packages/snaps-rpc-methods/src/permitted/setState.test.ts index e509a06f25..2433e52015 100644 --- a/packages/snaps-rpc-methods/src/permitted/setState.test.ts +++ b/packages/snaps-rpc-methods/src/permitted/setState.test.ts @@ -88,4 +88,16 @@ describe('set', () => { }, }); }); + + it('throws an error if the key is a prototype pollution attempt', () => { + expect(() => set({}, '__proto__.polluted', 'value')).toThrow( + 'Invalid params: Key contains forbidden characters.', + ); + }); + + it('throws an error if the key is a constructor pollution attempt', () => { + expect(() => set({}, 'constructor.polluted', 'value')).toThrow( + 'Invalid params: Key contains forbidden characters.', + ); + }); }); diff --git a/packages/snaps-rpc-methods/src/permitted/setState.ts b/packages/snaps-rpc-methods/src/permitted/setState.ts index 73e6cff2fa..3de458ff9e 100644 --- a/packages/snaps-rpc-methods/src/permitted/setState.ts +++ b/packages/snaps-rpc-methods/src/permitted/setState.ts @@ -202,7 +202,9 @@ export function set( for (let i = 0; i < keys.length; i++) { const currentKey = keys[i]; if (['__proto__', 'constructor'].includes(currentKey)) { - return {}; + throw rpcErrors.invalidParams( + 'Invalid params: Key contains forbidden characters.', + ); } if (i === keys.length - 1) { From 78658fcf8e38494b2368908cbf5c93574d084943 Mon Sep 17 00:00:00 2001 From: Maarten Zuidhoorn Date: Tue, 10 Dec 2024 10:02:25 +0100 Subject: [PATCH 05/31] Fix `@metamask/utils` version --- .../packages/manage-state/package.json | 2 +- yarn.lock | 19 +------------------ 2 files changed, 2 insertions(+), 19 deletions(-) diff --git a/packages/examples/packages/manage-state/package.json b/packages/examples/packages/manage-state/package.json index 6a7065d795..343524d21d 100644 --- a/packages/examples/packages/manage-state/package.json +++ b/packages/examples/packages/manage-state/package.json @@ -44,7 +44,7 @@ }, "dependencies": { "@metamask/snaps-sdk": "workspace:^", - "@metamask/utils": "^11.0.0" + "@metamask/utils": "^10.0.0" }, "devDependencies": { "@jest/globals": "^29.5.0", diff --git a/yarn.lock b/yarn.lock index 17aa0af215..0b7a429e29 100644 --- a/yarn.lock +++ b/yarn.lock @@ -5144,7 +5144,7 @@ __metadata: "@metamask/snaps-cli": "workspace:^" "@metamask/snaps-jest": "workspace:^" "@metamask/snaps-sdk": "workspace:^" - "@metamask/utils": "npm:^11.0.0" + "@metamask/utils": "npm:^10.0.0" "@swc/core": "npm:1.3.78" "@swc/jest": "npm:^0.2.26" "@typescript-eslint/eslint-plugin": "npm:^5.42.1" @@ -6476,23 +6476,6 @@ __metadata: languageName: node linkType: hard -"@metamask/utils@npm:^11.0.0": - version: 11.0.0 - resolution: "@metamask/utils@npm:11.0.0" - dependencies: - "@ethereumjs/tx": "npm:^4.2.0" - "@metamask/superstruct": "npm:^3.1.0" - "@noble/hashes": "npm:^1.3.1" - "@scure/base": "npm:^1.1.3" - "@types/debug": "npm:^4.1.7" - debug: "npm:^4.3.4" - pony-cause: "npm:^2.1.10" - semver: "npm:^7.5.4" - uuid: "npm:^9.0.1" - checksum: 10/a836fcbbc1e88af9c514f643bf5c6b85a9c9f433a6f43184e3c8e8d464390c6eafd788737904564c86b4ac81c4d3e0f070408ce6ae5b71f6abbd59a94f6137de - languageName: node - linkType: hard - "@metamask/utils@npm:^9.0.0, @metamask/utils@npm:^9.1.0": version: 9.2.1 resolution: "@metamask/utils@npm:9.2.1" From 4b1576a2f168b47dc8c9b30b60a1b4a14d6e6f84 Mon Sep 17 00:00:00 2001 From: Maarten Zuidhoorn Date: Tue, 10 Dec 2024 10:05:19 +0100 Subject: [PATCH 06/31] Use `snap_manageState`'s target name instead of hardcoding permission name --- packages/snaps-rpc-methods/src/permitted/clearState.ts | 3 ++- packages/snaps-rpc-methods/src/permitted/getState.ts | 3 ++- packages/snaps-rpc-methods/src/permitted/setState.ts | 3 ++- 3 files changed, 6 insertions(+), 3 deletions(-) diff --git a/packages/snaps-rpc-methods/src/permitted/clearState.ts b/packages/snaps-rpc-methods/src/permitted/clearState.ts index 176b6b8d12..bac8ff19f2 100644 --- a/packages/snaps-rpc-methods/src/permitted/clearState.ts +++ b/packages/snaps-rpc-methods/src/permitted/clearState.ts @@ -12,6 +12,7 @@ import { } from '@metamask/superstruct'; import type { PendingJsonRpcResponse, JsonRpcRequest } from '@metamask/utils'; +import { manageStateBuilder } from '../restricted/manageState'; import type { MethodHooksObject } from '../utils'; const hookNames: MethodHooksObject = { @@ -80,7 +81,7 @@ async function clearStateImplementation( ): Promise { const { params } = request; - if (!hasPermission('snap_manageState')) { + if (!hasPermission(manageStateBuilder.targetName)) { return end(providerErrors.unauthorized()); } diff --git a/packages/snaps-rpc-methods/src/permitted/getState.ts b/packages/snaps-rpc-methods/src/permitted/getState.ts index f4efe3a976..c5f83a7d52 100644 --- a/packages/snaps-rpc-methods/src/permitted/getState.ts +++ b/packages/snaps-rpc-methods/src/permitted/getState.ts @@ -14,6 +14,7 @@ import { 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'; const hookNames: MethodHooksObject = { @@ -96,7 +97,7 @@ async function getStateImplementation( ): Promise { const { params } = request; - if (!hasPermission('snap_manageState')) { + if (!hasPermission(manageStateBuilder.targetName)) { return end(providerErrors.unauthorized()); } diff --git a/packages/snaps-rpc-methods/src/permitted/setState.ts b/packages/snaps-rpc-methods/src/permitted/setState.ts index 3de458ff9e..a77ddd9668 100644 --- a/packages/snaps-rpc-methods/src/permitted/setState.ts +++ b/packages/snaps-rpc-methods/src/permitted/setState.ts @@ -14,6 +14,7 @@ import { import type { PendingJsonRpcResponse, JsonRpcRequest } from '@metamask/utils'; import { JsonStruct, isPlainObject, type Json } from '@metamask/utils'; +import { manageStateBuilder } from '../restricted/manageState'; import type { MethodHooksObject } from '../utils'; const hookNames: MethodHooksObject = { @@ -119,7 +120,7 @@ async function setStateImplementation( ): Promise { const { params } = request; - if (!hasPermission('snap_manageState')) { + if (!hasPermission(manageStateBuilder.targetName)) { return end(providerErrors.unauthorized()); } From 978eeaee21bb74071cf56fba21d0a867df7896f7 Mon Sep 17 00:00:00 2001 From: Maarten Zuidhoorn Date: Tue, 10 Dec 2024 10:07:54 +0100 Subject: [PATCH 07/31] Move check for invalid state --- packages/snaps-rpc-methods/src/permitted/setState.ts | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/packages/snaps-rpc-methods/src/permitted/setState.ts b/packages/snaps-rpc-methods/src/permitted/setState.ts index a77ddd9668..6e747918ca 100644 --- a/packages/snaps-rpc-methods/src/permitted/setState.ts +++ b/packages/snaps-rpc-methods/src/permitted/setState.ts @@ -136,13 +136,15 @@ async function setStateImplementation( const { origin } = request as JsonRpcRequest & { origin: string }; const state = await getSnapState(origin, encrypted); - const newState = set(state, key, value); - if (!isPlainObject(newState)) { + if (key === undefined && !isPlainObject(value)) { return end( - rpcErrors.invalidParams('Invalid state: Root must be an object.'), + rpcErrors.invalidParams( + 'Invalid params: Value must be an object if key is not provided.', + ), ); } + const newState = set(state, key, value); await updateSnapState(origin, newState, encrypted); response.result = null; } catch (error) { From 9d71f21d11abb7f453d8e8bbf40f4f35c1dcd7ce Mon Sep 17 00:00:00 2001 From: Maarten Zuidhoorn Date: Tue, 10 Dec 2024 10:43:05 +0100 Subject: [PATCH 08/31] Validate keys --- .../src/permitted/getState.ts | 4 +- .../src/permitted/setState.ts | 14 +++--- packages/snaps-rpc-methods/src/utils.test.ts | 43 ++++++++++++++++++- packages/snaps-rpc-methods/src/utils.ts | 23 ++++++++++ 4 files changed, 75 insertions(+), 9 deletions(-) diff --git a/packages/snaps-rpc-methods/src/permitted/getState.ts b/packages/snaps-rpc-methods/src/permitted/getState.ts index c5f83a7d52..0018959e84 100644 --- a/packages/snaps-rpc-methods/src/permitted/getState.ts +++ b/packages/snaps-rpc-methods/src/permitted/getState.ts @@ -8,7 +8,6 @@ import { create, object, optional, - string, StructError, } from '@metamask/superstruct'; import type { PendingJsonRpcResponse, JsonRpcRequest } from '@metamask/utils'; @@ -16,6 +15,7 @@ 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 = { 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 index 6e747918ca..f28463412d 100644 --- 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 = { 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 | 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 index 6241b90e4e..211be807ef 100644 --- 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 index 500038ff5b..57f8e67ed3 100644 --- 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; +}); From b61509c6e70a6bc702aef5bc76a9f0b98b437834 Mon Sep 17 00:00:00 2001 From: Maarten Zuidhoorn Date: Tue, 10 Dec 2024 14:02:55 +0100 Subject: [PATCH 09/31] Add tests for new methods --- packages/snaps-rpc-methods/jest.config.js | 8 +- .../src/permitted/clearState.test.ts | 207 ++++++++++ .../src/permitted/clearState.ts | 3 +- .../src/permitted/getState.test.ts | 281 +++++++++++++- .../src/permitted/getState.ts | 3 +- .../src/permitted/setState.test.ts | 365 +++++++++++++++++- .../src/permitted/setState.ts | 4 +- 7 files changed, 862 insertions(+), 9 deletions(-) create mode 100644 packages/snaps-rpc-methods/src/permitted/clearState.test.ts diff --git a/packages/snaps-rpc-methods/jest.config.js b/packages/snaps-rpc-methods/jest.config.js index f3813d88c0..918d7696ca 100644 --- a/packages/snaps-rpc-methods/jest.config.js +++ b/packages/snaps-rpc-methods/jest.config.js @@ -10,10 +10,10 @@ module.exports = deepmerge(baseConfig, { ], coverageThreshold: { global: { - branches: 92.94, - functions: 97.29, - lines: 97.88, - statements: 97.41, + branches: 92.95, + functions: 97.44, + lines: 97.93, + statements: 97.52, }, }, }); diff --git a/packages/snaps-rpc-methods/src/permitted/clearState.test.ts b/packages/snaps-rpc-methods/src/permitted/clearState.test.ts new file mode 100644 index 0000000000..1154341659 --- /dev/null +++ b/packages/snaps-rpc-methods/src/permitted/clearState.test.ts @@ -0,0 +1,207 @@ +import { JsonRpcEngine } from '@metamask/json-rpc-engine'; +import { errorCodes } from '@metamask/rpc-errors'; +import type { ClearStateResult } from '@metamask/snaps-sdk'; +import { MOCK_SNAP_ID } from '@metamask/snaps-utils/test-utils'; +import type { JsonRpcRequest, PendingJsonRpcResponse } from '@metamask/utils'; + +import type { ClearStateParameters } from './clearState'; +import { clearStateHandler } from './clearState'; + +describe('snap_clearState', () => { + describe('clearStateHandler', () => { + it('has the expected shape', () => { + expect(clearStateHandler).toMatchObject({ + methodNames: ['snap_clearState'], + implementation: expect.any(Function), + hookNames: { + clearSnapState: true, + hasPermission: true, + }, + }); + }); + }); + + describe('implementation', () => { + const createOriginMiddleware = + (origin: string) => + (request: any, _response: unknown, next: () => void, _end: unknown) => { + request.origin = origin; + next(); + }; + + it('returns the result from the `clearSnapState` hook', async () => { + const { implementation } = clearStateHandler; + + const clearSnapState = jest.fn().mockReturnValue(null); + const hasPermission = jest.fn().mockReturnValue(true); + + const hooks = { + clearSnapState, + hasPermission, + }; + + const engine = new JsonRpcEngine(); + + engine.push(createOriginMiddleware(MOCK_SNAP_ID)); + engine.push((request, response, next, end) => { + const result = implementation( + request as JsonRpcRequest, + response as PendingJsonRpcResponse, + next, + end, + hooks, + ); + + result?.catch(end); + }); + + const response = await engine.handle({ + jsonrpc: '2.0', + id: 1, + method: 'snap_clearState', + params: {}, + }); + + expect(clearSnapState).toHaveBeenCalledWith(MOCK_SNAP_ID, true); + expect(response).toStrictEqual({ + jsonrpc: '2.0', + id: 1, + result: null, + }); + }); + + it('clears unencrypted state if specified', async () => { + const { implementation } = clearStateHandler; + + const clearSnapState = jest.fn().mockReturnValue(null); + const hasPermission = jest.fn().mockReturnValue(true); + + const hooks = { + clearSnapState, + hasPermission, + }; + + const engine = new JsonRpcEngine(); + + engine.push(createOriginMiddleware(MOCK_SNAP_ID)); + engine.push((request, response, next, end) => { + const result = implementation( + request as JsonRpcRequest, + response as PendingJsonRpcResponse, + next, + end, + hooks, + ); + + result?.catch(end); + }); + + const response = await engine.handle({ + jsonrpc: '2.0', + id: 1, + method: 'snap_clearState', + params: { + encrypted: false, + }, + }); + + expect(clearSnapState).toHaveBeenCalledWith(MOCK_SNAP_ID, false); + expect(response).toStrictEqual({ + jsonrpc: '2.0', + id: 1, + result: null, + }); + }); + + it('throws if the requesting origin does not have the required permission', async () => { + const { implementation } = clearStateHandler; + + const clearSnapState = jest.fn(); + const hasPermission = jest.fn().mockReturnValue(false); + + const hooks = { + clearSnapState, + hasPermission, + }; + + const engine = new JsonRpcEngine(); + + engine.push((request, response, next, end) => { + const result = implementation( + request as JsonRpcRequest, + response as PendingJsonRpcResponse, + next, + end, + hooks, + ); + + result?.catch(end); + }); + + const response = await engine.handle({ + jsonrpc: '2.0', + id: 1, + method: 'snap_clearState', + params: {}, + }); + + expect(clearSnapState).not.toHaveBeenCalled(); + expect(response).toStrictEqual({ + jsonrpc: '2.0', + id: 1, + error: { + code: errorCodes.provider.unauthorized, + message: + 'The requested account and/or method has not been authorized by the user.', + stack: expect.any(String), + }, + }); + }); + + it('throws if the parameters are invalid', async () => { + const { implementation } = clearStateHandler; + + const clearSnapState = jest.fn(); + const hasPermission = jest.fn().mockReturnValue(true); + + const hooks = { + clearSnapState, + hasPermission, + }; + + const engine = new JsonRpcEngine(); + + engine.push((request, response, next, end) => { + const result = implementation( + request as JsonRpcRequest, + response as PendingJsonRpcResponse, + next, + end, + hooks, + ); + + result?.catch(end); + }); + + const response = await engine.handle({ + jsonrpc: '2.0', + id: 1, + method: 'snap_clearState', + params: { + encrypted: 'foo', + }, + }); + + expect(response).toStrictEqual({ + jsonrpc: '2.0', + id: 1, + error: { + code: errorCodes.rpc.invalidParams, + message: + 'Invalid params: At path: encrypted -- Expected a value of type `boolean`, but received: `"foo"`.', + stack: expect.any(String), + }, + }); + }); + }); +}); diff --git a/packages/snaps-rpc-methods/src/permitted/clearState.ts b/packages/snaps-rpc-methods/src/permitted/clearState.ts index bac8ff19f2..ae32a6ebe5 100644 --- a/packages/snaps-rpc-methods/src/permitted/clearState.ts +++ b/packages/snaps-rpc-methods/src/permitted/clearState.ts @@ -117,6 +117,7 @@ function getValidatedParams(params?: unknown) { }); } - throw error; + /* istanbul ignore next */ + throw rpcErrors.internal(); } } diff --git a/packages/snaps-rpc-methods/src/permitted/getState.test.ts b/packages/snaps-rpc-methods/src/permitted/getState.test.ts index 67a5df5a8d..252ab95d4d 100644 --- a/packages/snaps-rpc-methods/src/permitted/getState.test.ts +++ b/packages/snaps-rpc-methods/src/permitted/getState.test.ts @@ -1,4 +1,283 @@ -import { get } from './getState'; +import { JsonRpcEngine } from '@metamask/json-rpc-engine'; +import { errorCodes } from '@metamask/rpc-errors'; +import type { GetStateResult } from '@metamask/snaps-sdk'; +import { MOCK_SNAP_ID } from '@metamask/snaps-utils/test-utils'; +import type { JsonRpcRequest, PendingJsonRpcResponse } from '@metamask/utils'; + +import type { GetStateParameters } from './getState'; +import { get, getStateHandler } from './getState'; + +describe('snap_getState', () => { + describe('getStateHandler', () => { + it('has the expected shape', () => { + expect(getStateHandler).toMatchObject({ + methodNames: ['snap_getState'], + implementation: expect.any(Function), + hookNames: { + getSnapState: true, + hasPermission: true, + }, + }); + }); + }); + + describe('implementation', () => { + const createOriginMiddleware = + (origin: string) => + (request: any, _response: unknown, next: () => void, _end: unknown) => { + request.origin = origin; + next(); + }; + + it('returns the encrypted state', async () => { + const { implementation } = getStateHandler; + + const getSnapState = jest.fn().mockReturnValue({ + foo: 'bar', + }); + + const getUnlockPromise = jest.fn().mockResolvedValue(undefined); + const hasPermission = jest.fn().mockReturnValue(true); + + const hooks = { + getSnapState, + getUnlockPromise, + hasPermission, + }; + + const engine = new JsonRpcEngine(); + + engine.push(createOriginMiddleware(MOCK_SNAP_ID)); + engine.push((request, response, next, end) => { + const result = implementation( + request as JsonRpcRequest, + response as PendingJsonRpcResponse, + next, + end, + hooks, + ); + + result?.catch(end); + }); + + const response = await engine.handle({ + jsonrpc: '2.0', + id: 1, + method: 'snap_getState', + params: { + key: 'foo', + }, + }); + + expect(getSnapState).toHaveBeenCalledWith(MOCK_SNAP_ID, true); + expect(response).toStrictEqual({ + jsonrpc: '2.0', + id: 1, + result: 'bar', + }); + }); + + it('returns the entire state if no key is specified', async () => { + const { implementation } = getStateHandler; + + const getSnapState = jest.fn().mockReturnValue({ + foo: 'bar', + }); + + const getUnlockPromise = jest.fn().mockResolvedValue(undefined); + const hasPermission = jest.fn().mockReturnValue(true); + + const hooks = { + getSnapState, + getUnlockPromise, + hasPermission, + }; + + const engine = new JsonRpcEngine(); + + engine.push(createOriginMiddleware(MOCK_SNAP_ID)); + engine.push((request, response, next, end) => { + const result = implementation( + request as JsonRpcRequest, + response as PendingJsonRpcResponse, + next, + end, + hooks, + ); + + result?.catch(end); + }); + + const response = await engine.handle({ + jsonrpc: '2.0', + id: 1, + method: 'snap_getState', + params: {}, + }); + + expect(getSnapState).toHaveBeenCalledWith(MOCK_SNAP_ID, true); + expect(response).toStrictEqual({ + jsonrpc: '2.0', + id: 1, + result: { + foo: 'bar', + }, + }); + }); + + it('returns the unencrypted state', async () => { + const { implementation } = getStateHandler; + + const getSnapState = jest.fn().mockReturnValue({ + foo: 'bar', + }); + + const getUnlockPromise = jest.fn().mockResolvedValue(undefined); + const hasPermission = jest.fn().mockReturnValue(true); + + const hooks = { + getSnapState, + getUnlockPromise, + hasPermission, + }; + + const engine = new JsonRpcEngine(); + + engine.push(createOriginMiddleware(MOCK_SNAP_ID)); + engine.push((request, response, next, end) => { + const result = implementation( + request as JsonRpcRequest, + response as PendingJsonRpcResponse, + next, + end, + hooks, + ); + + result?.catch(end); + }); + + const response = await engine.handle({ + jsonrpc: '2.0', + id: 1, + method: 'snap_getState', + params: { + key: 'foo', + encrypted: false, + }, + }); + + expect(getSnapState).toHaveBeenCalledWith(MOCK_SNAP_ID, false); + expect(response).toStrictEqual({ + jsonrpc: '2.0', + id: 1, + result: 'bar', + }); + }); + + it('throws if the requesting origin does not have the required permission', async () => { + const { implementation } = getStateHandler; + + const getSnapState = jest.fn().mockReturnValue({ + foo: 'bar', + }); + + const getUnlockPromise = jest.fn().mockResolvedValue(undefined); + const hasPermission = jest.fn().mockReturnValue(false); + + const hooks = { + getSnapState, + getUnlockPromise, + hasPermission, + }; + + const engine = new JsonRpcEngine(); + + engine.push(createOriginMiddleware(MOCK_SNAP_ID)); + engine.push((request, response, next, end) => { + const result = implementation( + request as JsonRpcRequest, + response as PendingJsonRpcResponse, + next, + end, + hooks, + ); + + result?.catch(end); + }); + + const response = await engine.handle({ + jsonrpc: '2.0', + id: 1, + method: 'snap_getState', + params: {}, + }); + + expect(getSnapState).not.toHaveBeenCalled(); + expect(response).toStrictEqual({ + jsonrpc: '2.0', + id: 1, + error: { + code: errorCodes.provider.unauthorized, + message: + 'The requested account and/or method has not been authorized by the user.', + stack: expect.any(String), + }, + }); + }); + + it('throws if the parameters are invalid', async () => { + const { implementation } = getStateHandler; + + const getSnapState = jest.fn().mockReturnValue({ + foo: 'bar', + }); + + const getUnlockPromise = jest.fn().mockResolvedValue(undefined); + const hasPermission = jest.fn().mockReturnValue(true); + + const hooks = { + getSnapState, + getUnlockPromise, + hasPermission, + }; + + const engine = new JsonRpcEngine(); + + engine.push(createOriginMiddleware(MOCK_SNAP_ID)); + engine.push((request, response, next, end) => { + const result = implementation( + request as JsonRpcRequest, + response as PendingJsonRpcResponse, + next, + end, + hooks, + ); + + result?.catch(end); + }); + + const response = await engine.handle({ + jsonrpc: '2.0', + id: 1, + method: 'snap_getState', + params: { + encrypted: 'foo', + }, + }); + + expect(response).toStrictEqual({ + jsonrpc: '2.0', + id: 1, + error: { + code: errorCodes.rpc.invalidParams, + message: + 'Invalid params: At path: encrypted -- Expected a value of type `boolean`, but received: `"foo"`.', + stack: expect.any(String), + }, + }); + }); + }); +}); describe('get', () => { const object = { diff --git a/packages/snaps-rpc-methods/src/permitted/getState.ts b/packages/snaps-rpc-methods/src/permitted/getState.ts index 0018959e84..7ab1cd90f8 100644 --- a/packages/snaps-rpc-methods/src/permitted/getState.ts +++ b/packages/snaps-rpc-methods/src/permitted/getState.ts @@ -137,7 +137,8 @@ function getValidatedParams(params?: unknown) { }); } - throw error; + /* istanbul ignore next */ + throw rpcErrors.internal(); } } diff --git a/packages/snaps-rpc-methods/src/permitted/setState.test.ts b/packages/snaps-rpc-methods/src/permitted/setState.test.ts index 2433e52015..29e6357707 100644 --- a/packages/snaps-rpc-methods/src/permitted/setState.test.ts +++ b/packages/snaps-rpc-methods/src/permitted/setState.test.ts @@ -1,4 +1,367 @@ -import { set } from './setState'; +import { JsonRpcEngine } from '@metamask/json-rpc-engine'; +import { errorCodes } from '@metamask/rpc-errors'; +import type { SetStateResult } from '@metamask/snaps-sdk'; +import { MOCK_SNAP_ID } from '@metamask/snaps-utils/test-utils'; +import type { JsonRpcRequest, PendingJsonRpcResponse } from '@metamask/utils'; + +import { setStateHandler, type SetStateParameters, set } from './setState'; + +describe('snap_setState', () => { + describe('setStateHandler', () => { + it('has the expected shape', () => { + expect(setStateHandler).toMatchObject({ + methodNames: ['snap_setState'], + implementation: expect.any(Function), + hookNames: { + getSnapState: true, + hasPermission: true, + }, + }); + }); + }); + + describe('implementation', () => { + const createOriginMiddleware = + (origin: string) => + (request: any, _response: unknown, next: () => void, _end: unknown) => { + request.origin = origin; + next(); + }; + + it('sets the encrypted state', async () => { + const { implementation } = setStateHandler; + + const getSnapState = jest.fn().mockReturnValue({ + foo: 'bar', + }); + + const updateSnapState = jest.fn().mockReturnValue(null); + const getUnlockPromise = jest.fn().mockResolvedValue(undefined); + const hasPermission = jest.fn().mockReturnValue(true); + + const hooks = { + getSnapState, + updateSnapState, + getUnlockPromise, + hasPermission, + }; + + const engine = new JsonRpcEngine(); + + engine.push(createOriginMiddleware(MOCK_SNAP_ID)); + engine.push((request, response, next, end) => { + const result = implementation( + request as JsonRpcRequest, + response as PendingJsonRpcResponse, + next, + end, + hooks, + ); + + result?.catch(end); + }); + + const response = await engine.handle({ + jsonrpc: '2.0', + id: 1, + method: 'snap_setState', + params: { + key: 'foo', + value: 'baz', + }, + }); + + expect(getUnlockPromise).toHaveBeenCalled(); + expect(updateSnapState).toHaveBeenCalledWith( + MOCK_SNAP_ID, + { foo: 'baz' }, + true, + ); + + expect(response).toStrictEqual({ + jsonrpc: '2.0', + id: 1, + result: null, + }); + }); + + it('sets the entire state if no key is specified', async () => { + const { implementation } = setStateHandler; + + const getSnapState = jest.fn().mockReturnValue({ + foo: 'bar', + }); + + const updateSnapState = jest.fn().mockReturnValue(null); + const getUnlockPromise = jest.fn().mockResolvedValue(undefined); + const hasPermission = jest.fn().mockReturnValue(true); + + const hooks = { + getSnapState, + updateSnapState, + getUnlockPromise, + hasPermission, + }; + + const engine = new JsonRpcEngine(); + + engine.push(createOriginMiddleware(MOCK_SNAP_ID)); + engine.push((request, response, next, end) => { + const result = implementation( + request as JsonRpcRequest, + response as PendingJsonRpcResponse, + next, + end, + hooks, + ); + + result?.catch(end); + }); + + const response = await engine.handle({ + jsonrpc: '2.0', + id: 1, + method: 'snap_setState', + params: { + value: { + foo: 'baz', + }, + }, + }); + + expect(getUnlockPromise).toHaveBeenCalled(); + expect(updateSnapState).toHaveBeenCalledWith( + MOCK_SNAP_ID, + { foo: 'baz' }, + true, + ); + + expect(response).toStrictEqual({ + jsonrpc: '2.0', + id: 1, + result: null, + }); + }); + + it('sets the unencrypted state', async () => { + const { implementation } = setStateHandler; + + const getSnapState = jest.fn().mockReturnValue({ + foo: 'bar', + }); + + const updateSnapState = jest.fn().mockReturnValue(null); + const getUnlockPromise = jest.fn().mockResolvedValue(undefined); + const hasPermission = jest.fn().mockReturnValue(true); + + const hooks = { + getSnapState, + updateSnapState, + getUnlockPromise, + hasPermission, + }; + + const engine = new JsonRpcEngine(); + + engine.push(createOriginMiddleware(MOCK_SNAP_ID)); + engine.push((request, response, next, end) => { + const result = implementation( + request as JsonRpcRequest, + response as PendingJsonRpcResponse, + next, + end, + hooks, + ); + + result?.catch(end); + }); + + const response = await engine.handle({ + jsonrpc: '2.0', + id: 1, + method: 'snap_setState', + params: { + key: 'foo', + value: 'baz', + encrypted: false, + }, + }); + + expect(getUnlockPromise).not.toHaveBeenCalled(); + expect(updateSnapState).toHaveBeenCalledWith( + MOCK_SNAP_ID, + { + foo: 'baz', + }, + false, + ); + + expect(response).toStrictEqual({ + jsonrpc: '2.0', + id: 1, + result: null, + }); + }); + + it('throws if the requesting origin does not have the required permission', async () => { + const { implementation } = setStateHandler; + + const getSnapState = jest.fn().mockReturnValue({ + foo: 'bar', + }); + + const updateSnapState = jest.fn().mockReturnValue(null); + const getUnlockPromise = jest.fn().mockResolvedValue(undefined); + const hasPermission = jest.fn().mockReturnValue(false); + + const hooks = { + getSnapState, + updateSnapState, + getUnlockPromise, + hasPermission, + }; + + const engine = new JsonRpcEngine(); + + engine.push(createOriginMiddleware(MOCK_SNAP_ID)); + engine.push((request, response, next, end) => { + const result = implementation( + request as JsonRpcRequest, + response as PendingJsonRpcResponse, + next, + end, + hooks, + ); + + result?.catch(end); + }); + + const response = await engine.handle({ + jsonrpc: '2.0', + id: 1, + method: 'snap_setState', + params: {}, + }); + + expect(updateSnapState).not.toHaveBeenCalled(); + expect(response).toStrictEqual({ + jsonrpc: '2.0', + id: 1, + error: { + code: errorCodes.provider.unauthorized, + message: + 'The requested account and/or method has not been authorized by the user.', + stack: expect.any(String), + }, + }); + }); + + it('throws if the parameters are invalid', async () => { + const { implementation } = setStateHandler; + + const getSnapState = jest.fn().mockReturnValue({ + foo: 'bar', + }); + + const updateSnapState = jest.fn().mockReturnValue(null); + const getUnlockPromise = jest.fn().mockResolvedValue(undefined); + const hasPermission = jest.fn().mockReturnValue(true); + + const hooks = { + getSnapState, + updateSnapState, + getUnlockPromise, + hasPermission, + }; + + const engine = new JsonRpcEngine(); + + engine.push(createOriginMiddleware(MOCK_SNAP_ID)); + engine.push((request, response, next, end) => { + const result = implementation( + request as JsonRpcRequest, + response as PendingJsonRpcResponse, + next, + end, + hooks, + ); + + result?.catch(end); + }); + + const response = await engine.handle({ + jsonrpc: '2.0', + id: 1, + method: 'snap_setState', + params: {}, + }); + + expect(response).toStrictEqual({ + jsonrpc: '2.0', + id: 1, + error: { + code: errorCodes.rpc.internal, + message: 'Internal JSON-RPC error.', + stack: expect.any(String), + }, + }); + }); + + it('throws if `key` is not provided and `value` is not an object', async () => { + const { implementation } = setStateHandler; + + const getSnapState = jest.fn().mockReturnValue({ + foo: 'bar', + }); + + const updateSnapState = jest.fn().mockReturnValue(null); + const getUnlockPromise = jest.fn().mockResolvedValue(undefined); + const hasPermission = jest.fn().mockReturnValue(true); + + const hooks = { + getSnapState, + updateSnapState, + getUnlockPromise, + hasPermission, + }; + + const engine = new JsonRpcEngine(); + + engine.push(createOriginMiddleware(MOCK_SNAP_ID)); + engine.push((request, response, next, end) => { + const result = implementation( + request as JsonRpcRequest, + response as PendingJsonRpcResponse, + next, + end, + hooks, + ); + + result?.catch(end); + }); + + const response = await engine.handle({ + jsonrpc: '2.0', + id: 1, + method: 'snap_setState', + params: { + value: 'foo', + }, + }); + + expect(response).toStrictEqual({ + jsonrpc: '2.0', + id: 1, + error: { + code: errorCodes.rpc.invalidParams, + message: + 'Invalid params: Value must be an object if key is not provided.', + stack: expect.any(String), + }, + }); + }); + }); +}); describe('set', () => { it('sets the state in an empty object', () => { diff --git a/packages/snaps-rpc-methods/src/permitted/setState.ts b/packages/snaps-rpc-methods/src/permitted/setState.ts index f28463412d..64530320f2 100644 --- a/packages/snaps-rpc-methods/src/permitted/setState.ts +++ b/packages/snaps-rpc-methods/src/permitted/setState.ts @@ -171,7 +171,8 @@ function getValidatedParams(params?: unknown) { }); } - throw error; + /* istanbul ignore next */ + throw rpcErrors.internal(); } } @@ -225,5 +226,6 @@ export function set( } // This should never be reached. + /* istanbul ignore next */ return {}; } From f5b1c151da7df251115000c1b7658c562d4e22b5 Mon Sep 17 00:00:00 2001 From: Maarten Zuidhoorn Date: Tue, 17 Dec 2024 11:20:12 +0100 Subject: [PATCH 10/31] Update types in example --- packages/examples/packages/manage-state/src/index.ts | 11 ++++++++--- packages/examples/packages/manage-state/src/types.ts | 4 +++- 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/packages/examples/packages/manage-state/src/index.ts b/packages/examples/packages/manage-state/src/index.ts index 12e1cb24c4..6b3fec410c 100644 --- a/packages/examples/packages/manage-state/src/index.ts +++ b/packages/examples/packages/manage-state/src/index.ts @@ -3,7 +3,12 @@ import { type OnRpcRequestHandler, } from '@metamask/snaps-sdk'; -import type { BaseParams, LegacySetStateParams, SetStateParams } from './types'; +import type { + BaseParams, + LegacyParams, + LegacySetStateParams, + SetStateParams, +} from './types'; import { clearState, getState, setState } from './utils'; /** @@ -75,12 +80,12 @@ export const onRpcRequest: OnRpcRequestHandler = async ({ request }) => { } case 'legacy_getState': { - const params = request.params as BaseParams; + const params = request.params as LegacyParams; return await getState(params?.encrypted); } case 'legacy_clearState': { - const params = request.params as BaseParams; + const params = request.params as LegacyParams; await clearState(params?.encrypted); return true; diff --git a/packages/examples/packages/manage-state/src/types.ts b/packages/examples/packages/manage-state/src/types.ts index c433b4b867..897d2a48a2 100644 --- a/packages/examples/packages/manage-state/src/types.ts +++ b/packages/examples/packages/manage-state/src/types.ts @@ -2,6 +2,8 @@ import type { Json } from '@metamask/utils'; import type { State } from './utils'; +export type LegacyParams = { encrypted?: boolean }; + export type BaseParams = { key?: string; encrypted?: boolean }; /** @@ -14,4 +16,4 @@ export type SetStateParams = BaseParams & { /** * The parameters for the `legacy_setState` JSON-RPC method. */ -export type LegacySetStateParams = BaseParams & Partial; +export type LegacySetStateParams = LegacyParams & Partial; From 9144281787f9bf6f7e217989d643da5d87455605 Mon Sep 17 00:00:00 2001 From: Maarten Zuidhoorn Date: Tue, 17 Dec 2024 11:21:11 +0100 Subject: [PATCH 11/31] Add missing hook --- packages/snaps-rpc-methods/src/permitted/index.ts | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/packages/snaps-rpc-methods/src/permitted/index.ts b/packages/snaps-rpc-methods/src/permitted/index.ts index cadbf563fc..47febaba09 100644 --- a/packages/snaps-rpc-methods/src/permitted/index.ts +++ b/packages/snaps-rpc-methods/src/permitted/index.ts @@ -1,3 +1,4 @@ +import type { ClearStateHooks } from './clearState'; import type { CreateInterfaceMethodHooks } from './createInterface'; import type { ProviderRequestMethodHooks } from './experimentalProviderRequest'; import type { GetAllSnapsHooks } from './getAllSnaps'; @@ -11,7 +12,8 @@ import type { ResolveInterfaceMethodHooks } from './resolveInterface'; import type { SetStateHooks } from './setState'; import type { UpdateInterfaceMethodHooks } from './updateInterface'; -export type PermittedRpcMethodHooks = GetAllSnapsHooks & +export type PermittedRpcMethodHooks = ClearStateHooks & + GetAllSnapsHooks & GetClientStatusHooks & GetSnapsHooks & GetStateHooks & From 549142ae662ed8bcd3849683a7d2b301cfde1c2e Mon Sep 17 00:00:00 2001 From: Maarten Zuidhoorn Date: Tue, 17 Dec 2024 11:31:54 +0100 Subject: [PATCH 12/31] Assume methods bound to origin --- .../src/permitted/clearState.test.ts | 14 +----- .../src/permitted/clearState.ts | 7 +-- .../src/permitted/getState.test.ts | 19 ++------ .../src/permitted/getState.ts | 10 +---- .../src/permitted/setState.test.ts | 27 +----------- .../src/permitted/setState.ts | 43 ++++++++++++++----- 6 files changed, 43 insertions(+), 77 deletions(-) diff --git a/packages/snaps-rpc-methods/src/permitted/clearState.test.ts b/packages/snaps-rpc-methods/src/permitted/clearState.test.ts index 1154341659..d3ca4987e2 100644 --- a/packages/snaps-rpc-methods/src/permitted/clearState.test.ts +++ b/packages/snaps-rpc-methods/src/permitted/clearState.test.ts @@ -1,7 +1,6 @@ import { JsonRpcEngine } from '@metamask/json-rpc-engine'; import { errorCodes } from '@metamask/rpc-errors'; import type { ClearStateResult } from '@metamask/snaps-sdk'; -import { MOCK_SNAP_ID } from '@metamask/snaps-utils/test-utils'; import type { JsonRpcRequest, PendingJsonRpcResponse } from '@metamask/utils'; import type { ClearStateParameters } from './clearState'; @@ -22,13 +21,6 @@ describe('snap_clearState', () => { }); describe('implementation', () => { - const createOriginMiddleware = - (origin: string) => - (request: any, _response: unknown, next: () => void, _end: unknown) => { - request.origin = origin; - next(); - }; - it('returns the result from the `clearSnapState` hook', async () => { const { implementation } = clearStateHandler; @@ -42,7 +34,6 @@ describe('snap_clearState', () => { const engine = new JsonRpcEngine(); - engine.push(createOriginMiddleware(MOCK_SNAP_ID)); engine.push((request, response, next, end) => { const result = implementation( request as JsonRpcRequest, @@ -62,7 +53,7 @@ describe('snap_clearState', () => { params: {}, }); - expect(clearSnapState).toHaveBeenCalledWith(MOCK_SNAP_ID, true); + expect(clearSnapState).toHaveBeenCalledWith(true); expect(response).toStrictEqual({ jsonrpc: '2.0', id: 1, @@ -83,7 +74,6 @@ describe('snap_clearState', () => { const engine = new JsonRpcEngine(); - engine.push(createOriginMiddleware(MOCK_SNAP_ID)); engine.push((request, response, next, end) => { const result = implementation( request as JsonRpcRequest, @@ -105,7 +95,7 @@ describe('snap_clearState', () => { }, }); - expect(clearSnapState).toHaveBeenCalledWith(MOCK_SNAP_ID, false); + expect(clearSnapState).toHaveBeenCalledWith(false); expect(response).toStrictEqual({ jsonrpc: '2.0', id: 1, diff --git a/packages/snaps-rpc-methods/src/permitted/clearState.ts b/packages/snaps-rpc-methods/src/permitted/clearState.ts index ae32a6ebe5..0c4738ca66 100644 --- a/packages/snaps-rpc-methods/src/permitted/clearState.ts +++ b/packages/snaps-rpc-methods/src/permitted/clearState.ts @@ -37,7 +37,7 @@ export type ClearStateHooks = { /** * A function that clears the state of the requesting Snap. */ - clearSnapState: (snapId: string, encrypted: boolean) => void; + clearSnapState: (encrypted: boolean) => void; /** * Check if the requesting origin has a given permission. @@ -89,10 +89,7 @@ async function clearStateImplementation( const validatedParams = getValidatedParams(params); const { encrypted = true } = validatedParams; - // We expect the MM middleware stack to always add the origin to requests - const { origin } = request as JsonRpcRequest & { origin: string }; - clearSnapState(origin, encrypted); - + clearSnapState(encrypted); response.result = null; } catch (error) { return end(error); diff --git a/packages/snaps-rpc-methods/src/permitted/getState.test.ts b/packages/snaps-rpc-methods/src/permitted/getState.test.ts index 252ab95d4d..ad7b256137 100644 --- a/packages/snaps-rpc-methods/src/permitted/getState.test.ts +++ b/packages/snaps-rpc-methods/src/permitted/getState.test.ts @@ -1,7 +1,6 @@ import { JsonRpcEngine } from '@metamask/json-rpc-engine'; import { errorCodes } from '@metamask/rpc-errors'; import type { GetStateResult } from '@metamask/snaps-sdk'; -import { MOCK_SNAP_ID } from '@metamask/snaps-utils/test-utils'; import type { JsonRpcRequest, PendingJsonRpcResponse } from '@metamask/utils'; import type { GetStateParameters } from './getState'; @@ -22,13 +21,6 @@ describe('snap_getState', () => { }); describe('implementation', () => { - const createOriginMiddleware = - (origin: string) => - (request: any, _response: unknown, next: () => void, _end: unknown) => { - request.origin = origin; - next(); - }; - it('returns the encrypted state', async () => { const { implementation } = getStateHandler; @@ -47,7 +39,6 @@ describe('snap_getState', () => { const engine = new JsonRpcEngine(); - engine.push(createOriginMiddleware(MOCK_SNAP_ID)); engine.push((request, response, next, end) => { const result = implementation( request as JsonRpcRequest, @@ -69,7 +60,7 @@ describe('snap_getState', () => { }, }); - expect(getSnapState).toHaveBeenCalledWith(MOCK_SNAP_ID, true); + expect(getSnapState).toHaveBeenCalledWith(true); expect(response).toStrictEqual({ jsonrpc: '2.0', id: 1, @@ -95,7 +86,6 @@ describe('snap_getState', () => { const engine = new JsonRpcEngine(); - engine.push(createOriginMiddleware(MOCK_SNAP_ID)); engine.push((request, response, next, end) => { const result = implementation( request as JsonRpcRequest, @@ -115,7 +105,7 @@ describe('snap_getState', () => { params: {}, }); - expect(getSnapState).toHaveBeenCalledWith(MOCK_SNAP_ID, true); + expect(getSnapState).toHaveBeenCalledWith(true); expect(response).toStrictEqual({ jsonrpc: '2.0', id: 1, @@ -143,7 +133,6 @@ describe('snap_getState', () => { const engine = new JsonRpcEngine(); - engine.push(createOriginMiddleware(MOCK_SNAP_ID)); engine.push((request, response, next, end) => { const result = implementation( request as JsonRpcRequest, @@ -166,7 +155,7 @@ describe('snap_getState', () => { }, }); - expect(getSnapState).toHaveBeenCalledWith(MOCK_SNAP_ID, false); + expect(getSnapState).toHaveBeenCalledWith(false); expect(response).toStrictEqual({ jsonrpc: '2.0', id: 1, @@ -192,7 +181,6 @@ describe('snap_getState', () => { const engine = new JsonRpcEngine(); - engine.push(createOriginMiddleware(MOCK_SNAP_ID)); engine.push((request, response, next, end) => { const result = implementation( request as JsonRpcRequest, @@ -243,7 +231,6 @@ describe('snap_getState', () => { const engine = new JsonRpcEngine(); - engine.push(createOriginMiddleware(MOCK_SNAP_ID)); engine.push((request, response, next, end) => { const result = implementation( request as JsonRpcRequest, diff --git a/packages/snaps-rpc-methods/src/permitted/getState.ts b/packages/snaps-rpc-methods/src/permitted/getState.ts index 7ab1cd90f8..27e1d72f22 100644 --- a/packages/snaps-rpc-methods/src/permitted/getState.ts +++ b/packages/snaps-rpc-methods/src/permitted/getState.ts @@ -50,10 +50,7 @@ export type GetStateHooks = { * * @returns The current state of the Snap. */ - getSnapState: ( - snapId: string, - encrypted: boolean, - ) => Promise>; + getSnapState: (encrypted: boolean) => Promise>; /** * Wait for the extension to be unlocked. @@ -109,10 +106,7 @@ async function getStateImplementation( await getUnlockPromise(true); } - // We expect the MM middleware stack to always add the origin to requests - const { origin } = request as JsonRpcRequest & { origin: string }; - const state = await getSnapState(origin, encrypted); - + const state = await getSnapState(encrypted); response.result = get(state, key); } catch (error) { return end(error); diff --git a/packages/snaps-rpc-methods/src/permitted/setState.test.ts b/packages/snaps-rpc-methods/src/permitted/setState.test.ts index 29e6357707..9441f2c7b6 100644 --- a/packages/snaps-rpc-methods/src/permitted/setState.test.ts +++ b/packages/snaps-rpc-methods/src/permitted/setState.test.ts @@ -1,7 +1,6 @@ import { JsonRpcEngine } from '@metamask/json-rpc-engine'; import { errorCodes } from '@metamask/rpc-errors'; import type { SetStateResult } from '@metamask/snaps-sdk'; -import { MOCK_SNAP_ID } from '@metamask/snaps-utils/test-utils'; import type { JsonRpcRequest, PendingJsonRpcResponse } from '@metamask/utils'; import { setStateHandler, type SetStateParameters, set } from './setState'; @@ -21,13 +20,6 @@ describe('snap_setState', () => { }); describe('implementation', () => { - const createOriginMiddleware = - (origin: string) => - (request: any, _response: unknown, next: () => void, _end: unknown) => { - request.origin = origin; - next(); - }; - it('sets the encrypted state', async () => { const { implementation } = setStateHandler; @@ -48,7 +40,6 @@ describe('snap_setState', () => { const engine = new JsonRpcEngine(); - engine.push(createOriginMiddleware(MOCK_SNAP_ID)); engine.push((request, response, next, end) => { const result = implementation( request as JsonRpcRequest, @@ -72,11 +63,7 @@ describe('snap_setState', () => { }); expect(getUnlockPromise).toHaveBeenCalled(); - expect(updateSnapState).toHaveBeenCalledWith( - MOCK_SNAP_ID, - { foo: 'baz' }, - true, - ); + expect(updateSnapState).toHaveBeenCalledWith({ foo: 'baz' }, true); expect(response).toStrictEqual({ jsonrpc: '2.0', @@ -105,7 +92,6 @@ describe('snap_setState', () => { const engine = new JsonRpcEngine(); - engine.push(createOriginMiddleware(MOCK_SNAP_ID)); engine.push((request, response, next, end) => { const result = implementation( request as JsonRpcRequest, @@ -130,11 +116,7 @@ describe('snap_setState', () => { }); expect(getUnlockPromise).toHaveBeenCalled(); - expect(updateSnapState).toHaveBeenCalledWith( - MOCK_SNAP_ID, - { foo: 'baz' }, - true, - ); + expect(updateSnapState).toHaveBeenCalledWith({ foo: 'baz' }, true); expect(response).toStrictEqual({ jsonrpc: '2.0', @@ -163,7 +145,6 @@ describe('snap_setState', () => { const engine = new JsonRpcEngine(); - engine.push(createOriginMiddleware(MOCK_SNAP_ID)); engine.push((request, response, next, end) => { const result = implementation( request as JsonRpcRequest, @@ -189,7 +170,6 @@ describe('snap_setState', () => { expect(getUnlockPromise).not.toHaveBeenCalled(); expect(updateSnapState).toHaveBeenCalledWith( - MOCK_SNAP_ID, { foo: 'baz', }, @@ -223,7 +203,6 @@ describe('snap_setState', () => { const engine = new JsonRpcEngine(); - engine.push(createOriginMiddleware(MOCK_SNAP_ID)); engine.push((request, response, next, end) => { const result = implementation( request as JsonRpcRequest, @@ -276,7 +255,6 @@ describe('snap_setState', () => { const engine = new JsonRpcEngine(); - engine.push(createOriginMiddleware(MOCK_SNAP_ID)); engine.push((request, response, next, end) => { const result = implementation( request as JsonRpcRequest, @@ -327,7 +305,6 @@ describe('snap_setState', () => { const engine = new JsonRpcEngine(); - engine.push(createOriginMiddleware(MOCK_SNAP_ID)); engine.push((request, response, next, end) => { const result = implementation( request as JsonRpcRequest, diff --git a/packages/snaps-rpc-methods/src/permitted/setState.ts b/packages/snaps-rpc-methods/src/permitted/setState.ts index 64530320f2..f2e7392f04 100644 --- a/packages/snaps-rpc-methods/src/permitted/setState.ts +++ b/packages/snaps-rpc-methods/src/permitted/setState.ts @@ -54,10 +54,7 @@ export type SetStateHooks = { * @param encrypted - Whether the state is encrypted. * @returns The current state of the Snap. */ - getSnapState: ( - snapId: string, - encrypted: boolean, - ) => Promise>; + getSnapState: (encrypted: boolean) => Promise>; /** * Wait for the extension to be unlocked. @@ -74,7 +71,6 @@ export type SetStateHooks = { * @param encrypted - Whether the state should be encrypted. */ updateSnapState: ( - snapId: string, newState: Record, encrypted: boolean, ) => Promise; @@ -133,10 +129,6 @@ async function setStateImplementation( await getUnlockPromise(true); } - // We expect the MM middleware stack to always add the origin to requests - const { origin } = request as JsonRpcRequest & { origin: string }; - const state = await getSnapState(origin, encrypted); - if (key === undefined && !isPlainObject(value)) { return end( rpcErrors.invalidParams( @@ -145,8 +137,8 @@ async function setStateImplementation( ); } - const newState = set(state, key, value); - await updateSnapState(origin, newState, encrypted); + const newState = await getNewState(key, value, getSnapState); + await updateSnapState(newState, encrypted); response.result = null; } catch (error) { return end(error); @@ -176,6 +168,35 @@ function getValidatedParams(params?: unknown) { } } +/** + * Get the new state of the Snap. + * + * 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). + * + * @param key - The key to set. + * @param value - The value to set the key to. + * @param getSnapState - The `getSnapState` hook. + * @returns The new state of the Snap. + */ +async function getNewState( + key: string | undefined, + value: Json, + getSnapState: SetStateHooks['getSnapState'], +) { + if (key === undefined) { + assert(isPlainObject(value)); + return value; + } + + const state = await getSnapState(false); + return set(state, key, value); +} + /** * Set 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 From b5ac14e4f4c0653e257826b88d302d892a804a34 Mon Sep 17 00:00:00 2001 From: Maarten Zuidhoorn Date: Tue, 17 Dec 2024 12:35:35 +0100 Subject: [PATCH 13/31] Update packages/snaps-sdk/src/types/methods/get-state.ts Co-authored-by: Frederik Bolding --- packages/snaps-sdk/src/types/methods/get-state.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/snaps-sdk/src/types/methods/get-state.ts b/packages/snaps-sdk/src/types/methods/get-state.ts index 1692e7ea52..e3931bca6a 100644 --- a/packages/snaps-sdk/src/types/methods/get-state.ts +++ b/packages/snaps-sdk/src/types/methods/get-state.ts @@ -8,8 +8,8 @@ import type { Json } from '@metamask/utils'; * `a.b.c`, with the exception of array syntax. * @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 extension is unlocked, while unencrypted state can be used - * whether the extension is locked or unlocked. + * be used if the client is unlocked, while unencrypted state can be used + * whether the client is locked or unlocked. */ export type GetStateParams = { key?: string; From b324ad2ff6645d784c71a9841f101c386dd83730 Mon Sep 17 00:00:00 2001 From: Maarten Zuidhoorn Date: Tue, 17 Dec 2024 12:45:19 +0100 Subject: [PATCH 14/31] Update packages/snaps-sdk/src/types/methods/set-state.ts Co-authored-by: Frederik Bolding --- packages/snaps-sdk/src/types/methods/set-state.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/snaps-sdk/src/types/methods/set-state.ts b/packages/snaps-sdk/src/types/methods/set-state.ts index d5c30df46c..d0975cc78b 100644 --- a/packages/snaps-sdk/src/types/methods/set-state.ts +++ b/packages/snaps-sdk/src/types/methods/set-state.ts @@ -9,8 +9,8 @@ import type { Json } from '@metamask/utils'; * @property value - The value to set the state to. * @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 extension is unlocked, while unencrypted state can be used - * whether the extension is locked or unlocked. + * be used if the client is unlocked, while unencrypted state can be used + * whether the client is locked or unlocked. */ export type SetStateParams = { key?: string; From 6d1af36ee2a11ab4e74f26166234a3895e4f2b4f Mon Sep 17 00:00:00 2001 From: Maarten Zuidhoorn Date: Tue, 17 Dec 2024 12:48:28 +0100 Subject: [PATCH 15/31] Throw error for invalid branch --- packages/snaps-rpc-methods/src/permitted/setState.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/snaps-rpc-methods/src/permitted/setState.ts b/packages/snaps-rpc-methods/src/permitted/setState.ts index f2e7392f04..4f2e8907dd 100644 --- a/packages/snaps-rpc-methods/src/permitted/setState.ts +++ b/packages/snaps-rpc-methods/src/permitted/setState.ts @@ -248,5 +248,5 @@ export function set( // This should never be reached. /* istanbul ignore next */ - return {}; + throw new Error('Unexpected error while setting the state.'); } From 05ffe82e25c09312e1442a3dfac953474bb4918e Mon Sep 17 00:00:00 2001 From: Maarten Zuidhoorn Date: Wed, 18 Dec 2024 14:04:49 +0100 Subject: [PATCH 16/31] Use `isObject` for better performance --- packages/snaps-rpc-methods/jest.config.js | 4 ++-- packages/snaps-rpc-methods/src/permitted/setState.ts | 10 +++++----- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/packages/snaps-rpc-methods/jest.config.js b/packages/snaps-rpc-methods/jest.config.js index 918d7696ca..0e25d0bc3b 100644 --- a/packages/snaps-rpc-methods/jest.config.js +++ b/packages/snaps-rpc-methods/jest.config.js @@ -10,8 +10,8 @@ module.exports = deepmerge(baseConfig, { ], coverageThreshold: { global: { - branches: 92.95, - functions: 97.44, + branches: 92.98, + functions: 97.46, lines: 97.93, statements: 97.52, }, diff --git a/packages/snaps-rpc-methods/src/permitted/setState.ts b/packages/snaps-rpc-methods/src/permitted/setState.ts index 4f2e8907dd..b54dcdc772 100644 --- a/packages/snaps-rpc-methods/src/permitted/setState.ts +++ b/packages/snaps-rpc-methods/src/permitted/setState.ts @@ -12,7 +12,7 @@ import { StructError, } from '@metamask/superstruct'; import type { PendingJsonRpcResponse, JsonRpcRequest } from '@metamask/utils'; -import { assert, JsonStruct, isPlainObject, type Json } from '@metamask/utils'; +import { isObject, assert, JsonStruct, type Json } from '@metamask/utils'; import { manageStateBuilder } from '../restricted/manageState'; import type { MethodHooksObject } from '../utils'; @@ -129,7 +129,7 @@ async function setStateImplementation( await getUnlockPromise(true); } - if (key === undefined && !isPlainObject(value)) { + if (key === undefined && !isObject(value)) { return end( rpcErrors.invalidParams( 'Invalid params: Value must be an object if key is not provided.', @@ -189,7 +189,7 @@ async function getNewState( getSnapState: SetStateHooks['getSnapState'], ) { if (key === undefined) { - assert(isPlainObject(value)); + assert(isObject(value)); return value; } @@ -218,7 +218,7 @@ export function set( value: Json, ): JsonObject { if (key === undefined) { - assert(isPlainObject(value)); + assert(isObject(value)); return value; } @@ -239,7 +239,7 @@ export function set( return requiredObject; } - currentObject[currentKey] = isPlainObject(currentObject[currentKey]) + currentObject[currentKey] = isObject(currentObject[currentKey]) ? currentObject[currentKey] : {}; From 4b1f033f2d36a18e90f301b10efd56a087601505 Mon Sep 17 00:00:00 2001 From: Maarten Zuidhoorn Date: Wed, 18 Dec 2024 14:07:42 +0100 Subject: [PATCH 17/31] Use `isObject` for better performance in getState handler --- packages/snaps-rpc-methods/src/permitted/getState.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/snaps-rpc-methods/src/permitted/getState.ts b/packages/snaps-rpc-methods/src/permitted/getState.ts index 27e1d72f22..a740c0720f 100644 --- a/packages/snaps-rpc-methods/src/permitted/getState.ts +++ b/packages/snaps-rpc-methods/src/permitted/getState.ts @@ -11,7 +11,7 @@ import { StructError, } from '@metamask/superstruct'; import type { PendingJsonRpcResponse, JsonRpcRequest } from '@metamask/utils'; -import { hasProperty, isPlainObject, type Json } from '@metamask/utils'; +import { hasProperty, isObject, type Json } from '@metamask/utils'; import { manageStateBuilder } from '../restricted/manageState'; import type { MethodHooksObject } from '../utils'; @@ -167,7 +167,7 @@ export function get( ); } - if (isPlainObject(result)) { + if (isObject(result)) { if (!hasProperty(result, currentKey)) { return null; } From 9e3c1db68748d12cfb880c80dcd0a09d6351742b Mon Sep 17 00:00:00 2001 From: Maarten Zuidhoorn Date: Wed, 18 Dec 2024 14:13:48 +0100 Subject: [PATCH 18/31] More small performance optimisations --- packages/snaps-rpc-methods/src/permitted/getState.ts | 5 ++++- packages/snaps-rpc-methods/src/permitted/setState.ts | 8 ++++---- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/packages/snaps-rpc-methods/src/permitted/getState.ts b/packages/snaps-rpc-methods/src/permitted/getState.ts index a740c0720f..a4ce156021 100644 --- a/packages/snaps-rpc-methods/src/permitted/getState.ts +++ b/packages/snaps-rpc-methods/src/permitted/getState.ts @@ -160,7 +160,10 @@ export function get( const keys = key.split('.'); let result: Json = value; - for (const currentKey of keys) { + // Intentionally using a classic for loop here for performance reasons. + // eslint-disable-next-line @typescript-eslint/prefer-for-of + for (let i = 0; i < keys.length; i++) { + const currentKey = keys[i]; if (['__proto__', 'constructor'].includes(currentKey)) { throw rpcErrors.invalidParams( 'Invalid params: Key contains forbidden characters.', diff --git a/packages/snaps-rpc-methods/src/permitted/setState.ts b/packages/snaps-rpc-methods/src/permitted/setState.ts index b54dcdc772..a048206e1e 100644 --- a/packages/snaps-rpc-methods/src/permitted/setState.ts +++ b/packages/snaps-rpc-methods/src/permitted/setState.ts @@ -125,10 +125,6 @@ async function setStateImplementation( const validatedParams = getValidatedParams(params); const { key, value, encrypted = true } = validatedParams; - if (encrypted) { - await getUnlockPromise(true); - } - if (key === undefined && !isObject(value)) { return end( rpcErrors.invalidParams( @@ -137,6 +133,10 @@ async function setStateImplementation( ); } + if (encrypted) { + await getUnlockPromise(true); + } + const newState = await getNewState(key, value, getSnapState); await updateSnapState(newState, encrypted); response.result = null; From ce07fbbe61ad0b83aa0a6bbc69dba373aeecc0ac Mon Sep 17 00:00:00 2001 From: Maarten Zuidhoorn Date: Wed, 18 Dec 2024 14:14:10 +0100 Subject: [PATCH 19/31] Update coverage --- packages/snaps-rpc-methods/jest.config.js | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/snaps-rpc-methods/jest.config.js b/packages/snaps-rpc-methods/jest.config.js index 0e25d0bc3b..8d5a260d32 100644 --- a/packages/snaps-rpc-methods/jest.config.js +++ b/packages/snaps-rpc-methods/jest.config.js @@ -12,7 +12,7 @@ module.exports = deepmerge(baseConfig, { global: { branches: 92.98, functions: 97.46, - lines: 97.93, + lines: 97.94, statements: 97.52, }, }, From 01c3b7e4ab2f86026e2c77f7751fa0083b2d94a9 Mon Sep 17 00:00:00 2001 From: Maarten Zuidhoorn Date: Wed, 18 Dec 2024 14:36:12 +0100 Subject: [PATCH 20/31] Reuse some logic in test-snaps --- .../test-snaps/src/features/snaps/index.ts | 2 +- .../LegacyState.tsx} | 14 ++++---- .../components/ClearData.tsx | 2 +- .../components/SendData.tsx | 12 +++++-- .../components/index.ts | 0 .../src/features/snaps/legacy-state/index.ts | 1 + .../features/snaps/manage-state/constants.ts | 5 --- .../snaps/manage-state/hooks/index.ts | 1 - .../snaps/manage-state/hooks/useSnapState.ts | 32 ------------------- .../src/features/snaps/manage-state/index.ts | 1 - .../src/features/snaps/state/State.tsx | 4 +-- .../snaps/state/hooks/useSnapState.ts | 13 ++++---- .../src/features/snaps/state/index.ts | 2 ++ 13 files changed, 30 insertions(+), 59 deletions(-) rename packages/test-snaps/src/features/snaps/{manage-state/ManageState.tsx => legacy-state/LegacyState.tsx} (78%) rename packages/test-snaps/src/features/snaps/{manage-state => legacy-state}/components/ClearData.tsx (94%) rename packages/test-snaps/src/features/snaps/{manage-state => legacy-state}/components/SendData.tsx (91%) rename packages/test-snaps/src/features/snaps/{manage-state => legacy-state}/components/index.ts (100%) create mode 100644 packages/test-snaps/src/features/snaps/legacy-state/index.ts delete mode 100644 packages/test-snaps/src/features/snaps/manage-state/constants.ts delete mode 100644 packages/test-snaps/src/features/snaps/manage-state/hooks/index.ts delete mode 100644 packages/test-snaps/src/features/snaps/manage-state/hooks/useSnapState.ts delete mode 100644 packages/test-snaps/src/features/snaps/manage-state/index.ts diff --git a/packages/test-snaps/src/features/snaps/index.ts b/packages/test-snaps/src/features/snaps/index.ts index d563710530..7eadf6599e 100644 --- a/packages/test-snaps/src/features/snaps/index.ts +++ b/packages/test-snaps/src/features/snaps/index.ts @@ -16,7 +16,7 @@ export * from './images'; export * from './json-rpc'; export * from './jsx'; export * from './lifecycle-hooks'; -export * from './manage-state'; +export * from './legacy-state'; export * from './multi-install'; export * from './name-lookup'; export * from './network-access'; diff --git a/packages/test-snaps/src/features/snaps/manage-state/ManageState.tsx b/packages/test-snaps/src/features/snaps/legacy-state/LegacyState.tsx similarity index 78% rename from packages/test-snaps/src/features/snaps/manage-state/ManageState.tsx rename to packages/test-snaps/src/features/snaps/legacy-state/LegacyState.tsx index 386e56a704..37ad672112 100644 --- a/packages/test-snaps/src/features/snaps/manage-state/ManageState.tsx +++ b/packages/test-snaps/src/features/snaps/legacy-state/LegacyState.tsx @@ -1,21 +1,21 @@ import type { FunctionComponent } from 'react'; import { Result, Snap } from '../../../components'; -import { ClearData, SendData } from './components'; import { MANAGE_STATE_SNAP_ID, MANAGE_STATE_PORT, MANAGE_STATE_VERSION, -} from './constants'; -import { useSnapState } from './hooks'; + useSnapState, +} from '../state'; +import { ClearData, SendData } from './components'; -export const ManageState: FunctionComponent = () => { - const encryptedState = useSnapState(true); - const unencryptedState = useSnapState(false); +export const LegacyState: FunctionComponent = () => { + const encryptedState = useSnapState('legacy_getState', true); + const unencryptedState = useSnapState('legacy_getState', false); return ( = ({ encrypted, diff --git a/packages/test-snaps/src/features/snaps/manage-state/components/SendData.tsx b/packages/test-snaps/src/features/snaps/legacy-state/components/SendData.tsx similarity index 91% rename from packages/test-snaps/src/features/snaps/manage-state/components/SendData.tsx rename to packages/test-snaps/src/features/snaps/legacy-state/components/SendData.tsx index cb3949c817..53c2e89cf6 100644 --- a/packages/test-snaps/src/features/snaps/manage-state/components/SendData.tsx +++ b/packages/test-snaps/src/features/snaps/legacy-state/components/SendData.tsx @@ -6,15 +6,21 @@ import { Button, Form } from 'react-bootstrap'; import { Tag, useInvokeMutation } from '../../../../api'; import { Result } from '../../../../components'; import { getSnapId } from '../../../../utils'; -import { MANAGE_STATE_PORT, MANAGE_STATE_SNAP_ID } from '../constants'; -import { useSnapState } from '../hooks'; +import { + MANAGE_STATE_PORT, + MANAGE_STATE_SNAP_ID, + useSnapState, +} from '../../state'; export const SendData: FunctionComponent<{ encrypted: boolean }> = ({ encrypted, }) => { const [value, setValue] = useState(''); const [invokeSnap, { isLoading, data, error }] = useInvokeMutation(); - const snapState = useSnapState(encrypted); + const snapState = useSnapState<{ items: string[] }>( + 'legacy_getState', + encrypted, + ); const handleChange = (event: ChangeEvent) => { setValue(event.target.value); diff --git a/packages/test-snaps/src/features/snaps/manage-state/components/index.ts b/packages/test-snaps/src/features/snaps/legacy-state/components/index.ts similarity index 100% rename from packages/test-snaps/src/features/snaps/manage-state/components/index.ts rename to packages/test-snaps/src/features/snaps/legacy-state/components/index.ts diff --git a/packages/test-snaps/src/features/snaps/legacy-state/index.ts b/packages/test-snaps/src/features/snaps/legacy-state/index.ts new file mode 100644 index 0000000000..4e42145e67 --- /dev/null +++ b/packages/test-snaps/src/features/snaps/legacy-state/index.ts @@ -0,0 +1 @@ +export * from './LegacyState'; diff --git a/packages/test-snaps/src/features/snaps/manage-state/constants.ts b/packages/test-snaps/src/features/snaps/manage-state/constants.ts deleted file mode 100644 index cd2622fdbc..0000000000 --- a/packages/test-snaps/src/features/snaps/manage-state/constants.ts +++ /dev/null @@ -1,5 +0,0 @@ -import packageJson from '@metamask/manage-state-example-snap/package.json'; - -export const MANAGE_STATE_SNAP_ID = 'npm:@metamask/manage-state-example-snap'; -export const MANAGE_STATE_PORT = 8014; -export const MANAGE_STATE_VERSION = packageJson.version; diff --git a/packages/test-snaps/src/features/snaps/manage-state/hooks/index.ts b/packages/test-snaps/src/features/snaps/manage-state/hooks/index.ts deleted file mode 100644 index 3a6999aa10..0000000000 --- a/packages/test-snaps/src/features/snaps/manage-state/hooks/index.ts +++ /dev/null @@ -1 +0,0 @@ -export * from './useSnapState'; diff --git a/packages/test-snaps/src/features/snaps/manage-state/hooks/useSnapState.ts b/packages/test-snaps/src/features/snaps/manage-state/hooks/useSnapState.ts deleted file mode 100644 index d0139f9bc1..0000000000 --- a/packages/test-snaps/src/features/snaps/manage-state/hooks/useSnapState.ts +++ /dev/null @@ -1,32 +0,0 @@ -import { Tag, useInvokeQuery } from '../../../../api'; -import { getSnapId, useInstalled } from '../../../../utils'; -import { MANAGE_STATE_PORT, MANAGE_STATE_SNAP_ID } from '../constants'; - -export type State = { - items: string[]; -}; - -/** - * Hook to retrieve the state of the snap. - * - * @param encrypted - A flag to indicate whether to use encrypted storage or not. - * @returns The state of the snap. - */ -export function useSnapState(encrypted: boolean) { - const snapId = getSnapId(MANAGE_STATE_SNAP_ID, MANAGE_STATE_PORT); - const isInstalled = useInstalled(snapId); - - const { data: state } = useInvokeQuery<{ data: State }>( - { - snapId, - method: 'legacy_getState', - params: { encrypted }, - tags: [encrypted ? Tag.TestState : Tag.UnencryptedTestState], - }, - { - skip: !isInstalled, - }, - ); - - return state; -} diff --git a/packages/test-snaps/src/features/snaps/manage-state/index.ts b/packages/test-snaps/src/features/snaps/manage-state/index.ts deleted file mode 100644 index 65d95d15b4..0000000000 --- a/packages/test-snaps/src/features/snaps/manage-state/index.ts +++ /dev/null @@ -1 +0,0 @@ -export * from './ManageState'; diff --git a/packages/test-snaps/src/features/snaps/state/State.tsx b/packages/test-snaps/src/features/snaps/state/State.tsx index 65b7f40daa..2e29b4d93a 100644 --- a/packages/test-snaps/src/features/snaps/state/State.tsx +++ b/packages/test-snaps/src/features/snaps/state/State.tsx @@ -10,8 +10,8 @@ import { import { useSnapState } from './hooks'; export const State: FunctionComponent = () => { - const encryptedState = useSnapState(true); - const unencryptedState = useSnapState(false); + const encryptedState = useSnapState('getState', true); + const unencryptedState = useSnapState('getState', false); return ( = Record, +>(method: string, encrypted: boolean): State { const snapId = getSnapId(MANAGE_STATE_SNAP_ID, MANAGE_STATE_PORT); const isInstalled = useInstalled(snapId); const { data: state } = useInvokeQuery<{ data: State }>( { snapId, - method: 'getState', + method, params: { encrypted }, tags: [encrypted ? Tag.TestState : Tag.UnencryptedTestState], }, diff --git a/packages/test-snaps/src/features/snaps/state/index.ts b/packages/test-snaps/src/features/snaps/state/index.ts index e11bf86ccf..c06d9069e4 100644 --- a/packages/test-snaps/src/features/snaps/state/index.ts +++ b/packages/test-snaps/src/features/snaps/state/index.ts @@ -1 +1,3 @@ +export * from './constants'; +export * from './hooks'; export * from './State'; From bafff42d430c45c554d90ed4a94aea0b3abbc374 Mon Sep 17 00:00:00 2001 From: Maarten Zuidhoorn Date: Wed, 18 Dec 2024 14:48:27 +0100 Subject: [PATCH 21/31] Add headings to test-snaps --- packages/examples/packages/browserify/snap.manifest.json | 2 +- packages/examples/packages/manage-state/snap.manifest.json | 2 +- .../src/features/snaps/legacy-state/LegacyState.tsx | 6 ++++-- .../features/snaps/legacy-state/components/ClearData.tsx | 2 +- .../features/snaps/legacy-state/components/SendData.tsx | 7 ++----- packages/test-snaps/src/features/snaps/state/State.tsx | 2 ++ packages/test-snaps/src/features/snaps/state/index.ts | 2 -- 7 files changed, 11 insertions(+), 12 deletions(-) diff --git a/packages/examples/packages/browserify/snap.manifest.json b/packages/examples/packages/browserify/snap.manifest.json index 6f6b0f910f..f12d23acb3 100644 --- a/packages/examples/packages/browserify/snap.manifest.json +++ b/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", diff --git a/packages/examples/packages/manage-state/snap.manifest.json b/packages/examples/packages/manage-state/snap.manifest.json index bd60d0cc21..5baed7f48e 100644 --- a/packages/examples/packages/manage-state/snap.manifest.json +++ b/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", diff --git a/packages/test-snaps/src/features/snaps/legacy-state/LegacyState.tsx b/packages/test-snaps/src/features/snaps/legacy-state/LegacyState.tsx index 37ad672112..665369b800 100644 --- a/packages/test-snaps/src/features/snaps/legacy-state/LegacyState.tsx +++ b/packages/test-snaps/src/features/snaps/legacy-state/LegacyState.tsx @@ -5,8 +5,8 @@ import { MANAGE_STATE_SNAP_ID, MANAGE_STATE_PORT, MANAGE_STATE_VERSION, - useSnapState, -} from '../state'; +} from '../state/constants'; +import { useSnapState } from '../state/hooks'; import { ClearData, SendData } from './components'; export const LegacyState: FunctionComponent = () => { @@ -21,6 +21,7 @@ export const LegacyState: FunctionComponent = () => { version={MANAGE_STATE_VERSION} testId="manage-state" > +

Encrypted state

{JSON.stringify(encryptedState, null, 2)} @@ -30,6 +31,7 @@ export const LegacyState: FunctionComponent = () => { +

Unencrypted state

{JSON.stringify(unencryptedState, null, 2)} diff --git a/packages/test-snaps/src/features/snaps/legacy-state/components/ClearData.tsx b/packages/test-snaps/src/features/snaps/legacy-state/components/ClearData.tsx index 31965c07da..e1cf0d88db 100644 --- a/packages/test-snaps/src/features/snaps/legacy-state/components/ClearData.tsx +++ b/packages/test-snaps/src/features/snaps/legacy-state/components/ClearData.tsx @@ -5,7 +5,7 @@ import { Button } from 'react-bootstrap'; import { Tag, useInvokeMutation } from '../../../../api'; import { Result } from '../../../../components'; import { getSnapId } from '../../../../utils'; -import { MANAGE_STATE_PORT, MANAGE_STATE_SNAP_ID } from '../../state'; +import { MANAGE_STATE_PORT, MANAGE_STATE_SNAP_ID } from '../../state/constants'; export const ClearData: FunctionComponent<{ encrypted: boolean }> = ({ encrypted, diff --git a/packages/test-snaps/src/features/snaps/legacy-state/components/SendData.tsx b/packages/test-snaps/src/features/snaps/legacy-state/components/SendData.tsx index 53c2e89cf6..70eea82e0b 100644 --- a/packages/test-snaps/src/features/snaps/legacy-state/components/SendData.tsx +++ b/packages/test-snaps/src/features/snaps/legacy-state/components/SendData.tsx @@ -6,11 +6,8 @@ import { Button, Form } from 'react-bootstrap'; import { Tag, useInvokeMutation } from '../../../../api'; import { Result } from '../../../../components'; import { getSnapId } from '../../../../utils'; -import { - MANAGE_STATE_PORT, - MANAGE_STATE_SNAP_ID, - useSnapState, -} from '../../state'; +import { MANAGE_STATE_PORT, MANAGE_STATE_SNAP_ID } from '../../state/constants'; +import { useSnapState } from '../../state/hooks'; export const SendData: FunctionComponent<{ encrypted: boolean }> = ({ encrypted, diff --git a/packages/test-snaps/src/features/snaps/state/State.tsx b/packages/test-snaps/src/features/snaps/state/State.tsx index 2e29b4d93a..d0e7dcc5a5 100644 --- a/packages/test-snaps/src/features/snaps/state/State.tsx +++ b/packages/test-snaps/src/features/snaps/state/State.tsx @@ -21,6 +21,7 @@ export const State: FunctionComponent = () => { version={MANAGE_STATE_VERSION} testId="state" > +

Encrypted state

 {
       
       
 
+      
Unencrypted state

       
         
Date: Wed, 18 Dec 2024 15:27:28 +0100
Subject: [PATCH 22/31] Move forbidden keys to shared array

---
 packages/snaps-rpc-methods/jest.config.js            | 2 +-
 packages/snaps-rpc-methods/src/permitted/getState.ts | 4 ++--
 packages/snaps-rpc-methods/src/permitted/setState.ts | 4 ++--
 packages/snaps-rpc-methods/src/utils.ts              | 2 ++
 4 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/packages/snaps-rpc-methods/jest.config.js b/packages/snaps-rpc-methods/jest.config.js
index 8d5a260d32..bcc30dee74 100644
--- a/packages/snaps-rpc-methods/jest.config.js
+++ b/packages/snaps-rpc-methods/jest.config.js
@@ -13,7 +13,7 @@ module.exports = deepmerge(baseConfig, {
       branches: 92.98,
       functions: 97.46,
       lines: 97.94,
-      statements: 97.52,
+      statements: 97.53,
     },
   },
 });
diff --git a/packages/snaps-rpc-methods/src/permitted/getState.ts b/packages/snaps-rpc-methods/src/permitted/getState.ts
index a4ce156021..60ddcd2555 100644
--- a/packages/snaps-rpc-methods/src/permitted/getState.ts
+++ b/packages/snaps-rpc-methods/src/permitted/getState.ts
@@ -15,7 +15,7 @@ import { hasProperty, isObject, type Json } from '@metamask/utils';
 
 import { manageStateBuilder } from '../restricted/manageState';
 import type { MethodHooksObject } from '../utils';
-import { StateKeyStruct } from '../utils';
+import { FORBIDDEN_KEYS, StateKeyStruct } from '../utils';
 
 const hookNames: MethodHooksObject = {
   hasPermission: true,
@@ -164,7 +164,7 @@ export function get(
   // eslint-disable-next-line @typescript-eslint/prefer-for-of
   for (let i = 0; i < keys.length; i++) {
     const currentKey = keys[i];
-    if (['__proto__', 'constructor'].includes(currentKey)) {
+    if (FORBIDDEN_KEYS.includes(currentKey)) {
       throw rpcErrors.invalidParams(
         'Invalid params: Key contains forbidden characters.',
       );
diff --git a/packages/snaps-rpc-methods/src/permitted/setState.ts b/packages/snaps-rpc-methods/src/permitted/setState.ts
index a048206e1e..be7f119a7f 100644
--- a/packages/snaps-rpc-methods/src/permitted/setState.ts
+++ b/packages/snaps-rpc-methods/src/permitted/setState.ts
@@ -16,7 +16,7 @@ import { isObject, assert, JsonStruct, type Json } from '@metamask/utils';
 
 import { manageStateBuilder } from '../restricted/manageState';
 import type { MethodHooksObject } from '../utils';
-import { StateKeyStruct } from '../utils';
+import { FORBIDDEN_KEYS, StateKeyStruct } from '../utils';
 
 const hookNames: MethodHooksObject = {
   hasPermission: true,
@@ -228,7 +228,7 @@ export function set(
 
   for (let i = 0; i < keys.length; i++) {
     const currentKey = keys[i];
-    if (['__proto__', 'constructor'].includes(currentKey)) {
+    if (FORBIDDEN_KEYS.includes(currentKey)) {
       throw rpcErrors.invalidParams(
         'Invalid params: Key contains forbidden characters.',
       );
diff --git a/packages/snaps-rpc-methods/src/utils.ts b/packages/snaps-rpc-methods/src/utils.ts
index 57f8e67ed3..ae76a7f16a 100644
--- a/packages/snaps-rpc-methods/src/utils.ts
+++ b/packages/snaps-rpc-methods/src/utils.ts
@@ -21,6 +21,8 @@ import { keccak_256 as keccak256 } from '@noble/hashes/sha3';
 
 const HARDENED_VALUE = 0x80000000;
 
+export const FORBIDDEN_KEYS = ['constructor', '__proto__', 'prototype'];
+
 /**
  * Maps an interface with method hooks to an object, using the keys of the
  * interface, and `true` as value. This ensures that the `methodHooks` object

From 3ad1e63cc5cb5615a4249678ef567a0ec2217ff9 Mon Sep 17 00:00:00 2001
From: Maarten Zuidhoorn 
Date: Wed, 18 Dec 2024 16:01:28 +0100
Subject: [PATCH 23/31] Fix some minor issues

---
 .../browserify-plugin/snap.manifest.json      |  2 +-
 packages/snaps-rpc-methods/jest.config.js     |  6 +-
 .../src/permitted/getState.test.ts            |  4 --
 .../src/permitted/getState.ts                 |  2 +-
 .../src/permitted/setState.test.ts            | 55 +++++++++++++++++++
 .../src/permitted/setState.ts                 |  6 +-
 6 files changed, 64 insertions(+), 11 deletions(-)

diff --git a/packages/examples/packages/browserify-plugin/snap.manifest.json b/packages/examples/packages/browserify-plugin/snap.manifest.json
index ee74bc2427..551a11c430 100644
--- a/packages/examples/packages/browserify-plugin/snap.manifest.json
+++ b/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",
diff --git a/packages/snaps-rpc-methods/jest.config.js b/packages/snaps-rpc-methods/jest.config.js
index bcc30dee74..c2944819dc 100644
--- a/packages/snaps-rpc-methods/jest.config.js
+++ b/packages/snaps-rpc-methods/jest.config.js
@@ -10,10 +10,10 @@ module.exports = deepmerge(baseConfig, {
   ],
   coverageThreshold: {
     global: {
-      branches: 92.98,
+      branches: 93.28,
       functions: 97.46,
-      lines: 97.94,
-      statements: 97.53,
+      lines: 98.03,
+      statements: 97.61,
     },
   },
 });
diff --git a/packages/snaps-rpc-methods/src/permitted/getState.test.ts b/packages/snaps-rpc-methods/src/permitted/getState.test.ts
index ad7b256137..4c4ba226cb 100644
--- a/packages/snaps-rpc-methods/src/permitted/getState.test.ts
+++ b/packages/snaps-rpc-methods/src/permitted/getState.test.ts
@@ -279,10 +279,6 @@ describe('get', () => {
     expect(get(object, 'a.b.c')).toBe('value');
   });
 
-  it('returns the object if the key is empty', () => {
-    expect(get(object, '')).toBe(object);
-  });
-
   it('returns `null` if the object is `null`', () => {
     expect(get(null, '')).toBeNull();
   });
diff --git a/packages/snaps-rpc-methods/src/permitted/getState.ts b/packages/snaps-rpc-methods/src/permitted/getState.ts
index 60ddcd2555..92603d0232 100644
--- a/packages/snaps-rpc-methods/src/permitted/getState.ts
+++ b/packages/snaps-rpc-methods/src/permitted/getState.ts
@@ -153,7 +153,7 @@ export function get(
   value: Record | null,
   key?: string | undefined,
 ): Json {
-  if (key === undefined || key === '') {
+  if (key === undefined) {
     return value;
   }
 
diff --git a/packages/snaps-rpc-methods/src/permitted/setState.test.ts b/packages/snaps-rpc-methods/src/permitted/setState.test.ts
index 9441f2c7b6..e8f2a94fea 100644
--- a/packages/snaps-rpc-methods/src/permitted/setState.test.ts
+++ b/packages/snaps-rpc-methods/src/permitted/setState.test.ts
@@ -285,6 +285,61 @@ describe('snap_setState', () => {
       });
     });
 
+    it('throws if the encrypted parameter is invalid', async () => {
+      const { implementation } = setStateHandler;
+
+      const getSnapState = jest.fn().mockReturnValue({
+        foo: 'bar',
+      });
+
+      const updateSnapState = jest.fn().mockReturnValue(null);
+      const getUnlockPromise = jest.fn().mockResolvedValue(undefined);
+      const hasPermission = jest.fn().mockReturnValue(true);
+
+      const hooks = {
+        getSnapState,
+        updateSnapState,
+        getUnlockPromise,
+        hasPermission,
+      };
+
+      const engine = new JsonRpcEngine();
+
+      engine.push((request, response, next, end) => {
+        const result = implementation(
+          request as JsonRpcRequest,
+          response as PendingJsonRpcResponse,
+          next,
+          end,
+          hooks,
+        );
+
+        result?.catch(end);
+      });
+
+      const response = await engine.handle({
+        jsonrpc: '2.0',
+        id: 1,
+        method: 'snap_setState',
+        params: {
+          key: 'foo',
+          value: 'bar',
+          encrypted: 'baz',
+        },
+      });
+
+      expect(response).toStrictEqual({
+        jsonrpc: '2.0',
+        id: 1,
+        error: {
+          code: errorCodes.rpc.invalidParams,
+          message:
+            'Invalid params: At path: encrypted -- Expected a value of type `boolean`, but received: `"baz"`.',
+          stack: expect.any(String),
+        },
+      });
+    });
+
     it('throws if `key` is not provided and `value` is not an object', async () => {
       const { implementation } = setStateHandler;
 
diff --git a/packages/snaps-rpc-methods/src/permitted/setState.ts b/packages/snaps-rpc-methods/src/permitted/setState.ts
index be7f119a7f..521c7398a2 100644
--- a/packages/snaps-rpc-methods/src/permitted/setState.ts
+++ b/packages/snaps-rpc-methods/src/permitted/setState.ts
@@ -137,7 +137,7 @@ async function setStateImplementation(
       await getUnlockPromise(true);
     }
 
-    const newState = await getNewState(key, value, getSnapState);
+    const newState = await getNewState(key, value, encrypted, getSnapState);
     await updateSnapState(newState, encrypted);
     response.result = null;
   } catch (error) {
@@ -180,12 +180,14 @@ function getValidatedParams(params?: unknown) {
  *
  * @param key - The key to set.
  * @param value - The value to set the key to.
+ * @param encrypted - Whether the state is encrypted.
  * @param getSnapState - The `getSnapState` hook.
  * @returns The new state of the Snap.
  */
 async function getNewState(
   key: string | undefined,
   value: Json,
+  encrypted: boolean,
   getSnapState: SetStateHooks['getSnapState'],
 ) {
   if (key === undefined) {
@@ -193,7 +195,7 @@ async function getNewState(
     return value;
   }
 
-  const state = await getSnapState(false);
+  const state = await getSnapState(encrypted);
   return set(state, key, value);
 }
 

From bf52606a65b7460b9d17e3dfa56789359616259c Mon Sep 17 00:00:00 2001
From: Maarten Zuidhoorn 
Date: Wed, 18 Dec 2024 16:03:48 +0100
Subject: [PATCH 24/31] Test calls to getSnapState hook

---
 packages/snaps-rpc-methods/src/permitted/setState.test.ts | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/packages/snaps-rpc-methods/src/permitted/setState.test.ts b/packages/snaps-rpc-methods/src/permitted/setState.test.ts
index e8f2a94fea..5c7fa82a30 100644
--- a/packages/snaps-rpc-methods/src/permitted/setState.test.ts
+++ b/packages/snaps-rpc-methods/src/permitted/setState.test.ts
@@ -63,6 +63,7 @@ describe('snap_setState', () => {
       });
 
       expect(getUnlockPromise).toHaveBeenCalled();
+      expect(getSnapState).toHaveBeenCalledWith(true);
       expect(updateSnapState).toHaveBeenCalledWith({ foo: 'baz' }, true);
 
       expect(response).toStrictEqual({
@@ -116,6 +117,7 @@ describe('snap_setState', () => {
       });
 
       expect(getUnlockPromise).toHaveBeenCalled();
+      expect(getSnapState).not.toHaveBeenCalled();
       expect(updateSnapState).toHaveBeenCalledWith({ foo: 'baz' }, true);
 
       expect(response).toStrictEqual({
@@ -169,6 +171,7 @@ describe('snap_setState', () => {
       });
 
       expect(getUnlockPromise).not.toHaveBeenCalled();
+      expect(getSnapState).toHaveBeenCalledWith(false);
       expect(updateSnapState).toHaveBeenCalledWith(
         {
           foo: 'baz',

From 05929baf725ebb42083bf4a3ce1d254350182c1d Mon Sep 17 00:00:00 2001
From: Maarten Zuidhoorn 
Date: Wed, 18 Dec 2024 16:07:44 +0100
Subject: [PATCH 25/31] Update
 packages/snaps-rpc-methods/src/permitted/getState.test.ts

Co-authored-by: Frederik Bolding 
---
 packages/snaps-rpc-methods/src/permitted/getState.test.ts | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/packages/snaps-rpc-methods/src/permitted/getState.test.ts b/packages/snaps-rpc-methods/src/permitted/getState.test.ts
index 4c4ba226cb..7be63e258f 100644
--- a/packages/snaps-rpc-methods/src/permitted/getState.test.ts
+++ b/packages/snaps-rpc-methods/src/permitted/getState.test.ts
@@ -297,7 +297,7 @@ describe('get', () => {
     );
   });
 
-  it('returns `null` if the key is a constructor pollution attempt', () => {
+  it('throws an error if the key is a constructor pollution attempt', () => {
     expect(() => get(object, 'constructor.polluted')).toThrow(
       'Invalid params: Key contains forbidden characters.',
     );

From b99cb200fda624dd284353e429d8bd79a9458c6a Mon Sep 17 00:00:00 2001
From: Maarten Zuidhoorn 
Date: Wed, 18 Dec 2024 16:09:57 +0100
Subject: [PATCH 26/31] Update
 packages/snaps-rpc-methods/src/permitted/setState.ts

Co-authored-by: Frederik Bolding 
---
 packages/snaps-rpc-methods/src/permitted/setState.ts | 1 -
 1 file changed, 1 deletion(-)

diff --git a/packages/snaps-rpc-methods/src/permitted/setState.ts b/packages/snaps-rpc-methods/src/permitted/setState.ts
index 521c7398a2..94f46634f5 100644
--- a/packages/snaps-rpc-methods/src/permitted/setState.ts
+++ b/packages/snaps-rpc-methods/src/permitted/setState.ts
@@ -66,7 +66,6 @@ export type SetStateHooks = {
   /**
    * Update the state of the requesting Snap.
    *
-   * @param snapId - The ID of the Snap.
    * @param newState - The new state of the Snap.
    * @param encrypted - Whether the state should be encrypted.
    */

From d9c7cd80fc73e85af81f2fd63a2fa07c004c3f80 Mon Sep 17 00:00:00 2001
From: Maarten Zuidhoorn 
Date: Wed, 18 Dec 2024 16:10:18 +0100
Subject: [PATCH 27/31] Update
 packages/snaps-rpc-methods/src/permitted/setState.ts

Co-authored-by: Frederik Bolding 
---
 packages/snaps-rpc-methods/src/permitted/setState.ts | 1 -
 1 file changed, 1 deletion(-)

diff --git a/packages/snaps-rpc-methods/src/permitted/setState.ts b/packages/snaps-rpc-methods/src/permitted/setState.ts
index 94f46634f5..c65b30cca6 100644
--- a/packages/snaps-rpc-methods/src/permitted/setState.ts
+++ b/packages/snaps-rpc-methods/src/permitted/setState.ts
@@ -50,7 +50,6 @@ export type SetStateHooks = {
   /**
    * Get the state of the requesting Snap.
    *
-   * @param snapId - The ID of the Snap.
    * @param encrypted - Whether the state is encrypted.
    * @returns The current state of the Snap.
    */

From 0c8be2d12805eec822a7a1befc75de170992cbfc Mon Sep 17 00:00:00 2001
From: Maarten Zuidhoorn 
Date: Wed, 18 Dec 2024 16:32:10 +0100
Subject: [PATCH 28/31] Don't allow overwriting non-objects

---
 packages/snaps-rpc-methods/jest.config.js     |  4 ++--
 .../src/permitted/setState.test.ts            | 22 ++++++++++++++++++
 .../src/permitted/setState.ts                 | 23 +++++++++++++++----
 3 files changed, 43 insertions(+), 6 deletions(-)

diff --git a/packages/snaps-rpc-methods/jest.config.js b/packages/snaps-rpc-methods/jest.config.js
index c2944819dc..fb93d5d6f9 100644
--- a/packages/snaps-rpc-methods/jest.config.js
+++ b/packages/snaps-rpc-methods/jest.config.js
@@ -10,10 +10,10 @@ module.exports = deepmerge(baseConfig, {
   ],
   coverageThreshold: {
     global: {
-      branches: 93.28,
+      branches: 93.33,
       functions: 97.46,
       lines: 98.03,
-      statements: 97.61,
+      statements: 97.62,
     },
   },
 });
diff --git a/packages/snaps-rpc-methods/src/permitted/setState.test.ts b/packages/snaps-rpc-methods/src/permitted/setState.test.ts
index 5c7fa82a30..f14373e980 100644
--- a/packages/snaps-rpc-methods/src/permitted/setState.test.ts
+++ b/packages/snaps-rpc-methods/src/permitted/setState.test.ts
@@ -487,6 +487,28 @@ describe('set', () => {
     });
   });
 
+  it('allows overwriting if a parent key is `null`', () => {
+    const object = {
+      nested: null,
+    };
+
+    expect(set(object, 'nested.key', 'newValue')).toStrictEqual({
+      nested: {
+        key: 'newValue',
+      },
+    });
+  });
+
+  it('throws if a parent key is not an object', () => {
+    const object = {
+      nested: 'value',
+    };
+
+    expect(() => set(object, 'nested.key', 'newValue')).toThrow(
+      'Invalid params: Cannot overwrite non-object value.',
+    );
+  });
+
   it('throws an error if the key is a prototype pollution attempt', () => {
     expect(() => set({}, '__proto__.polluted', 'value')).toThrow(
       'Invalid params: Key contains forbidden characters.',
diff --git a/packages/snaps-rpc-methods/src/permitted/setState.ts b/packages/snaps-rpc-methods/src/permitted/setState.ts
index c65b30cca6..0692f231cd 100644
--- a/packages/snaps-rpc-methods/src/permitted/setState.ts
+++ b/packages/snaps-rpc-methods/src/permitted/setState.ts
@@ -12,7 +12,13 @@ import {
   StructError,
 } from '@metamask/superstruct';
 import type { PendingJsonRpcResponse, JsonRpcRequest } from '@metamask/utils';
-import { isObject, assert, JsonStruct, type Json } from '@metamask/utils';
+import {
+  hasProperty,
+  isObject,
+  assert,
+  JsonStruct,
+  type Json,
+} from '@metamask/utils';
 
 import { manageStateBuilder } from '../restricted/manageState';
 import type { MethodHooksObject } from '../utils';
@@ -239,9 +245,18 @@ export function set(
       return requiredObject;
     }
 
-    currentObject[currentKey] = isObject(currentObject[currentKey])
-      ? currentObject[currentKey]
-      : {};
+    if (
+      !hasProperty(currentObject, currentKey) ||
+      currentObject[currentKey] === null
+    ) {
+      currentObject[currentKey] = {};
+    }
+
+    if (!isObject(currentObject[currentKey])) {
+      throw rpcErrors.invalidParams(
+        'Invalid params: Cannot overwrite non-object value.',
+      );
+    }
 
     currentObject = currentObject[currentKey] as Record;
   }

From bc395c12cdb88d4298ef288506dc458dc10559bf Mon Sep 17 00:00:00 2001
From: Maarten Zuidhoorn 
Date: Wed, 18 Dec 2024 16:36:37 +0100
Subject: [PATCH 29/31] Remove unnecessary code

---
 packages/snaps-rpc-methods/jest.config.js            |  4 ++--
 .../snaps-rpc-methods/src/permitted/setState.test.ts | 12 ------------
 packages/snaps-rpc-methods/src/permitted/setState.ts |  7 +------
 3 files changed, 3 insertions(+), 20 deletions(-)

diff --git a/packages/snaps-rpc-methods/jest.config.js b/packages/snaps-rpc-methods/jest.config.js
index fb93d5d6f9..21a9f370bb 100644
--- a/packages/snaps-rpc-methods/jest.config.js
+++ b/packages/snaps-rpc-methods/jest.config.js
@@ -10,10 +10,10 @@ module.exports = deepmerge(baseConfig, {
   ],
   coverageThreshold: {
     global: {
-      branches: 93.33,
+      branches: 93.3,
       functions: 97.46,
       lines: 98.03,
-      statements: 97.62,
+      statements: 97.61,
     },
   },
 });
diff --git a/packages/snaps-rpc-methods/src/permitted/setState.test.ts b/packages/snaps-rpc-methods/src/permitted/setState.test.ts
index f14373e980..759722d9bf 100644
--- a/packages/snaps-rpc-methods/src/permitted/setState.test.ts
+++ b/packages/snaps-rpc-methods/src/permitted/setState.test.ts
@@ -461,18 +461,6 @@ describe('set', () => {
     });
   });
 
-  it('overwrites the state in an existing object', () => {
-    const object = {
-      nested: {
-        key: 'oldValue',
-      },
-    };
-
-    expect(set(object, undefined, { foo: 'bar' })).toStrictEqual({
-      foo: 'bar',
-    });
-  });
-
   it('overwrites the nested state in an existing object', () => {
     const object = {
       nested: {
diff --git a/packages/snaps-rpc-methods/src/permitted/setState.ts b/packages/snaps-rpc-methods/src/permitted/setState.ts
index 0692f231cd..a402b2e920 100644
--- a/packages/snaps-rpc-methods/src/permitted/setState.ts
+++ b/packages/snaps-rpc-methods/src/permitted/setState.ts
@@ -220,14 +220,9 @@ async function getNewState(
 export function set(
   // eslint-disable-next-line @typescript-eslint/default-param-last
   object: Record | null,
-  key: string | undefined,
+  key: string,
   value: Json,
 ): JsonObject {
-  if (key === undefined) {
-    assert(isObject(value));
-    return value;
-  }
-
   const keys = key.split('.');
   const requiredObject = object ?? {};
   let currentObject: Record = requiredObject;

From baa00901b20722fe72c79a5a87547485963214dc Mon Sep 17 00:00:00 2001
From: Maarten Zuidhoorn 
Date: Wed, 18 Dec 2024 18:40:40 +0100
Subject: [PATCH 30/31] Update
 packages/snaps-rpc-methods/src/permitted/setState.ts

Co-authored-by: Frederik Bolding 
---
 packages/snaps-rpc-methods/src/permitted/setState.ts | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/packages/snaps-rpc-methods/src/permitted/setState.ts b/packages/snaps-rpc-methods/src/permitted/setState.ts
index a402b2e920..5227d69391 100644
--- a/packages/snaps-rpc-methods/src/permitted/setState.ts
+++ b/packages/snaps-rpc-methods/src/permitted/setState.ts
@@ -245,9 +245,7 @@ export function set(
       currentObject[currentKey] === null
     ) {
       currentObject[currentKey] = {};
-    }
-
-    if (!isObject(currentObject[currentKey])) {
+    } else if (!isObject(currentObject[currentKey])) {
       throw rpcErrors.invalidParams(
         'Invalid params: Cannot overwrite non-object value.',
       );

From 65625fb8fc9e8079a2b6575d329c618ef5051c8e Mon Sep 17 00:00:00 2001
From: Maarten Zuidhoorn 
Date: Wed, 18 Dec 2024 18:48:35 +0100
Subject: [PATCH 31/31] Coverage

---
 packages/snaps-rpc-methods/jest.config.js | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/packages/snaps-rpc-methods/jest.config.js b/packages/snaps-rpc-methods/jest.config.js
index 21a9f370bb..3a4658da9f 100644
--- a/packages/snaps-rpc-methods/jest.config.js
+++ b/packages/snaps-rpc-methods/jest.config.js
@@ -10,7 +10,7 @@ module.exports = deepmerge(baseConfig, {
   ],
   coverageThreshold: {
     global: {
-      branches: 93.3,
+      branches: 93.33,
       functions: 97.46,
       lines: 98.03,
       statements: 97.61,