Add support for chain spec modifier commands

.gitlab-ci.yml

@@ -459,3 +459,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

@@ -32,6 +32,7 @@ The network config can be provided both in `json` or `toml` format and each sect
   - `local_path`: string;
   - `remote_name`: string;
 - `default_resources`: (Object) **Only** available in `kubernetes`, represent the resources `limits`/`reservations` needed by the nodes by default.
+- `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`:
@@ -72,6 +73,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 }];
@@ -663,64 +683,130 @@ export async function getChainIdFromSpec(specPath: string): Promise<string> {
   });
 }
 
-export async function customizePlainRelayChain(
-  specPath: string,
-  networkSpec: ComputedNetwork,
-): Promise<void> {
+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 {
-    // Relay-chain spec customization logic
-    const plainRelayChainSpec = readAndParseChainSpec(specPath);
-    const keyType = specHaveSessionsKeys(plainRelayChainSpec)
-      ? "session"
-      : "aura";
-
-    // Clear all defaults
-    clearAuthorities(specPath);
-
-    // add balances for nodes
-    await addBalances(specPath, networkSpec.relaychain.nodes);
-
-    // add authorities for nodes
-    const validatorKeys = [];
-    for (const node of networkSpec.relaychain.nodes) {
-      if (node.validator) {
-        validatorKeys.push(node.accounts.sr_stash.address);
-
-        if (keyType === "session") {
-          const key = getNodeKey(node);
-          await addAuthority(specPath, node, key);
+    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 {
-          await addAuraAuthority(specPath, node.name, node.accounts!);
-          await addGrandpaAuthority(specPath, node.name, node.accounts!);
+          reject(new Error(`Process returned error code ${code}!`));
         }
+      });
 
-        await addStaking(specPath, node);
-      }
-    }
+      processes["mutator"].on("error", (err) => {
+        reject(err);
+      });
+    });
 
-    if (networkSpec.relaychain.randomNominatorsCount) {
-      await generateNominators(
-        specPath,
-        networkSpec.relaychain.randomNominatorsCount,
-        networkSpec.relaychain.maxNominations,
-        validatorKeys,
-      );
-    }
+    // 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;
+  }
+}
 
-    if (networkSpec.relaychain.genesis) {
-      await changeGenesisConfig(specPath, networkSpec.relaychain.genesis);
-    }
+export async function customizePlainRelayChain(
+  specPath: string,
+  networkSpec: ComputedNetwork,
+): Promise<void> {
+  // Relay-chain spec customization logic
+  const plainRelayChainSpec = readAndParseChainSpec(specPath);
+  const keyType = specHaveSessionsKeys(plainRelayChainSpec)
+    ? "session"
+    : "aura";
+
+  // Clear all defaults
+  clearAuthorities(specPath);
+
+  // add balances for nodes
+  await addBalances(specPath, networkSpec.relaychain.nodes);
+
+  // add authorities for nodes
+  const validatorKeys = [];
+  for (const node of networkSpec.relaychain.nodes) {
+    if (node.validator) {
+      validatorKeys.push(node.accounts.sr_stash.address);
+
+      if (keyType === "session") {
+        const key = getNodeKey(node);
+        await addAuthority(specPath, node, key);
+      } else {
+        await addAuraAuthority(specPath, node.name, node.accounts!);
+        await addGrandpaAuthority(specPath, node.name, node.accounts!);
+      }
 
-    if (networkSpec.hrmp_channels) {
-      await addHrmpChannelsToGenesis(specPath, networkSpec.hrmp_channels);
+      await addStaking(specPath, node);
     }
-  } catch (err) {
-    console.log(
-      `\n ${decorators.red("Unexpected error: ")} \t ${decorators.bright(
-        err,
-      )}\n`,
+  }
+
+  if (networkSpec.relaychain.randomNominatorsCount) {
+    await generateNominators(
+      specPath,
+      networkSpec.relaychain.randomNominatorsCount,
+      networkSpec.relaychain.maxNominations,
+      validatorKeys,
     );
   }
+
+  if (networkSpec.relaychain.genesis) {
+    await changeGenesisConfig(specPath, networkSpec.relaychain.genesis);
+  }
+
+  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);
+  }
 }
 export default {
   addAuraAuthority,