diff --git a/crates/configuration/src/shared/node.rs b/crates/configuration/src/shared/node.rs index 03613d28f..59239d0ba 100644 --- a/crates/configuration/src/shared/node.rs +++ b/crates/configuration/src/shared/node.rs @@ -106,6 +106,12 @@ pub struct NodeConfig { /// or long form (e.g., "audi_sr") with explicit schema (sr, ed, ec). #[serde(default)] keystore_key_types: Vec, + /// Chain spec session key types to inject. + /// Supports short form (e.g., "aura") using predefined schemas, + /// or long form (e.g., "aura_sr") with explicit schema (sr, ed, ec). + /// When empty, uses the default session keys from the chain spec. + #[serde(default)] + chain_spec_key_types: Vec, } impl Serialize for NodeConfig { @@ -194,6 +200,12 @@ impl Serialize for NodeConfig { state.serialize_field("keystore_key_types", &self.keystore_key_types)?; } + if self.chain_spec_key_types.is_empty() { + state.skip_field("chain_spec_key_typese")?; + } else { + state.serialize_field("chain_spec_key_types", &self.chain_spec_key_types)?; + } + state.skip_field("chain_context")?; state.end() } @@ -364,6 +376,15 @@ impl NodeConfig { pub fn keystore_key_types(&self) -> Vec<&str> { self.keystore_key_types.iter().map(String::as_str).collect() } + + /// Chain spec session key types to inject. + /// Returns the list of key type specifications (short form like "aura" or long form like "aura_sr"). + pub fn chain_spec_key_types(&self) -> Vec<&str> { + self.chain_spec_key_types + .iter() + .map(String::as_str) + .collect() + } } /// A node configuration builder, used to build a [`NodeConfig`] declaratively with fields validation. @@ -401,6 +422,7 @@ impl Default for NodeConfigBuilder { node_log_path: None, keystore_path: None, keystore_key_types: vec![], + chain_spec_key_types: vec![], }, validation_context: Default::default(), errors: vec![], @@ -834,6 +856,42 @@ impl NodeConfigBuilder { ) } + /// Set the chain spec session key types to inject. + /// + /// Each key type can be specified in short form (e.g., "aura") using predefined schemas + /// (defaults to `sr` if no predefined schema exists for the key type), + /// or in long form (e.g., "aura_sr") with an explicit schema (sr, ed, ec). + /// + /// When specified, only these keys will be injected into the chain spec session keys. + /// When empty, uses the default session keys from the chain spec. + /// + /// # Examples + /// + /// ``` + /// use zombienet_configuration::shared::{node::NodeConfigBuilder, types::ChainDefaultContext}; + /// + /// let config = NodeConfigBuilder::new(ChainDefaultContext::default(), Default::default()) + /// .with_name("node") + /// .with_chain_spec_key_types(vec!["aura", "grandpa", "babe_sr"]) + /// .build() + /// .unwrap(); + /// + /// assert_eq!( + /// config.chain_spec_key_types(), + /// &["aura", "grandpa", "babe_sr"] + /// ); + /// ``` + pub fn with_chain_spec_key_types(self, key_types: Vec>) -> Self { + Self::transition( + NodeConfig { + chain_spec_key_types: key_types.into_iter().map(|k| k.into()).collect(), + ..self.config + }, + self.validation_context, + self.errors, + ) + } + /// Seals the builder and returns a [`NodeConfig`] if there are no validation errors, else returns errors. pub fn build(self) -> Result)> { if !self.errors.is_empty() { diff --git a/crates/examples/examples/chain_spec_key_types.rs b/crates/examples/examples/chain_spec_key_types.rs new file mode 100644 index 000000000..0b9edf458 --- /dev/null +++ b/crates/examples/examples/chain_spec_key_types.rs @@ -0,0 +1,116 @@ +//! Example demonstrating custom chain spec session key types. +//! +//! This example shows how to customize the session keys that are injected into the chain spec +//! for each validator node. This is useful when you need specific key types with specific +//! cryptographic schemes for your runtime. +//! +//! # Chain Spec Key Types +//! +//! There are 2 ways to specify key types: +//! +//! - **Short form**: `aura` - uses the predefined schema for the key type +//! - **Long form**: `aura_sr` - uses the explicitly specified schema (sr, ed, ec) +//! +//! ## Schemas +//! +//! - `sr` - Sr25519 +//! - `ed` - Ed25519 +//! - `ec` - ECDSA +//! +//! ## Predefined Key Type Schemas +//! +//! | Key Type | Default Schema | Description | +//! |----------|---------------|-------------| +//! | `babe` | sr | BABE consensus | +//! | `im_online` | sr | I'm Online | +//! | `parachain_validator` | sr | Parachain validator | +//! | `authority_discovery` | sr | Authority discovery | +//! | `para_validator` | sr | Para validator | +//! | `para_assignment` | sr | Para assignment | +//! | `aura` | sr (ed for asset-hub-polkadot) | AURA consensus | +//! | `nimbus` | sr | Nimbus consensus | +//! | `vrf` | sr | VRF | +//! | `grandpa` | ed | GRANDPA finality | +//! | `beefy` | ec | BEEFY | +//! +//! # Usage +//! +//! ```ignore +//! .with_validator(|node| { +//! node.with_name("alice") +//! // Only inject aura and grandpa keys into chain spec +//! .with_chain_spec_key_types(vec!["aura", "grandpa"]) +//! }) +//! .with_validator(|node| { +//! node.with_name("bob") +//! // Override grandpa to use sr25519 instead of ed25519 +//! .with_chain_spec_key_types(vec!["aura", "grandpa_sr", "babe"]) +//! }) +//! ``` +use futures::StreamExt; +use zombienet_sdk::{subxt, NetworkConfigBuilder, NetworkConfigExt}; + +#[tokio::main] +async fn main() -> Result<(), Box> { + tracing_subscriber::fmt::init(); + + let keys = vec![ + "babe", // default sr scheme + "grandpa", // default ed scheme + "im_online", // default sr scheme + "authority_discovery", // default sr scheme + "para_validator", // default sr scheme + "para_assignment", // default sr scheme + "beefy", // default ec scheme + // add a custom key type with explicit scheme + "custom_ec", // custom key with ecdsa scheme + ]; + let network = NetworkConfigBuilder::new() + .with_relaychain(|r| { + r.with_chain("rococo-local") + .with_default_command("polkadot") + .with_default_image("docker.io/parity/polkadot:v1.20.2") + .with_validator(|node| node.with_name("alice")) + .with_validator(|node| { + node.with_name("bob") + .with_chain_spec_key_types(keys.clone()) + }) + .with_validator(|node| { + node.with_name("charlie") + .with_chain_spec_key_types(keys.clone()) + }) + }) + .build() + .unwrap() + .spawn_docker() + .await?; + + println!("🚀🚀🚀🚀 network deployed"); + + let base_dir = network + .base_dir() + .ok_or("Failed to get network base directory")?; + + println!("📁 Network base directory: {}", base_dir); + println!(); + + println!("📋 Chain spec key types configuration:"); + println!(" - alice: Default keys (all standard session keys)"); + println!(" - bob: Custom keys (babe, grandpa, im_online, authority_discovery)"); + println!(" - charlie: Custom keys with overrides (babe, grandpa_sr, custom_ec)"); + println!(); + + let alice = network.get_node("alice")?; + let client = alice.wait_client::().await?; + let mut finalized_blocks = client.blocks().subscribe_finalized().await?.take(3); + + println!("⏳ Waiting for finalized blocks..."); + + while let Some(block) = finalized_blocks.next().await { + println!("✅ Finalized block {}", block?.header().number); + } + + network.destroy().await?; + + Ok(()) +} diff --git a/crates/orchestrator/src/generators.rs b/crates/orchestrator/src/generators.rs index e7913eabf..950e429a5 100644 --- a/crates/orchestrator/src/generators.rs +++ b/crates/orchestrator/src/generators.rs @@ -5,6 +5,7 @@ pub mod para_artifact; mod arg_filter; mod bootnode_addr; +mod chain_spec_key_types; mod command; mod identity; mod keystore; diff --git a/crates/orchestrator/src/generators/chain_spec.rs b/crates/orchestrator/src/generators/chain_spec.rs index d10a272db..ca7e67ba6 100644 --- a/crates/orchestrator/src/generators/chain_spec.rs +++ b/crates/orchestrator/src/generators/chain_spec.rs @@ -20,8 +20,12 @@ use support::{constants::THIS_IS_A_BUG, fs::FileSystem, replacer::apply_replacem use tokio::process::Command; use tracing::{debug, info, trace, warn}; -use super::errors::GeneratorError; +use super::{ + chain_spec_key_types::{parse_chain_spec_key_types, ChainSpecKeyType}, + errors::GeneratorError, +}; use crate::{ + generators::keystore_key_types::KeyScheme, network_spec::{node::NodeSpec, parachain::ParachainSpec, relaychain::RelaychainSpec}, ScopedFilesystem, }; @@ -1386,6 +1390,20 @@ fn add_balances( } } +/// Gets the address for a given key scheme from the node's accounts. +fn get_address_for_scheme(node: &NodeSpec, scheme: KeyScheme) -> String { + let account_key = scheme.account_key(); + node.accounts + .accounts + .get(account_key) + .expect(&format!( + "'{}' account should be set at spec computation {THIS_IS_A_BUG}", + account_key + )) + .address + .clone() +} + fn get_node_keys( node: &NodeSpec, session_key: SessionKeyType, @@ -1427,6 +1445,45 @@ fn get_node_keys( (account_to_use.clone(), account_to_use, keys) } + +/// Generates session keys for a node with custom key types. +/// Returns (account, account, keys_map) tuple. +fn get_node_keys_with_custom_types( + node: &NodeSpec, + session_key: SessionKeyType, + custom_key_types: &[ChainSpecKeyType], +) -> GenesisNodeKey { + let sr_account = node.accounts.accounts.get("sr").unwrap(); + let sr_stash = node.accounts.accounts.get("sr_stash").unwrap(); + let eth_account = node.accounts.accounts.get("eth").unwrap(); + + // key_name -> address + let mut keys = HashMap::new(); + for key_type in custom_key_types { + let scheme = key_type.scheme; + let account_key = scheme.account_key(); + let address = node + .accounts + .accounts + .get(account_key) + .expect(&format!( + "'{}' account should be set at spec computation {THIS_IS_A_BUG}", + account_key + )) + .address + .clone(); + keys.insert(key_type.key_name.clone(), address); + } + + let account_to_use = match session_key { + SessionKeyType::Default => sr_account.address.clone(), + SessionKeyType::Stash => sr_stash.address.clone(), + SessionKeyType::Evm => format!("0x{}", eth_account.public_key), + }; + + (account_to_use.clone(), account_to_use, keys) +} + fn add_authorities( runtime_config_ptr: &str, chain_spec_json: &mut serde_json::Value, @@ -1442,7 +1499,15 @@ fn add_authorities( if let Some(session_keys) = val.pointer_mut("/session/keys") { let keys: Vec = nodes .iter() - .map(|node| get_node_keys(node, session_key, asset_hub_polkadot)) + .map(|node| { + if let Some(custom_key_types) = + parse_chain_spec_key_types(&node.chain_spec_key_types, asset_hub_polkadot) + { + get_node_keys_with_custom_types(node, session_key, &custom_key_types) + } else { + get_node_keys(node, session_key, asset_hub_polkadot) + } + }) .collect(); *session_keys = json!(keys); } else { @@ -2175,4 +2240,64 @@ mod tests { assert_eq!(new_balances_map.len(), balances_map.len() + 1); } + + #[test] + fn get_node_keys_with_custom_types_works() { + use super::super::{chain_spec_key_types::ChainSpecKeyType, keystore_key_types::KeyScheme}; + + let mut name = String::from("alice"); + let seed = format!("//{}{name}", name.remove(0).to_uppercase()); + let accounts = NodeAccounts { + accounts: generators::generate_node_keys(&seed).unwrap(), + seed, + }; + let node = NodeSpec { + name, + accounts, + ..Default::default() + }; + + let custom_key_types = vec![ + ChainSpecKeyType::new("aura", KeyScheme::Ed), + ChainSpecKeyType::new("grandpa", KeyScheme::Sr), + ]; + + let node_key = + get_node_keys_with_custom_types(&node, SessionKeyType::Default, &custom_key_types); + + // Account should be sr (default) + assert_eq!(node_key.0, node.accounts.accounts["sr"].address); + assert_eq!(node_key.1, node.accounts.accounts["sr"].address); + + // Keys should use custom schemes + assert_eq!(node_key.2["aura"], node.accounts.accounts["ed"].address); + assert_eq!(node_key.2["grandpa"], node.accounts.accounts["sr"].address); + } + + #[test] + fn get_node_keys_with_custom_types_stash_works() { + use super::super::{chain_spec_key_types::ChainSpecKeyType, keystore_key_types::KeyScheme}; + + let mut name = String::from("alice"); + let seed = format!("//{}{name}", name.remove(0).to_uppercase()); + let accounts = NodeAccounts { + accounts: generators::generate_node_keys(&seed).unwrap(), + seed, + }; + let node = NodeSpec { + name, + accounts, + ..Default::default() + }; + + let custom_key_types = vec![ChainSpecKeyType::new("aura", KeyScheme::Sr)]; + + let node_key = + get_node_keys_with_custom_types(&node, SessionKeyType::Stash, &custom_key_types); + + // Account should be sr_stash (stash derivation) + assert_eq!(node_key.0, node.accounts.accounts["sr_stash"].address); + assert_eq!(node_key.1, node.accounts.accounts["sr_stash"].address); + assert_eq!(node_key.2["aura"], node.accounts.accounts["sr"].address); + } } diff --git a/crates/orchestrator/src/generators/chain_spec_key_types.rs b/crates/orchestrator/src/generators/chain_spec_key_types.rs new file mode 100644 index 000000000..0cce80097 --- /dev/null +++ b/crates/orchestrator/src/generators/chain_spec_key_types.rs @@ -0,0 +1,205 @@ +use std::collections::HashMap; + +use super::keystore_key_types::KeyScheme; + +/// A parsed chain spec session key type. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ChainSpecKeyType { + /// The key type name as it appears in the chain spec (e.g., "aura", "grandpa", "babe"). + pub key_name: String, + /// The cryptographic scheme to use for this key type. + pub scheme: KeyScheme, +} + +impl ChainSpecKeyType { + pub fn new(key_name: impl Into, scheme: KeyScheme) -> Self { + Self { + key_name: key_name.into(), + scheme, + } + } +} + +/// Returns the default predefined key schemes for known chain spec key types. +/// Special handling for `aura` when `is_asset_hub_polkadot` is true. +fn get_predefined_schemes(is_asset_hub_polkadot: bool) -> HashMap<&'static str, KeyScheme> { + let mut schemes = HashMap::new(); + + // aura has special handling for asset-hub-polkadot + if is_asset_hub_polkadot { + schemes.insert("aura", KeyScheme::Ed); + } else { + schemes.insert("aura", KeyScheme::Sr); + } + + // SR25519 keys + schemes.insert("babe", KeyScheme::Sr); + schemes.insert("im_online", KeyScheme::Sr); + schemes.insert("parachain_validator", KeyScheme::Sr); + schemes.insert("authority_discovery", KeyScheme::Sr); + schemes.insert("para_validator", KeyScheme::Sr); + schemes.insert("para_assignment", KeyScheme::Sr); + schemes.insert("nimbus", KeyScheme::Sr); + schemes.insert("vrf", KeyScheme::Sr); + + // ED25519 keys + schemes.insert("grandpa", KeyScheme::Ed); + + // ECDSA keys + schemes.insert("beefy", KeyScheme::Ec); + + schemes +} + +/// Parses a single chain spec key type specification string. +/// +/// Supports two formats: +/// - Short: `aura` - uses predefined default scheme (defaults to `sr` if not predefined) +/// - Long: `aura_sr` - uses explicit scheme +/// +/// Returns `None` if the spec is empty. +pub fn parse_key_spec( + spec: &str, + predefined: &HashMap<&str, KeyScheme>, +) -> Option { + let spec = spec.trim(); + + if spec.is_empty() { + return None; + } + + if let Some((key_name, scheme_str)) = spec.rsplit_once('_') { + if let Ok(scheme) = KeyScheme::try_from(scheme_str) { + return Some(ChainSpecKeyType::new(key_name, scheme)); + } + // If not a valid scheme, define whole string as the key name + } + + let scheme = predefined.get(spec).copied().unwrap_or(KeyScheme::Sr); + Some(ChainSpecKeyType::new(spec, scheme)) +} + +/// Parses a list of chain spec key type specifications. +/// +/// Each spec can be in short form (`aura`) or long form (`aura_sr`). +/// Invalid specs are silently ignored. +/// +/// If the input list is empty, returns `None` to indicate that default +/// chain spec behavior should be used. +pub fn parse_chain_spec_key_types>( + specs: &[T], + is_asset_hub_polkadot: bool, +) -> Option> { + if specs.is_empty() { + return None; + } + + let predefined_schemes = get_predefined_schemes(is_asset_hub_polkadot); + + let parsed: Vec = specs + .iter() + .filter_map(|spec| parse_key_spec(spec.as_ref(), &predefined_schemes)) + .collect(); + + if parsed.is_empty() { + None + } else { + Some(parsed) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parse_chain_spec_key_types_returns_none_when_empty() { + let specs: Vec = vec![]; + let result = parse_chain_spec_key_types(&specs, false); + assert!(result.is_none()); + } + + #[test] + fn parse_chain_spec_key_types_parses_short_form() { + let specs = vec!["aura".to_string(), "grandpa".to_string()]; + let result = parse_chain_spec_key_types(&specs, false).unwrap(); + + assert_eq!(result.len(), 2); + assert_eq!(result[0], ChainSpecKeyType::new("aura", KeyScheme::Sr)); + assert_eq!(result[1], ChainSpecKeyType::new("grandpa", KeyScheme::Ed)); + } + + #[test] + fn parse_chain_spec_key_types_parses_long_form() { + let specs = vec![ + "aura_ed".to_string(), // Override aura to use ed + "grandpa_sr".to_string(), // Override grandpa to use sr + "custom_ec".to_string(), // Custom key with ec + ]; + let result = parse_chain_spec_key_types(&specs, false).unwrap(); + + assert_eq!(result.len(), 3); + assert_eq!(result[0], ChainSpecKeyType::new("aura", KeyScheme::Ed)); + assert_eq!(result[1], ChainSpecKeyType::new("grandpa", KeyScheme::Sr)); + assert_eq!(result[2], ChainSpecKeyType::new("custom", KeyScheme::Ec)); + } + + #[test] + fn parse_chain_spec_key_types_mixed_forms() { + let specs = vec![ + "aura".to_string(), // Short form - uses default sr + "grandpa_sr".to_string(), // Long form - override to sr + "babe".to_string(), // Short form - uses default sr + ]; + let result = parse_chain_spec_key_types(&specs, false).unwrap(); + + assert_eq!(result.len(), 3); + assert_eq!(result[0], ChainSpecKeyType::new("aura", KeyScheme::Sr)); + assert_eq!(result[1], ChainSpecKeyType::new("grandpa", KeyScheme::Sr)); + assert_eq!(result[2], ChainSpecKeyType::new("babe", KeyScheme::Sr)); + } + + #[test] + fn parse_chain_spec_key_types_asset_hub_polkadot() { + let specs = vec!["aura".to_string(), "babe".to_string()]; + + let result = parse_chain_spec_key_types(&specs, false).unwrap(); + assert_eq!(result[0].scheme, KeyScheme::Sr); + + let result = parse_chain_spec_key_types(&specs, true).unwrap(); + assert_eq!(result[0].scheme, KeyScheme::Ed); + } + + #[test] + fn parse_chain_spec_key_types_unknown_key_defaults_to_sr() { + let specs = vec!["unknown_key".to_string()]; + let result = parse_chain_spec_key_types(&specs, false).unwrap(); + + assert_eq!(result.len(), 1); + assert_eq!( + result[0], + ChainSpecKeyType::new("unknown_key", KeyScheme::Sr) + ); + } + + #[test] + fn parse_chain_spec_key_types_handles_underscore_in_key_name() { + let specs = vec![ + "im_online".to_string(), // Known key with underscore + "para_validator".to_string(), // Known key with underscore + "my_custom_key_sr".to_string(), // Custom key with underscores and explicit scheme + ]; + let result = parse_chain_spec_key_types(&specs, false).unwrap(); + + assert_eq!(result.len(), 3); + assert_eq!(result[0], ChainSpecKeyType::new("im_online", KeyScheme::Sr)); + assert_eq!( + result[1], + ChainSpecKeyType::new("para_validator", KeyScheme::Sr) + ); + assert_eq!( + result[2], + ChainSpecKeyType::new("my_custom_key", KeyScheme::Sr) + ); + } +} diff --git a/crates/orchestrator/src/network_spec.rs b/crates/orchestrator/src/network_spec.rs index 6efab45d6..44cbe9726 100644 --- a/crates/orchestrator/src/network_spec.rs +++ b/crates/orchestrator/src/network_spec.rs @@ -249,21 +249,20 @@ impl NetworkSpec { try_join_all( image_command_to_nodes_mapping .keys() - .map(|(image, command)| { - let ns_cloned = ns.clone(); - async move { - // get node available args output from image/command - let available_args = ns_cloned - .get_node_available_args((command.clone(), image.clone())) - .await?; - debug!( - "retrieved available args for image: {:?}, command: {}", - image, command - ); - - // map the result to include image and command - Ok::<_, OrchestratorError>((image.clone(), command.clone(), available_args)) - } + .map(|(image, command)| async { + let image = image.clone(); + let command = command.clone(); + // get node available args output from image/command + let available_args = ns + .get_node_available_args((command.clone(), image.clone())) + .await?; + debug!( + "retrieved available args for image: {:?}, command: {}", + image, command + ); + + // map the result to include image and command + Ok::<_, OrchestratorError>((image, command, available_args)) }) .collect::>(), ) diff --git a/crates/orchestrator/src/network_spec/node.rs b/crates/orchestrator/src/network_spec/node.rs index 98d6d4a42..4796b03a0 100644 --- a/crates/orchestrator/src/network_spec/node.rs +++ b/crates/orchestrator/src/network_spec/node.rs @@ -133,6 +133,12 @@ pub struct NodeSpec { /// Supports short form (e.g., "audi") using predefined schemas, /// or long form (e.g., "audi_sr") with explicit schema (sr, ed, ec). pub(crate) keystore_key_types: Vec, + + /// Chain spec session key types to inject. + /// Supports short form (e.g., "aura") using predefined schemas, + /// or long form (e.g., "aura_sr") with explicit schema (sr, ed, ec). + /// When empty, uses the default session keys from the chain spec. + pub(crate) chain_spec_key_types: Vec, } impl NodeSpec { @@ -237,6 +243,11 @@ impl NodeSpec { .into_iter() .map(str::to_string) .collect(), + chain_spec_key_types: node_config + .chain_spec_key_types() + .into_iter() + .map(str::to_string) + .collect(), }) } @@ -338,6 +349,7 @@ impl NodeSpec { node_log_path: None, keystore_path: None, keystore_key_types: vec![], + chain_spec_key_types: vec![], }) }