Add support for chain spec modifier commands

.gitlab-ci.yml

@@ -305,3 +305,34 @@ db-snapshot:
   retry: 2
   tags:
     - zombienet-polkadot-integration-test
+
+zombienet-chain-spec-custom-modification:
+  stage: deploy
+  <<: *kubernetes-env
+  image: "paritypr/zombienet:${CI_COMMIT_SHORT_SHA}"
+  rules:
+    - if: $CI_PIPELINE_SOURCE == "schedule"
+    - if: $CI_COMMIT_REF_NAME == "master"
+    - if: $CI_COMMIT_REF_NAME =~ /^[0-9]+$/ # PRs
+    - if: $CI_COMMIT_REF_NAME =~ /^v[0-9]+\.[0-9]+.*$/ # i.e. v1.0, v2.1rc1
+  # needs:
+  #   - job:                         publish-docker-pr
+
+  variables:
+    GH_DIR: "https://github.com/paritytech/zombienet/tree/${CI_COMMIT_SHORT_SHA}/tests"
+
+  before_script:
+    - echo "Zombienet Tests Custom Chain Spec Modification"
+    - echo "paritypr/zombienet:${CI_COMMIT_SHORT_SHA}"
+    - echo "${GH_DIR}"
+    - export DEBUG=zombie*
+    - export ZOMBIENET_INTEGRATION_TEST_IMAGE="docker.io/paritypr/polkadot-debug:master"
+    - export COL_IMAGE="docker.io/paritypr/colander:master"
+
+  script:
+    - /home/nonroot/zombie-net/scripts/ci/run-test-local-env-manager.sh
+      --test="0014-chain-spec-modification.zndsl"
+  allow_failure: true
+  retry: 2
+  tags:
+    - zombienet-polkadot-integration-test

README.md

@@ -270,6 +270,27 @@ export COL_IMAGE=docker.io/paritypr/colander:master
 ./zombienet-macos spawn examples/0001-small-network.toml
 ```
 
+#### Custom commands to modify the resulting chain-specs
+
+If you want to customize the chain specification, plain or raw, that one of your chains will be launched with, beyond
+what Zombienet provides by default, you can do so with `chain_spec_modifier_commands`.
+
+The `chain_spec_modifier_commands` option allows you to specify a list of CLI commands that will use, modify and return 
+the modified chain-spec before it is used to launch the network. These commands can modify the chain specification in any way you need.
+
+```toml
+...
+[[parachains]]
+id = 100
+chain_spec_modifier_commands = [[
+  "/path/to/your_custom_script.sh",
+  "{{'plain'|chainSpec}}"
+]]
+...
+```
+
+For a more in-depth example with a chain-querying tool `chainql`, you can visit the [examples](examples) directory.
+
 ##### Teardown
 
 You can teardown the network (and cleanup the used resources) by terminating the process (`ctrl+c`).

docs/src/network-definition-spec.md

@@ -34,6 +34,7 @@ The network config can be provided both in `json` or `toml` format and each sect
   - `remote_name`: string;
 - `default_resources`: (Object) **Only** available in `kubernetes`, represent the resources `limits`/`reservations` needed by the nodes by default.
 - `default_prometheus_prefix`: A parameter for customizing the metric's prefix. If parameter is placed in `relaychain` level, it will be "passed" to all `relaychain` nodes. Defaults to 'substrate'.
+- `chain_spec_modifier_commands`: (Array of `commands`, optional) Set of commands and their arguments to modify the resulting chain spec before the chain is launched with it. `Commands` are themselves arrays of strings (each argument is a string). In the arguments, `{{'plain'|chainSpec}}` will be substituted for the plain spec path and run before the raw chain spec is generated, and `{{'raw'|chainSpec}}` will be substituted for the raw chain spec path.
 - `random_nominators_count`: (number, optional), if is set _and the stacking pallet is enabled_ zombienet will generate `x` nominators and will be injected in the genesis.
 - `max_nominations`: (number, default 24), the max allowed number of nominations by a nominator. This should match the value set in the runtime (e.g Kusama is 24 and Polkadot 16).
 - `nodes`:
@@ -78,6 +79,7 @@ The network config can be provided both in `json` or `toml` format and each sect
   - `*id`: (Number) The id to assign to this parachain. Must be unique.
   - `add_to_genesis`: (Boolean, default true) flag to add parachain to genesis or register in runtime.
   - `cumulus_based`: (Boolean, default true) flag to use `cumulus` command generation.
+  - `chain_spec_modifier_commands`: (Array of `commands`, optional) Set of commands and their arguments to modify the resulting chain spec before the chain is launched with it. `Commands` are themselves arrays of strings (each argument is a string). In the arguments, `{{'plain'|chainSpec}}` will be substituted for the plain spec path and run before the raw chain spec is generated, and `{{'raw'|chainSpec}}` will be substituted for the raw chain spec path.
   - `genesis_wasm_path`: (String) Path to the wasm file to use.
   - `genesis_wasm_generator`: (String) Command to generate the wasm file.
   - `genesis_state_path`: (String) Path to the state file to use.

examples/0005-chain-spec-mutation-with-chainql.toml

@@ -0,0 +1,31 @@
+# examples/0005-chain-spec-mutation-with-chainql.toml
+[relaychain]
+default_image = "docker.io/parity/polkadot:latest"
+default_command = "polkadot"
+default_args = [ "-lparachain=debug" ]
+
+chain = "rococo-local"
+
+  [[relaychain.nodes]]
+  name = "alice"
+  validator = true
+
+  [[relaychain.nodes]]
+  name = "bob"
+  validator = true
+
+[[parachains]]
+id = 100
+# make sure to have chainql installed (`cargo install chainql`)
+# https://github.com/UniqueNetwork/chainql
+chain_spec_modifier_commands = [[
+  "chainql",
+  "--tla-code=rawSpec=import '{{'raw'|chainSpec}}'",
+  "--tla-str=pullFrom=wss://rococo-rockmine-rpc.polkadot.io:443",
+  "--trace-format=explaining",
+  "chainqlCopyBalances.jsonnet",
+]]
+
+  [[parachains.collator_groups]]
+  count = 2
+  name = "collator"

examples/chainqlCopyBalances.jsonnet

@@ -0,0 +1,53 @@
+// Get a live chain's system balances on the URL provided with `pullFrom`, and 
+// insert them into a raw spec to launch a new chain with.
+//
+// ### Arguments
+// - `rawSpec`: Path to the raw chain spec to modify
+// - `pullFrom`: URL of the chain's WS port to get the data from
+// 
+// ### Usage
+// `chainql --tla-code=rawSpec="import '/path/to/parachain-spec-raw.json'" --tla-str=pullFrom="wss://some-parachain.node:443" chainqlCopyBalances.jsonnet`
+//
+// Make sure to to have `chainql` installed: `cargo install chainql`
+
+function(rawSpec, pullFrom)
+// get the latest state of the blockchain
+local sourceChainState = cql.chain(pullFrom).latest;
+
+local
+	// store all keys under the `Account` storage of the `System` pallet
+	accounts = sourceChainState.System.Account._preloadKeys,
+	// get the encoded naming of `pallet_balances::TotalIssuance` for future use
+	totalIssuanceKey = sourceChainState.Balances._encodeKey.TotalIssuance([]),
+;
+
+// output the raw spec with the following changes
+rawSpec {
+	genesis+: {
+		raw+: {
+			// add the following entries to the `top` section
+			top+: 
+			{
+				// encode key and value of every account under `system.account` and add them to the chain spec
+				[sourceChainState.System._encodeKey.Account([key])]: 
+					sourceChainState.System._encodeValue.Account(accounts[key])
+				for key in std.objectFields(accounts)
+			} + {
+				// add to the local, already-existing total issuance the issuance of all incoming accounts.
+				// NOTE: we do not take into consideration for total issuance's funds potential overlap with the testnet's present accounts.
+				[totalIssuanceKey]: sourceChainState.Balances._encodeValue.TotalIssuance(
+					// decode the chain-spec's already existing totalIssuance
+					sourceChainState.Balances._decodeValue.TotalIssuance(super[totalIssuanceKey])
+					// iterate over and sum up the total issuance of the incoming accounts
+					+ std.foldl(
+						function(issuance, acc)
+							issuance + acc.data.free + acc.data.reserved
+						,
+						std.objectValues(accounts),
+						std.bigint('0'),
+					)
+				)
+			},
+		},
+	},
+}

javascript/packages/orchestrator/src/chainSpec.ts

@@ -6,10 +6,21 @@ import {
   getRandom,
   readDataFile,
 } from "@zombienet/utils";
+import { ChildProcessWithoutNullStreams, spawn } from "child_process";
 import crypto from "crypto";
 import fs from "fs";
+import {
+  PLAIN_CHAIN_SPEC_IN_CMD_PATTERN,
+  RAW_CHAIN_SPEC_IN_CMD_PATTERN,
+} from "./constants";
 import { generateKeyFromSeed } from "./keys";
-import { ChainSpec, ComputedNetwork, HrmpChannelsConfig, Node } from "./types";
+import {
+  ChainSpec,
+  Command,
+  ComputedNetwork,
+  HrmpChannelsConfig,
+  Node,
+} from "./types";
 const JSONbig = require("json-bigint")({ useNativeBigInt: true });
 const debug = require("debug")("zombie::chain-spec");
 
@@ -18,6 +29,15 @@ const JSONStream = require("JSONStream");
 // track 1st staking as default;
 let stakingBond: number | undefined;
 
+const processes: { [key: string]: ChildProcessWithoutNullStreams } = {};
+
+// kill any runnning processes related to non-node chain spec processing
+export async function destroyChainSpecProcesses() {
+  for (const key of Object.keys(processes)) {
+    processes[key].kill();
+  }
+}
+
 export type KeyType = "session" | "aura" | "grandpa";
 
 export type GenesisNodeKey = [string, string, { [key: string]: string }];
@@ -664,6 +684,75 @@ export async function getChainIdFromSpec(specPath: string): Promise<string> {
   });
 }
 
+export async function runCommandWithChainSpec(
+  chainSpecFullPath: string,
+  commandWithArgs: Command,
+  workingDirectory: string | URL,
+) {
+  const chainSpecSubstitutePattern = new RegExp(
+    RAW_CHAIN_SPEC_IN_CMD_PATTERN?.source +
+      "|" +
+      PLAIN_CHAIN_SPEC_IN_CMD_PATTERN?.source,
+    "gi",
+  );
+
+  const substitutedCommandArgs = commandWithArgs.map(
+    (arg) => `${arg.replaceAll(chainSpecSubstitutePattern, chainSpecFullPath)}`,
+  );
+  const chainSpecModifiedPath = chainSpecFullPath.replace(
+    ".json",
+    "-modified.json",
+  );
+
+  new CreateLogTable({ colWidths: [30, 90] }).pushToPrint([
+    [
+      decorators.green("🧪 Mutating chain spec"),
+      decorators.white(substitutedCommandArgs.join(" ")),
+    ],
+  ]);
+
+  try {
+    await new Promise<void>(function (resolve, reject) {
+      if (processes["mutator"]) {
+        processes["mutator"].kill();
+      }
+
+      // spawn the chain spec mutator thread with the command and arguments
+      processes["mutator"] = spawn(
+        substitutedCommandArgs[0],
+        substitutedCommandArgs.slice(1),
+        { cwd: workingDirectory },
+      );
+      // flush the modified spec to a different file and then copy it back into the original path
+      const spec = fs.createWriteStream(chainSpecModifiedPath);
+
+      // `pipe` since it deals with flushing and we need to guarantee that the data is flushed
+      // before we resolve the promise.
+      processes["mutator"].stdout.pipe(spec);
+
+      processes["mutator"].stderr.pipe(process.stderr);
+
+      processes["mutator"].on("close", (code) => {
+        if (code === 0) {
+          resolve();
+        } else {
+          reject(new Error(`Process returned error code ${code}!`));
+        }
+      });
+
+      processes["mutator"].on("error", (err) => {
+        reject(err);
+      });
+    });
+
+    // copy the modified file back into the original path after the mutation has completed
+    fs.copyFileSync(chainSpecModifiedPath, chainSpecFullPath);
+  } catch (e: any) {
+    console.error(`\n${decorators.red("Failed to mutate chain spec!")}`);
+    throw e;
+  }
+}
+
 export async function customizePlainRelayChain(
   specPath: string,
   networkSpec: ComputedNetwork,
@@ -716,6 +805,11 @@ export async function customizePlainRelayChain(
     if (networkSpec.hrmp_channels) {
       await addHrmpChannelsToGenesis(specPath, networkSpec.hrmp_channels);
     }
+
+    // modify the plain chain spec with any custom commands
+    for (const cmd of networkSpec.relaychain.chainSpecModifierCommands) {
+      await runCommandWithChainSpec(specPath, cmd, networkSpec.configBasePath);
+    }
   } catch (err) {
     console.log(
       `\n ${decorators.red("Unexpected error: ")} \t ${decorators.bright(

javascript/packages/orchestrator/src/configGenerator.ts

@@ -25,6 +25,8 @@ import {
   DEFAULT_WASM_GENERATE_SUBCOMMAND,
   GENESIS_STATE_FILENAME,
   GENESIS_WASM_FILENAME,
+  PLAIN_CHAIN_SPEC_IN_CMD_PATTERN,
+  RAW_CHAIN_SPEC_IN_CMD_PATTERN,
   UNDYING_COLLATOR_BIN,
   ZOMBIE_WRAPPER,
 } from "./constants";
@@ -119,6 +121,8 @@ export async function generateNetworkSpec(
       defaultImage: config.relaychain.default_image || DEFAULT_IMAGE,
       defaultCommand: config.relaychain.default_command || DEFAULT_COMMAND,
       defaultArgs: config.relaychain.default_args || [],
+      chainSpecModifierCommands: [],
+      rawChainSpecModifierCommands: [],
       randomNominatorsCount: config.relaychain?.random_nominators_count || 0,
       maxNominations:
         config.relaychain?.max_nominations || DEFAULT_MAX_NOMINATIONS,
@@ -182,6 +186,35 @@ export async function generateNetworkSpec(
         ).replace("{{DEFAULT_COMMAND}}", networkSpec.relaychain.defaultCommand);
   }
 
+  for (const cmd of config.relaychain.chain_spec_modifier_commands || []) {
+    const cmdHasRawSpec = cmd.some((arg) =>
+      RAW_CHAIN_SPEC_IN_CMD_PATTERN.test(arg),
+    );
+    const cmdHasPlainSpec = cmd.some((arg) =>
+      PLAIN_CHAIN_SPEC_IN_CMD_PATTERN.test(arg),
+    );
+
+    if (cmdHasRawSpec && cmdHasPlainSpec) {
+      console.log(
+        decorators.yellow(
+          `Chain spec modifier command references both raw and plain chain specs! Only the raw chain spec will be modified.\n\t${cmd}`,
+        ),
+      );
+    }
+
+    if (cmdHasRawSpec) {
+      networkSpec.relaychain.rawChainSpecModifierCommands.push(cmd);
+    } else if (cmdHasPlainSpec) {
+      networkSpec.relaychain.chainSpecModifierCommands.push(cmd);
+    } else {
+      console.log(
+        decorators.yellow(
+          `Chain spec modifier command does not attempt to reference a chain spec path! It will not be executed.\n\t${cmd}`,
+        ),
+      );
+    }
+  }
+
   const relayChainBootnodes: string[] = [];
   for (const node of config.relaychain.nodes || []) {
     const nodeSetup = await getNodeFromConfig(
@@ -373,6 +406,8 @@ export async function generateNetworkSpec(
         name: getUniqueName(parachain.id.toString()),
         para,
         cumulusBased: isCumulusBased,
+        chainSpecModifierCommands: [],
+        rawChainSpecModifierCommands: [],
         addToGenesis:
           parachain.add_to_genesis === undefined
             ? true
@@ -405,6 +440,35 @@ export async function generateNetworkSpec(
         }
       }
 
+      for (const cmd of parachain.chain_spec_modifier_commands || []) {
+        const cmdHasRawSpec = cmd.some((arg) =>
+          RAW_CHAIN_SPEC_IN_CMD_PATTERN.test(arg),
+        );
+        const cmdHasPlainSpec = cmd.some((arg) =>
+          PLAIN_CHAIN_SPEC_IN_CMD_PATTERN.test(arg),
+        );
+
+        if (cmdHasRawSpec && cmdHasPlainSpec) {
+          console.log(
+            decorators.yellow(
+              `Chain spec modifier command references both raw and plain chain specs! Only the raw chain spec will be modified.\n\t${cmd}`,
+            ),
+          );
+        }
+
+        if (cmdHasRawSpec) {
+          parachainSetup.rawChainSpecModifierCommands.push(cmd);
+        } else if (cmdHasPlainSpec) {
+          parachainSetup.chainSpecModifierCommands.push(cmd);
+        } else {
+          console.log(
+            decorators.yellow(
+              `Chain spec modifier command does not attempt to reference a chain spec path! It will not be executed.\n\t${cmd}`,
+            ),
+          );
+        }
+      }
+
       parachainSetup = {
         ...parachainSetup,
         ...(parachain.balance ? { balance: parachain.balance } : {}),