feat: add chain_spec_key_types support for custom session keys

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<String>,
+    /// 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<String>,
 }
 
 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<Initial> {
                 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<Buildable> {
         )
     }
 
+    /// 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<impl Into<String>>) -> 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<NodeConfig, (String, Vec<anyhow::Error>)> {
         if !self.errors.is_empty() {

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<dyn std::error::Error>> {
+    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::<subxt::PolkadotConfig>().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(())
+}

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;

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,
 };
@@ -1392,6 +1396,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,
@@ -1433,6 +1451,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,
@@ -1448,7 +1505,15 @@ fn add_authorities(
         if let Some(session_keys) = val.pointer_mut("/session/keys") {
             let keys: Vec<GenesisNodeKey> = 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 {
@@ -2070,4 +2135,64 @@ mod tests {
         let node_key = get_node_keys(&node, SessionKeyType::default(), true);
         assert_eq!(node_key.2["aura"], node.accounts.accounts["ed"].address);
     }
+
+    #[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);
+    }
 }