Improve state size validation

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

@@ -10,10 +10,10 @@ module.exports = deepmerge(baseConfig, {
   ],
   coverageThreshold: {
     global: {
-      branches: 93.97,
+      branches: 94.89,
       functions: 98.05,
       lines: 98.67,
-      statements: 98.25,
+      statements: 98.34,
     },
   },
 });

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

@@ -1,7 +1,11 @@
 import { JsonRpcEngine } from '@metamask/json-rpc-engine';
 import { errorCodes } from '@metamask/rpc-errors';
 import type { SetStateResult } from '@metamask/snaps-sdk';
-import type { JsonRpcRequest, PendingJsonRpcResponse } from '@metamask/utils';
+import type {
+  Json,
+  JsonRpcRequest,
+  PendingJsonRpcResponse,
+} from '@metamask/utils';
 
 import { setStateHandler, type SetStateParameters, set } from './setState';
 
@@ -396,6 +400,111 @@ describe('snap_setState', () => {
         },
       });
     });
+
+    it('throws if the new state is not JSON serialisable', async () => {
+      const { implementation } = setStateHandler;
+
+      const getSnapState = jest.fn().mockReturnValue(null);
+      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<SetStateParameters>,
+          response as PendingJsonRpcResponse<SetStateResult>,
+          next,
+          end,
+          hooks,
+        );
+
+        result?.catch(end);
+      });
+
+      const response = await engine.handle({
+        jsonrpc: '2.0',
+        id: 1,
+        method: 'snap_setState',
+        params: {
+          value: {
+            // @ts-expect-error - BigInt is not JSON serializable.
+            foo: 1n as Json,
+          },
+        },
+      });
+
+      expect(response).toStrictEqual({
+        jsonrpc: '2.0',
+        id: 1,
+        error: {
+          code: errorCodes.rpc.invalidParams,
+          message:
+            'Invalid params: At path: value -- Expected a value of type `JSON`, but received: `[object Object]`.',
+          stack: expect.any(String),
+        },
+      });
+    });
+
+    it('throws if the new state exceeds the size limit', async () => {
+      const { implementation } = setStateHandler;
+
+      const getSnapState = jest.fn().mockReturnValue(null);
+      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<SetStateParameters>,
+          response as PendingJsonRpcResponse<SetStateResult>,
+          next,
+          end,
+          hooks,
+        );
+
+        result?.catch(end);
+      });
+
+      const response = await engine.handle({
+        jsonrpc: '2.0',
+        id: 1,
+        method: 'snap_setState',
+        params: {
+          value: {
+            foo: 'foo'.repeat(21_500_000), // 64.5 MB
+          },
+        },
+      });
+
+      expect(response).toStrictEqual({
+        jsonrpc: '2.0',
+        id: 1,
+        error: {
+          code: errorCodes.rpc.invalidParams,
+          message:
+            'Invalid params: The new state must not exceed 64 MB in size.',
+          stack: expect.any(String),
+        },
+      });
+    });
   });
 });
 

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

@@ -13,14 +13,18 @@ import {
 } from '@metamask/superstruct';
 import type { PendingJsonRpcResponse, JsonRpcRequest } from '@metamask/utils';
 import {
+  getJsonSize,
   hasProperty,
   isObject,
   assert,
   JsonStruct,
   type Json,
 } from '@metamask/utils';
 
-import { manageStateBuilder } from '../restricted/manageState';
+import {
+  manageStateBuilder,
+  STORAGE_SIZE_LIMIT,
+} from '../restricted/manageState';
 import type { MethodHooksObject } from '../utils';
 import { FORBIDDEN_KEYS, StateKeyStruct } from '../utils';
 
@@ -142,6 +146,16 @@ async function setStateImplementation(
     }
 
     const newState = await getNewState(key, value, encrypted, getSnapState);
+
+    const size = getJsonSize(newState);
+    if (size > STORAGE_SIZE_LIMIT) {
+      throw rpcErrors.invalidParams({
+        message: `Invalid params: The new state must not exceed ${
+          STORAGE_SIZE_LIMIT / 1_000_000
+        } MB in size.`,
+      });
+    }
+
     await updateSnapState(newState, encrypted);
     response.result = null;
   } catch (error) {

packages/snaps-rpc-methods/src/restricted/manageState.test.ts

@@ -498,7 +498,9 @@ describe('snap_manageState', () => {
           MOCK_SMALLER_STORAGE_SIZE_LIMIT,
         ),
       ).toThrow(
-        `Invalid snap_manageState "updateState" parameter: The new state must not exceed ${MOCK_SMALLER_STORAGE_SIZE_LIMIT} bytes in size.`,
+        `Invalid snap_manageState "updateState" parameter: The new state must not exceed ${
+          MOCK_SMALLER_STORAGE_SIZE_LIMIT / 1_000_000
+        } MB in size.`,
       );
     });
   });

packages/snaps-rpc-methods/src/restricted/manageState.ts

@@ -107,7 +107,7 @@ export const manageStateBuilder = Object.freeze({
   methodHooks,
 } as const);
 
-export const STORAGE_SIZE_LIMIT = 104857600; // In bytes (100MB)
+export const STORAGE_SIZE_LIMIT = 64_000_000; // In bytes (64 MB)
 
 type GetEncryptionKeyArgs = {
   snapId: string;
@@ -257,10 +257,6 @@ export function getValidatedParams(
     if (!isObject(newState)) {
       throw rpcErrors.invalidParams({
         message: `Invalid ${method} "updateState" parameter: The new state must be a plain object.`,
-        data: {
-          receivedNewState:
-            typeof newState === 'undefined' ? 'undefined' : newState,
-        },
       });
     }
 
@@ -271,20 +267,14 @@ export function getValidatedParams(
     } catch {
       throw rpcErrors.invalidParams({
         message: `Invalid ${method} "updateState" parameter: The new state must be JSON serializable.`,
-        data: {
-          receivedNewState:
-            typeof newState === 'undefined' ? 'undefined' : newState,
-        },
       });
     }
 
     if (size > storageSizeLimit) {
       throw rpcErrors.invalidParams({
-        message: `Invalid ${method} "updateState" parameter: The new state must not exceed ${storageSizeLimit} bytes in size.`,
-        data: {
-          receivedNewState:
-            typeof newState === 'undefined' ? 'undefined' : newState,
-        },
+        message: `Invalid ${method} "updateState" parameter: The new state must not exceed ${
+          storageSizeLimit / 1_000_000
+        } MB in size.`,
       });
     }
   }