Support forking a custom relay chain

README.md

@@ -103,6 +103,31 @@ zombie-bite bite -r kusama --rc-upgrade ./kusama_runtime.wasm --and-spawn --appl
 zombie-bite spawn -d /tmp/base_path --apply-upgrade
 ```
 
+#### Forking a relay that is not a public network
+
+`-r` also takes `custom%<name>%<rpc_endpoint>%<chain_spec_path>`, for a relay zombie-bite has no built-in knowledge of:
+
+```sh
+zombie-bite bite -d /tmp/base_path \
+  -r custom%previewnet%wss://previewnet.example.com%/path/to/previewnet.json \
+  -p custom%2000%wss://para.example.com%/path/to/para.json
+```
+
+The name is what the artifacts are named after, the endpoint is what the bite reads state and metadata from, and the chain-spec is what the node is started with. There is no built-in host config for such a relay, so the endpoint has to be reachable — `Configuration::ActiveConfig` is read from it.
+
+A relay name that is not one of `polkadot`, `kusama`, `paseo` or `westend` is treated as a custom relay rather than silently falling back to polkadot.
+
+#### Cores and messaging state
+
+- `--para-cores <para_id>=<cores>` overrides how many cores a parachain gets (defaults mirror the live networks, e.g. asset-hub takes 3 for elastic scaling). The relay's validator count follows the total.
+- `--keep-messaging-state` keeps the inherited HRMP/DMP state instead of clearing it. Only correct when the relay's parachains are exactly the ones being bitten, so both snapshots agree on channel heads; on a shared relay the mismatch makes cumulus panic with `HRMP head mismatch`.
+
+#### Overrides are checked against the runtime
+
+Storage keys are derived from pallet and item names, and every value is decoded against its real on-chain type and required to re-encode byte-identically, so a renamed item or changed type fails the bite instead of silently landing as something else. Items the runtime does not have are skipped — except ones you asked for explicitly (a carried upgrade, a wasm override, `ZOMBIE_SUDO`), which are errors. `HostConfiguration` is patched from the live value (only `num_cores` changes) rather than replaced, so executor params, async backing and `max_pov_size` of the bitten chain are preserved.
+
+Metadata and the live values are read at the block being bitten (`--rc-bite-at` / a para's `bite_at`), so they match the state being imported. Parachains use a default public endpoint when no `rpc_endpoint` is configured; if it can't be reached, the bite still runs with a warning and those overrides go unverified. Custom parachains are only verified when their config supplies an `rpc_endpoint`.
+
 #### Spawn
 
 Spawn a new instance of the _bited_ network with the following cmd:

src/cli.rs

@@ -1,3 +1,4 @@
+use anyhow::{anyhow, bail};
 use clap::{Parser, Subcommand};
 use std::{
     env,
@@ -7,7 +8,9 @@ use std::{
 };
 use tracing::{trace, warn};
 
-use crate::config::{Parachain, Relaychain, Upgrades, ZombieBiteConfig};
+use crate::config::{
+    BiteOptions, CoresOverride, Parachain, Relaychain, Upgrades, ZombieBiteConfig,
+};
 
 #[derive(Parser, Debug)]
 #[command(author, version, about, long_about = None)]
@@ -26,7 +29,9 @@ pub enum Commands {
         /// The network will be using for bite
         /// If not specified, will use the value from config.
         /// If not in config, defaults to polkadot.
-        #[arg(short = 'r', long = "rc", value_parser = clap::builder::PossibleValuesParser::new(["polkadot", "kusama", "paseo", "westend"]))]
+        /// For a relay that is not a public network use:
+        /// custom%<name>%<rpc_endpoint>%<chain_spec_path>
+        #[arg(short = 'r', long = "rc", verbatim_doc_comment)]
         relay: Option<String>,
         /// If provided we will override the runtime as part of the process of 'bite'
         /// The resulting network will be running with this runtime.
@@ -47,6 +52,15 @@ pub enum Commands {
         /// for every carried upgrade and wait until it enacts.
         #[arg(long, default_value_t = false, verbatim_doc_comment)]
         apply_upgrade: bool,
+        /// Keep the inherited HRMP/DMP state instead of clearing it. Only correct
+        /// when the relay's parachains are exactly the ones being bitten, so the
+        /// two snapshots agree on channel heads.
+        #[arg(long, default_value_t = false, verbatim_doc_comment)]
+        keep_messaging_state: bool,
+        /// Override the cores assigned to a parachain, format: <para_id>=<cores>
+        /// Can be set multiple times, once per para.
+        #[arg(long = "para-cores", verbatim_doc_comment)]
+        para_cores: Vec<String>,
         /// If provided we will _bite_ the live network at the supplied block hieght
         #[arg(long = "rc-bite-at", verbatim_doc_comment)]
         relay_bite_at: Option<u32>,
@@ -154,8 +168,8 @@ pub struct ResolvedBiteConfig {
     pub parachains: Vec<Parachain>,
     pub base_path: PathBuf,
     pub and_spawn: bool,
-    pub upgrades: Upgrades,
     pub apply_upgrade: bool,
+    pub opts: BiteOptions,
 }
 
 #[derive(Debug)]
@@ -178,6 +192,8 @@ pub fn resolve_bite_config(
     relay_upgrade: Option<String>,
     para_upgrade: Vec<String>,
     apply_upgrade: bool,
+    keep_messaging_state: bool,
+    para_cores: Vec<String>,
 ) -> Result<ResolvedBiteConfig, anyhow::Error> {
     // Load config file if provided
     let config_file = if let Some(path) = config_path {
@@ -196,8 +212,9 @@ pub fn resolve_bite_config(
         "polkadot".to_string()
     };
 
-    let relaychain = if relay_runtime.is_some() || rc_sync_url.is_some() || relay_bite_at.is_some()
-    {
+    let relaychain = if relay_network.starts_with("custom%") {
+        resolve_custom_relaychain(&relay_network, relay_runtime.clone(), relay_bite_at)?
+    } else if relay_runtime.is_some() || rc_sync_url.is_some() || relay_bite_at.is_some() {
         // CLI args provided, use them
         Relaychain::new_with_values(&relay_network, relay_runtime, rc_sync_url, relay_bite_at)
     } else if let Some(ref config) = config_file {
@@ -281,20 +298,19 @@ pub fn resolve_bite_config(
     };
 
     // Resolve upgrades (CLI overrides config file)
+    let mut para_upgrades = std::collections::HashMap::new();
+    for entry in &para_upgrade {
+        let (id, path) = entry.split_once('=').ok_or_else(|| {
+            anyhow!("--para-upgrade must be <para_id>=<wasm_path>, got '{entry}'")
+        })?;
+        let id: u32 = id
+            .parse()
+            .map_err(|_| anyhow!("invalid para_id '{id}' in --para-upgrade"))?;
+        para_upgrades.insert(id, path.to_string());
+    }
     let mut upgrades = Upgrades {
         relay: relay_upgrade,
-        paras: para_upgrade
-            .iter()
-            .map(|entry| {
-                let (id, path) = entry.split_once('=').unwrap_or_else(|| {
-                    panic!("--para-upgrade format must be <para_id>=<wasm_path>, got: {entry}")
-                });
-                let id: u32 = id
-                    .parse()
-                    .unwrap_or_else(|_| panic!("Invalid para_id '{id}' in --para-upgrade"));
-                (id, path.to_string())
-            })
-            .collect(),
+        paras: para_upgrades,
     };
     if let Some(ref config) = config_file {
         if upgrades.relay.is_none() {
@@ -315,13 +331,57 @@ pub fn resolve_bite_config(
         false
     };
 
+    // Per-para cores: CLI entries win over the config file's `cores`.
+    let mut cores: CoresOverride = CoresOverride::new();
+    if let Some(ref config) = config_file {
+        for para_cfg in config.parachains.as_deref().unwrap_or_default() {
+            if let (Some(c), Some(para)) = (para_cfg.cores, para_cfg.to_parachain()) {
+                cores.insert(para.id(), c);
+            }
+        }
+    }
+    for entry in &para_cores {
+        let (id, c) = entry
+            .split_once('=')
+            .ok_or_else(|| anyhow!("--para-cores must be <para_id>=<cores>, got '{entry}'"))?;
+        let id: u32 = id
+            .parse()
+            .map_err(|_| anyhow!("invalid para_id '{id}' in --para-cores"))?;
+        let c: u32 = c
+            .parse()
+            .map_err(|_| anyhow!("invalid cores '{c}' in --para-cores"))?;
+        if c == 0 {
+            bail!("--para-cores {id}=0: a parachain with no cores can't have blocks backed");
+        }
+        cores.insert(id, c);
+    }
+    // A core count for a para that is not part of the bite is a typo, not a
+    // silently ignorable no-op.
+    for id in cores.keys() {
+        if !resolved_parachains.iter().any(|para| para.id() == *id) {
+            bail!("--para-cores/config sets cores for para {id}, which is not part of this bite");
+        }
+    }
+
+    let resolved_keep_messaging = if keep_messaging_state {
+        true
+    } else if let Some(ref config) = config_file {
+        config.keep_messaging_state.unwrap_or(false)
+    } else {
+        false
+    };
+
     Ok(ResolvedBiteConfig {
         relaychain,
         parachains: resolved_parachains,
         base_path: resolved_base_path,
         and_spawn: resolved_and_spawn,
-        upgrades,
         apply_upgrade: resolved_apply_upgrade,
+        opts: BiteOptions {
+            upgrades,
+            cores,
+            keep_messaging_state: resolved_keep_messaging,
+        },
     })
 }
 
@@ -369,6 +429,29 @@ pub fn resolve_spawn_config(
     })
 }
 
+/// custom%<name>%<rpc_endpoint>%<chain_spec_path>
+fn resolve_custom_relaychain(
+    s: &str,
+    maybe_override: Option<String>,
+    maybe_bite_at: Option<u32>,
+) -> Result<Relaychain, anyhow::Error> {
+    let parts: Vec<&str> = s.splitn(4, '%').collect();
+    if parts.len() != 4 {
+        bail!("custom relay must be custom%<name>%<rpc_endpoint>%<chain_spec_path>, got '{s}'");
+    }
+    let (name, rpc, chain_spec) = (parts[1], parts[2], parts[3]);
+    if name.is_empty() || rpc.is_empty() || chain_spec.is_empty() {
+        bail!("custom relay needs a name, an rpc endpoint and a chain-spec path, got '{s}'");
+    }
+    Ok(Relaychain::new_custom(
+        name,
+        chain_spec,
+        rpc,
+        maybe_override,
+        maybe_bite_at,
+    ))
+}
+
 fn resolve_custom_parachain(s: &str) -> Parachain {
     let parts: Vec<&str> = s.splitn(5, '%').collect();
     trace!("custom parts: {parts:?}");
@@ -461,4 +544,45 @@ mod test {
         let s = "custom%3392%wss://kusama-yap-3392.example.com:1234%/path/to/chain-spec.json%abc";
         let _para = resolve_custom_parachain(s);
     }
+    #[test]
+    fn custom_relay_works() {
+        let rc = resolve_custom_relaychain(
+            "custom%previewnet%wss://previewnet.example.com%/path/to/previewnet.json",
+            None,
+            Some(42),
+        )
+        .unwrap();
+
+        assert_eq!(rc.as_chain_string(), "previewnet");
+        assert_eq!(rc.chain_spec(), Some("/path/to/previewnet.json"));
+        // a custom relay is passed to the node as a spec path, not a name
+        assert_eq!(rc.chain_arg(), "/path/to/previewnet.json");
+        assert_eq!(rc.rpc_endpoint(), "wss://previewnet.example.com");
+        assert_eq!(rc.sync_endpoint(), "wss://previewnet.example.com");
+        assert_eq!(rc.at_block(), Some(42));
+        assert!(rc.is_custom());
+    }
+
+    #[test]
+    fn custom_relay_needs_every_part() {
+        for bad in [
+            "custom%previewnet%wss://previewnet.example.com",
+            "custom%previewnet%%/path/to/spec.json",
+            "custom%%wss://x%/path/to/spec.json",
+        ] {
+            assert!(
+                resolve_custom_relaychain(bad, None, None).is_err(),
+                "should reject '{bad}'"
+            );
+        }
+    }
+
+    #[test]
+    fn unknown_relay_name_keeps_its_name() {
+        // helper subcommands only get the name back as a string, and the
+        // artifacts are named after it
+        let rc = Relaychain::new("previewnet");
+        assert_eq!(rc.as_chain_string(), "previewnet");
+        assert!(rc.is_custom());
+    }
 }