feat(dev_newBlock): Adds option to override relaychain state (#824)

* feat(dev_newBlock): Adds option to override relaychain state

* fix comment

* test: only test relayChainStateOverrides on parachains

* apply suggestions from review

packages/core/src/blockchain/inherent/parachain/validation-data.ts

@@ -98,7 +98,7 @@ export class SetValidationData implements InherentProvider {
 
     const extrinsic = await getValidationData(parent)
 
-    const newEntries: [HexString, HexString | null][] = []
+    let newEntries: [HexString, HexString | null][] = []
     const downwardMessages: DownwardMessage[] = []
     const horizontalMessages: Record<number, HorizontalMessage[]> = {}
 
@@ -268,6 +268,16 @@ export class SetValidationData implements InherentProvider {
       newEntries.push([upgradeKey, null])
     }
 
+    // Apply relay chain state overrides
+    if (params.relayChainStateOverrides) {
+      for (const [key, value] of params.relayChainStateOverrides) {
+        // Remove any entry that matches the key being overridden
+        newEntries = newEntries.filter(([k, _]) => k != key)
+        // Push override
+        newEntries.push([key, value])
+      }
+    }
+
     const { trieRootHash, nodes } = await createProof(extrinsic.relayChainState.trieNodes, newEntries)
 
     const newData = {

packages/core/src/blockchain/txpool.ts

@@ -39,6 +39,7 @@ export interface BuildBlockParams {
   horizontalMessages: Record<number, HorizontalMessage[]>
   transactions: HexString[]
   unsafeBlockHeight?: number
+  relayChainStateOverrides?: [HexString, HexString | null][]
 }
 
 export class TxPool {
@@ -177,6 +178,7 @@ export class TxPool {
     const downwardMessages = params?.downwardMessages || this.#dmp.splice(0)
     const horizontalMessages = params?.horizontalMessages || { ...this.#hrmp }
     const unsafeBlockHeight = params?.unsafeBlockHeight
+    const relayChainStateOverrides = params?.relayChainStateOverrides
     if (!params?.upwardMessages) {
       for (const id of Object.keys(this.#ump)) {
         delete this.#ump[id]
@@ -195,6 +197,7 @@ export class TxPool {
         downwardMessages,
         horizontalMessages,
         unsafeBlockHeight,
+        relayChainStateOverrides,
       })
 
       // with the latest message queue, messages could be processed in the upcoming block

packages/core/src/rpc/dev/new-block.ts

@@ -32,6 +32,7 @@ const schema = z.object({
     .optional(),
   transactions: z.array(zHex).min(1).optional(),
   unsafeBlockHeight: z.number().optional(),
+  relayChainStateOverrides: z.array(z.tuple([zHex, z.union([zHex, z.null()])])).optional(),
 })
 
 type Params = z.infer<typeof schema>
@@ -65,6 +66,10 @@ export interface NewBlockParams {
    * Build block using a specific block height (unsafe)
    */
   unsafeBlockHeight: Params['unsafeBlockHeight']
+  /**
+   * Build block using a custom relay chain state
+   */
+  relayChainStateOverrides: Params['relayChainStateOverrides']
 }
 
 /**
@@ -106,7 +111,9 @@ export interface NewBlockParams {
  * ```
  */
 export const dev_newBlock = async (context: Context, [params]: [NewBlockParams]) => {
-  const { count, to, hrmp, ump, dmp, transactions, unsafeBlockHeight } = schema.parse(params || {})
+  const { count, to, hrmp, ump, dmp, transactions, unsafeBlockHeight, relayChainStateOverrides } = schema.parse(
+    params || {},
+  )
   const now = context.chain.head.number
   const diff = to ? to - now : count
   const finalCount = diff !== undefined ? Math.max(diff, 1) : 1
@@ -124,6 +131,7 @@ export const dev_newBlock = async (context: Context, [params]: [NewBlockParams])
         upwardMessages: ump,
         downwardMessages: dmp,
         unsafeBlockHeight: i === 0 ? unsafeBlockHeight : undefined,
+        relayChainStateOverrides: relayChainStateOverrides,
       })
       .catch((error) => {
         throw new ResponseError(1, error.toString())

packages/core/src/utils/proof.ts

@@ -27,6 +27,11 @@ export const upgradeGoAheadSignal = (paraId: u32) => {
   return hash(prefix, paraId.toU8a())
 }
 
+export const upgradeRestrictionSignal = (paraId: u32) => {
+  const prefix = '0xcd710b30bd2eab0352ddcc26417aa194f27bbb460270642b5bcaf032ea04d56a'
+  return hash(prefix, paraId.toU8a())
+}
+
 export const hrmpIngressChannelIndex = (paraId: u32) => {
   const prefix = '0x6a0da05ca59913bc38a8630590f2627c1d3719f5b0b12c7105c073c507445948'
   return hash(prefix, paraId.toU8a())

packages/e2e/src/build-parachain-block.test.ts

@@ -0,0 +1,33 @@
+import { describe, expect, it } from 'vitest'
+
+import { TypeRegistry } from '@polkadot/types'
+import { decodeProof } from '@acala-network/chopsticks-core'
+import { upgradeRestrictionSignal } from '@acala-network/chopsticks-core/utils/proof.js'
+import networks from './networks.js'
+
+describe('override-relay-state-proof', async () => {
+  it('build block using relayChainStateOverrides', async () => {
+    const { ws, api, teardown } = await networks.acala()
+    const registry = new TypeRegistry()
+    const paraId = registry.createType('u32', 1000)
+
+    const keyToOverride = upgradeRestrictionSignal(paraId)
+    const value = '0x00'
+    const relayChainStateOverrides = [[keyToOverride, value]]
+
+    await ws.send('dev_newBlock', [{ relayChainStateOverrides }])
+    const block = await api.rpc.chain.getBlock()
+    const setValidationData = block.block.extrinsics
+      .find(({ method }) => method.method == 'setValidationData')
+      ?.method.toJSON().args.data
+
+    const relayParentStorageRoot = setValidationData.validationData.relayParentStorageRoot
+    const trieNodes = setValidationData.relayChainState.trieNodes
+
+    const relayChainState = await decodeProof(relayParentStorageRoot, trieNodes)
+
+    expect(relayChainState[keyToOverride]).to.be.eq(value)
+
+    await teardown()
+  })
+})