add rpc plugin schema check (#528)

packages/chopsticks/src/plugins/dry-run/rpc.ts

@@ -1,12 +1,10 @@
-import { HexString } from '@polkadot/util/types'
 import { z } from 'zod'
 
 import { Context, ResponseError } from '@acala-network/chopsticks-core'
 import { decodeStorageDiff } from '../../utils/decoder.js'
 import { generateHtmlDiff } from '../../utils/generate-html-diff.js'
+import { zHash, zHex } from '../../schema/index.js'
 
-const zHex = z.custom<HexString>((val: any) => /^0x\w+$/.test(val))
-const zHash = z.string().length(66).and(zHex)
 const zParaId = z.string().regex(/^\d+$/).transform(Number)
 
 const schema = z.object({

packages/chopsticks/src/plugins/new-block/rpc.ts

@@ -1,36 +1,70 @@
-import { Context, DownwardMessage, HorizontalMessage, ResponseError } from '@acala-network/chopsticks-core'
-import { HexString } from '@polkadot/util/types'
+import { Context, ResponseError } from '@acala-network/chopsticks-core'
+import { z } from 'zod'
+
 import { defaultLogger } from '../../logger.js'
+import { zHex } from '../../schema/index.js'
+
+const schema = z.object({
+  count: z.number().optional(),
+  to: z.number().optional(),
+  dmp: z
+    .array(
+      z.object({
+        sentAt: z.number(),
+        msg: zHex,
+      }),
+    )
+    .min(1)
+    .optional(),
+  ump: z.record(z.number(), z.array(zHex).min(1)).optional(),
+  hrmp: z
+    .record(
+      z.number(),
+      z
+        .array(
+          z.object({
+            sentAt: z.number(),
+            data: zHex,
+          }),
+        )
+        .min(1),
+    )
+    .optional(),
+  transactions: z.array(zHex).min(1).optional(),
+  unsafeBlockHeight: z.number().optional(),
+})
+
+type Params = z.infer<typeof schema>
 
 export interface NewBlockParams {
   /**
    * The number of blocks to build
    */
-  count: number
+  count: Params['count']
   /**
    * The block number to build to
    */
-  to: number
+  to: Params['to']
   /**
    * The downward messages to include in the block
    */
-  dmp: DownwardMessage[]
+  dmp: Params['dmp']
   /**
    * The upward messages to include in the block
    */
-  ump: Record<number, HexString[]>
+  ump: Params['ump']
   /**
    * The horizontal messages to include in the block
    */
-  hrmp: Record<number, HorizontalMessage[]>
+  hrmp: Params['hrmp']
   /**
    * The transactions to include in the block
    */
-  transactions: HexString[]
+  transactions: Params['transactions']
   /**
    * Build block using a specific block height (unsafe)
    */
-  unsafeBlockHeight: number
+  unsafeBlockHeight: Params['unsafeBlockHeight']
 }
 
 /**
@@ -71,16 +105,14 @@ export interface NewBlockParams {
  * await ws.send('dev_newBlock', [{ count: 2, unsafeBlockHeight: 100000001 }])
  * ```
  */
-export const rpc = async (context: Context, params: [NewBlockParams]) => {
-  const [param] = params
-  const { count, to, hrmp, ump, dmp, transactions, unsafeBlockHeight } = param || {}
+export const rpc = async (context: Context, [params]: [NewBlockParams]) => {
+  const { count, to, hrmp, ump, dmp, transactions, unsafeBlockHeight } = schema.parse(params || {})
   const now = context.chain.head.number
   const diff = to ? to - now : count
-  const finalCount = diff > 0 ? diff : 1
+  const finalCount = diff !== undefined ? Math.max(diff, 1) : 1
 
   let finalHash: string | undefined
-
-  if (unsafeBlockHeight < now) {
+  if (unsafeBlockHeight !== undefined && unsafeBlockHeight <= now) {
     throw new ResponseError(1, 'unsafeBlockHeight must be greater than current block height')
   }
 

packages/chopsticks/src/plugins/run-block/rpc.ts

@@ -14,9 +14,7 @@ import {
   runTask,
   taskHandler,
 } from '@acala-network/chopsticks-core'
-
-const zHex = z.custom<HexString>((val: any) => /^0x\w+$/.test(val))
-const zHash = z.string().length(66).and(zHex)
+import { zHash, zHex } from '../../schema/index.js'
 
 const schema = z.object({
   includeRaw: z.boolean().optional(),
@@ -116,6 +114,8 @@ export interface RunBlockResponse {
   }
 }
 
+export const name = 'runBlock'
+
 /**
  * Run a set of extrinsics on top of a block and get the storage diff
  * and optionally the parsed storage diff and block details.

packages/chopsticks/src/plugins/set-block-build-mode/rpc.ts

@@ -4,6 +4,8 @@ import { defaultLogger } from '../../logger.js'
 /**
  * Set a build block mode. See [BuildBlockMode](../core/enums/BuildBlockMode).
  *
+ * 1 - Batch, 2 - Instant, 3 - Manual
+ *
  * This function is a dev rpc handler. Use `dev_setBlockBuildMode` as the method name when calling it.
  *
  * @param context - The context object of the rpc handler
@@ -18,7 +20,7 @@ import { defaultLogger } from '../../logger.js'
  * ```
  */
 export const rpc = async (context: Context, [mode]: [BuildBlockMode]) => {
-  defaultLogger.debug({ mode }, 'dev_setBlockBuildMode')
+  defaultLogger.debug({ mode: BuildBlockMode[mode] }, 'dev_setBlockBuildMode')
 
   if (BuildBlockMode[mode] === undefined) {
     throw new ResponseError(1, `Invalid mode ${mode}`)

packages/chopsticks/src/plugins/set-head/rpc.ts

@@ -1,5 +1,10 @@
 import { Block, Context, ResponseError } from '@acala-network/chopsticks-core'
-import { HexString } from '@polkadot/util/types'
+import { z } from 'zod'
+
+import { zHash } from '../../schema/index.js'
+
+const schema = zHash.or(z.number())
+type Params = z.infer<typeof schema>
 
 /**
  * Set head.
@@ -16,7 +21,8 @@ import { HexString } from '@polkadot/util/types'
  * await ws.send('dev_setHead', [1000000])
  * ```
  */
-export const rpc = async (context: Context, [hashOrNumber]: [HexString | number]) => {
+export const rpc = async (context: Context, [params]: [Params]) => {
+  const hashOrNumber = schema.parse(params)
   let block: Block | undefined
   if (typeof hashOrNumber === 'number') {
     const blockNumber = hashOrNumber > 0 ? hashOrNumber : context.chain.head.number + hashOrNumber

packages/chopsticks/src/schema/index.ts

@@ -1,11 +1,15 @@
 import { BuildBlockMode, defaultLogger, genesisSchema, isUrl } from '@acala-network/chopsticks-core'
+import { HexString } from '@polkadot/util/types'
 import { basename, extname } from 'node:path'
 import { readFileSync } from 'node:fs'
 import { z } from 'zod'
 import _ from 'lodash'
 import axios from 'axios'
 import yaml from 'js-yaml'
 
+export const zHex = z.custom<HexString>((val: any) => /^0x\w+$/.test(val))
+export const zHash = z.string().length(66).and(zHex)
+
 export const configSchema = z
   .object({
     port: z.number().optional(),