add support for evm and fix stash account

paritytech · Sep 11, 2024 · 5a3dd46 · 5a3dd46
1 parent d15f8f0
commit 5a3dd46
Show file tree

Hide file tree

Showing 6 changed files with 81 additions and 19 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -57,6 +57,7 @@ tower = { version = "0.4" }
 tower-http = { version = "0.5" }
 tracing-subscriber = { version = "0.3" }
 glob-match = "0.2.1"
+libsecp256k1 = { version = "0.7.1", default-features = false }
 
 # Zombienet workspace crates:
 support = { package = "zombienet-support", version = "0.2.9", path = "crates/support" }

diff --git a/crates/orchestrator/Cargo.toml b/crates/orchestrator/Cargo.toml
@@ -34,6 +34,7 @@ regex = { workspace = true }
 glob-match = { workspace = true }
 async-trait = { workspace = true }
 serde = { workspace = true, features = ["derive"] }
+libsecp256k1 = { workspace = true  }
 
 # Zombienet deps
 configuration = { workspace = true }

diff --git a/crates/orchestrator/src/generators/chain_spec.rs b/crates/orchestrator/src/generators/chain_spec.rs
@@ -40,6 +40,19 @@ enum KeyType {
     Grandpa,
 }
 
+#[derive(Debug, Clone, Copy)]
+enum SessionKeyType {
+    Default,
+    Stash,
+    Evm,
+}
+
+impl Default for SessionKeyType {
+    fn default() -> Self {
+        Self::Default
+    }
+}
+
 #[derive(Debug, Clone, Serialize)]
 pub enum CommandInContext {
     Local(String),
@@ -448,7 +461,16 @@ impl ChainSpec {
                 .pointer(&format!("{}/session", pointer))
                 .is_some()
             {
-                add_authorities(&pointer, &mut chain_spec_json, &validators, false);
+                add_authorities(
+                    &pointer,
+                    &mut chain_spec_json,
+                    &validators,
+                    if para.is_evm_based {
+                        SessionKeyType::Evm
+                    } else {
+                        SessionKeyType::Default
+                    },
+                );
             } else if chain_spec_json
                 .pointer(&format!("{}/aura", pointer))
                 .is_some()
@@ -547,7 +569,12 @@ impl ChainSpec {
                 .pointer(&format!("{}/session", pointer))
                 .is_some()
             {
-                add_authorities(&pointer, &mut chain_spec_json, &validators, true);
+                add_authorities(
+                    &pointer,
+                    &mut chain_spec_json,
+                    &validators,
+                    SessionKeyType::Stash,
+                );
             }
 
             // staking && nominators
@@ -880,11 +907,16 @@ fn add_balances(
     }
 }
 
-fn get_node_keys(node: &NodeSpec, use_stash: bool, asset_hub_polkadot: bool) -> GenesisNodeKey {
+fn get_node_keys(
+    node: &NodeSpec,
+    session_key: SessionKeyType,
+    asset_hub_polkadot: bool,
+) -> GenesisNodeKey {
     let sr_account = node.accounts.accounts.get("sr").unwrap();
     let sr_stash = node.accounts.accounts.get("sr_stash").unwrap();
     let ed_account = node.accounts.accounts.get("ed").unwrap();
     let ec_account = node.accounts.accounts.get("ec").unwrap();
+    let eth_account = node.accounts.accounts.get("eth").unwrap();
     let mut keys = HashMap::new();
     for k in [
         "babe",
@@ -906,19 +938,21 @@ fn get_node_keys(node: &NodeSpec, use_stash: bool, asset_hub_polkadot: bool) ->
 
     keys.insert("grandpa".to_string(), ed_account.address.clone());
     keys.insert("beefy".to_string(), ec_account.address.clone());
+    keys.insert("eth".to_string(), eth_account.public_key.clone());
+
+    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),
+    };
 
-    let account_to_use = if use_stash { sr_stash } else { sr_account };
-    (
-        account_to_use.address.clone(),
-        account_to_use.address.clone(),
-        keys,
-    )
+    (account_to_use.clone(), account_to_use, keys)
 }
 fn add_authorities(
     runtime_config_ptr: &str,
     chain_spec_json: &mut serde_json::Value,
     nodes: &[&NodeSpec],
-    use_stash: bool,
+    session_key: SessionKeyType,
 ) {
     let asset_hub_polkadot = chain_spec_json
         .get("id")
@@ -928,7 +962,7 @@ fn add_authorities(
     if let Some(val) = chain_spec_json.pointer_mut(runtime_config_ptr) {
         let keys: Vec<GenesisNodeKey> = nodes
             .iter()
-            .map(|node| get_node_keys(node, use_stash, asset_hub_polkadot))
+            .map(|node| get_node_keys(node, session_key, asset_hub_polkadot))
             .collect();
         val["session"]["keys"] = json!(keys);
     } else {
@@ -1255,17 +1289,18 @@ mod tests {
                 node.accounts.accounts["ed"].address.clone(),
             ),
             ("beefy".into(), node.accounts.accounts["ec"].address.clone()),
+            ("eth".into(), node.accounts.accounts["eth"].address.clone()),
         ]
         .into();
 
         // Stash
         let sr_stash = &node.accounts.accounts["sr_stash"];
-        let node_key = get_node_keys(&node, true, false);
+        let node_key = get_node_keys(&node, SessionKeyType::Stash, false);
         assert_eq!(node_key.0, sr_stash.address);
         assert_eq!(node_key.1, sr_stash.address);
         assert_eq!(node_key.2, keys);
         // Non-stash
-        let node_key = get_node_keys(&node, false, false);
+        let node_key = get_node_keys(&node, SessionKeyType::Default, false);
         assert_eq!(node_key.0, sr.address);
         assert_eq!(node_key.1, sr.address);
         assert_eq!(node_key.2, keys);
@@ -1285,10 +1320,10 @@ mod tests {
             ..Default::default()
         };
 
-        let node_key = get_node_keys(&node, false, false);
+        let node_key = get_node_keys(&node, SessionKeyType::default(), false);
         assert_eq!(node_key.2["aura"], node.accounts.accounts["sr"].address);
 
-        let node_key = get_node_keys(&node, false, true);
+        let node_key = get_node_keys(&node, SessionKeyType::default(), true);
         assert_eq!(node_key.2["aura"], node.accounts.accounts["ed"].address);
     }
 }
diff --git a/crates/orchestrator/src/generators/key.rs b/crates/orchestrator/src/generators/key.rs
@@ -1,8 +1,8 @@
-use sp_core::{crypto::SecretStringError, ecdsa, ed25519, sr25519, Pair};
+use sp_core::{crypto::SecretStringError, ecdsa, ed25519, keccak_256, sr25519, Pair, H160, H256};
 
 use super::errors::GeneratorError;
 use crate::shared::types::{Accounts, NodeAccount};
-const KEYS: [&str; 4] = ["sr", "sr_stash", "ed", "ec"];
+const KEYS: [&str; 5] = ["sr", "sr_stash", "ed", "ec", "eth"];
 
 pub fn generate_pair<T: Pair>(seed: &str) -> Result<T::Pair, SecretStringError> {
     let pair = T::Pair::from_string(seed, None)?;
@@ -19,7 +19,7 @@ pub fn generate(seed: &str) -> Result<Accounts, GeneratorError> {
                 (pair.public().to_string(), hex::encode(pair.public()))
             },
             "sr_stash" => {
-                let pair = generate_pair::<sr25519::Pair>(&format!("{}/stash", seed))
+                let pair = generate_pair::<sr25519::Pair>(&format!("{}//stash", seed))
                     .map_err(|_| GeneratorError::KeyGeneration(k.into(), seed.into()))?;
                 (pair.public().to_string(), hex::encode(pair.public()))
             },
@@ -33,6 +33,19 @@ pub fn generate(seed: &str) -> Result<Accounts, GeneratorError> {
                     .map_err(|_| GeneratorError::KeyGeneration(k.into(), seed.into()))?;
                 (pair.public().to_string(), hex::encode(pair.public()))
             },
+            "eth" => {
+                let pair = generate_pair::<ecdsa::Pair>(seed)
+                    .map_err(|_| GeneratorError::KeyGeneration(k.into(), seed.into()))?;
+
+                let decompressed = libsecp256k1::PublicKey::parse_compressed(&pair.public().0)
+                    .map_err(|_| GeneratorError::KeyGeneration(k.into(), seed.into()))?
+                    .serialize();
+                let mut m = [0u8; 64];
+                m.copy_from_slice(&decompressed[1..65]);
+                let account = H160::from(H256::from(keccak_256(&m)));
+
+                (hex::encode(account), hex::encode(account))
+            },
             _ => unreachable!(),
         };
         accounts.insert(k.into(), NodeAccount::new(address, public_key));
@@ -111,13 +124,15 @@ mod tests {
         let sr_stash = pair.get("sr_stash").unwrap();
         let ed = pair.get("ed").unwrap();
         let ec = pair.get("ec").unwrap();
+        let eth = pair.get("eth").unwrap();
+
         assert_eq!(
             sr.address,
             "5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY"
         );
         assert_eq!(
             sr_stash.address,
-            "5DZnGRAr28KP4GvbuxW2cBNo9Aodcm4QKUMj3Zqj67YjYStr"
+            "5GNJqTPyNqANBkUVMN1LPPrxXnFouWXoe2wNSmmEoLctxiZY"
         );
         assert_eq!(
             ed.address,
@@ -127,5 +142,10 @@ mod tests {
             format!("0x{}", ec.public_key),
             "0x020a1091341fe5664bfa1782d5e04779689068c916b04cb365ec3153755684d9a1"
         );
+
+        assert_eq!(
+            format!("0x{}", eth.public_key),
+            "0xe04cc55ebee1cbce552f250e85c57b70b2e2625b"
+        )
     }
 }
diff --git a/crates/orchestrator/src/network_spec.rs b/crates/orchestrator/src/network_spec.rs
@@ -285,6 +285,7 @@ impl NetworkSpec {
 
 #[cfg(test)]
 mod tests {
+    use std::collections::HashMap;
 
     #[tokio::test]
     async fn small_network_config_get_spec() {

diff --git a/crates/orchestrator/src/network_spec/parachain.rs b/crates/orchestrator/src/network_spec/parachain.rs
@@ -55,6 +55,9 @@ pub struct ParachainSpec {
     /// Is the parachain cumulus-based
     pub(crate) is_cumulus_based: bool,
 
+    /// Is the parachain evm-based
+    pub(crate) is_evm_based: bool,
+
     /// Initial balance
     pub(crate) initial_balance: u128,
 
@@ -202,6 +205,7 @@ impl ParachainSpec {
                 .clone(),
             onboard_as_parachain: config.onboard_as_parachain(),
             is_cumulus_based: config.is_cumulus_based(),
+            is_evm_based: config.is_evm_based(),
             initial_balance: config.initial_balance(),
             genesis_state,
             genesis_wasm,