From 04791725259eafc86f1499bef5e56ea365d3ccf1 Mon Sep 17 00:00:00 2001
From: Javier Viola <363911+pepoviola@users.noreply.github.com>
Date: Thu, 12 Sep 2024 05:33:17 -0300
Subject: [PATCH] add support for evm and fix stash account (#256)

- Fix stash derivation (use `//`)
- Add `evm` support, you can set your parachain as evm based with
`evm_based(choice: bool)`.


Fix: #247
cc: @AlexD10S / @al3mart
---
 Cargo.toml                                    |  1 +
 crates/configuration/src/parachain.rs         | 52 ++++++++++++-
 crates/configuration/src/utils.rs             |  4 +
 .../testing/snapshots/0001-big-network.toml   |  2 +
 .../snapshots/0002-overridden-defaults.toml   |  1 +
 .../0003-small-network_w_parachain.toml       | 17 ++++
 crates/orchestrator/Cargo.toml                |  1 +
 .../orchestrator/src/generators/chain_spec.rs | 77 +++++++++++++++----
 crates/orchestrator/src/generators/key.rs     | 28 ++++++-
 .../src/network_spec/parachain.rs             |  4 +
 10 files changed, 165 insertions(+), 22 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 33ee922fb..cf83b4aea 100644
--- 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/configuration/src/parachain.rs b/crates/configuration/src/parachain.rs
index 98888e94d..33d6ff157 100644
--- a/crates/configuration/src/parachain.rs
+++ b/crates/configuration/src/parachain.rs
@@ -18,7 +18,7 @@ use crate::{
             Arg, AssetLocation, Chain, ChainDefaultContext, Command, Image, ValidationContext, U128,
         },
     },
-    utils::{default_as_true, default_initial_balance, is_false},
+    utils::{default_as_false, default_as_true, default_initial_balance, is_false},
 };
 
 /// The registration strategy that will be used for the parachain.
@@ -136,6 +136,8 @@ pub struct ParachainConfig {
     chain_spec_command_is_local: bool,
     #[serde(rename = "cumulus_based", default = "default_as_true")]
     is_cumulus_based: bool,
+    #[serde(rename = "evm_based", default = "default_as_false")]
+    is_evm_based: bool,
     #[serde(skip_serializing_if = "std::vec::Vec::is_empty", default)]
     bootnodes_addresses: Vec<Multiaddr>,
     genesis_overrides: Option<serde_json::Value>,
@@ -246,6 +248,11 @@ impl ParachainConfig {
         self.is_cumulus_based
     }
 
+    /// Whether the parachain is evm based (e.g frontier).
+    pub fn is_evm_based(&self) -> bool {
+        self.is_evm_based
+    }
+
     /// The bootnodes addresses the collators will connect to.
     pub fn bootnodes_addresses(&self) -> Vec<&Multiaddr> {
         self.bootnodes_addresses.iter().collect::<Vec<_>>()
@@ -313,6 +320,7 @@ impl<C: Context> Default for ParachainConfigBuilder<Initial, C> {
                 chain_spec_command: None,
                 chain_spec_command_is_local: false, // remote by default
                 is_cumulus_based: true,
+                is_evm_based: false,
                 bootnodes_addresses: vec![],
                 collators: vec![],
                 collator: None,
@@ -719,6 +727,18 @@ impl<C: Context> ParachainConfigBuilder<WithId, C> {
         )
     }
 
+    /// Set whether the parachain is evm based (e.g frontier /evm template)
+    pub fn evm_based(self, choice: bool) -> Self {
+        Self::transition(
+            ParachainConfig {
+                is_evm_based: choice,
+                ..self.config
+            },
+            self.validation_context,
+            self.errors,
+        )
+    }
+
     /// Set the bootnodes addresses the collators will connect to.
     pub fn with_bootnodes_addresses<T>(self, bootnodes_addresses: Vec<T>) -> Self
     where
@@ -858,6 +878,7 @@ mod tests {
             .with_genesis_state_generator("generator_state")
             .with_chain_spec_path("./path/to/chain/spec.json")
             .cumulus_based(false)
+            .evm_based(false)
             .with_bootnodes_addresses(vec![
                 "/ip4/10.41.122.55/tcp/45421",
                 "/ip4/51.144.222.10/tcp/2333",
@@ -949,6 +970,7 @@ mod tests {
             parachain_config.bootnodes_addresses(),
             bootnodes_addresses.iter().collect::<Vec<_>>()
         );
+        assert!(!parachain_config.is_evm_based());
     }
 
     #[test]
@@ -1243,9 +1265,12 @@ mod tests {
         .unwrap();
 
         let parachain = load_from_toml_small.parachains()[0];
+        let parachain_evm = load_from_toml_small.parachains()[1];
 
         assert_eq!(parachain.registration_strategy(), None);
+        assert!(!parachain.is_evm_based());
         assert_eq!(parachain.collators().len(), 1);
+        assert!(parachain_evm.is_evm_based());
     }
 
     #[test]
@@ -1260,6 +1285,31 @@ mod tests {
         assert!(config.onboard_as_parachain());
     }
 
+    #[test]
+    fn evm_based_default_to_false() {
+        let config = ParachainConfigBuilder::new(Default::default())
+            .with_id(2000)
+            .with_chain("myparachain")
+            .with_collator(|collator| collator.with_name("collator"))
+            .build()
+            .unwrap();
+
+        assert!(!config.is_evm_based());
+    }
+
+    #[test]
+    fn evm_based() {
+        let config = ParachainConfigBuilder::new(Default::default())
+            .with_id(2000)
+            .with_chain("myparachain")
+            .evm_based(true)
+            .with_collator(|collator| collator.with_name("collator"))
+            .build()
+            .unwrap();
+
+        assert!(config.is_evm_based());
+    }
+
     #[test]
     fn build_config_in_running_context() {
         let config = ParachainConfigBuilder::new_with_running(Default::default())
diff --git a/crates/configuration/src/utils.rs b/crates/configuration/src/utils.rs
index a26d7e883..42614f8b6 100644
--- a/crates/configuration/src/utils.rs
+++ b/crates/configuration/src/utils.rs
@@ -12,6 +12,10 @@ pub(crate) fn default_as_true() -> bool {
     true
 }
 
+pub(crate) fn default_as_false() -> bool {
+    false
+}
+
 pub(crate) fn default_initial_balance() -> crate::types::U128 {
     2_000_000_000_000.into()
 }
diff --git a/crates/configuration/testing/snapshots/0001-big-network.toml b/crates/configuration/testing/snapshots/0001-big-network.toml
index 53d22ca74..ef7c15163 100644
--- a/crates/configuration/testing/snapshots/0001-big-network.toml
+++ b/crates/configuration/testing/snapshots/0001-big-network.toml
@@ -38,6 +38,7 @@ balance = 2000000000000
 default_db_snapshot = "https://storage.com/path/to/db_snapshot.tgz"
 chain_spec_path = "/path/to/my/chain/spec.json"
 cumulus_based = true
+evm_based = false
 
 [[parachains.collators]]
 name = "john"
@@ -67,6 +68,7 @@ add_to_genesis = true
 balance = 2000000000000
 chain_spec_path = "/path/to/my/other/chain/spec.json"
 cumulus_based = true
+evm_based = false
 
 [[parachains.collators]]
 name = "mike"
diff --git a/crates/configuration/testing/snapshots/0002-overridden-defaults.toml b/crates/configuration/testing/snapshots/0002-overridden-defaults.toml
index 31698eeb1..1abe3bba1 100644
--- a/crates/configuration/testing/snapshots/0002-overridden-defaults.toml
+++ b/crates/configuration/testing/snapshots/0002-overridden-defaults.toml
@@ -56,6 +56,7 @@ default_image = "mydefaultimage:latest"
 default_db_snapshot = "https://storage.com/path/to/other_snapshot.tgz"
 chain_spec_path = "/path/to/my/chain/spec.json"
 cumulus_based = true
+evm_based = false
 
 [[parachains.collators]]
 name = "john"
diff --git a/crates/configuration/testing/snapshots/0003-small-network_w_parachain.toml b/crates/configuration/testing/snapshots/0003-small-network_w_parachain.toml
index a330ec7c8..8a1ba79fa 100644
--- a/crates/configuration/testing/snapshots/0003-small-network_w_parachain.toml
+++ b/crates/configuration/testing/snapshots/0003-small-network_w_parachain.toml
@@ -38,3 +38,20 @@ validator = true
 invulnerable = true
 bootnode = true
 balance = 5000000000
+
+[[parachains]]
+id = 1000
+chain = "myparachain"
+onboard_as_parachain = false
+balance = 2000000000000
+default_db_snapshot = "https://storage.com/path/to/db_snapshot.tgz"
+chain_spec_path = "/path/to/my/chain/spec.json"
+cumulus_based = true
+evm_based = true
+
+[[parachains.collators]]
+name = "john"
+validator = true
+invulnerable = true
+bootnode = true
+balance = 5000000000
diff --git a/crates/orchestrator/Cargo.toml b/crates/orchestrator/Cargo.toml
index 12a509daa..86ac12bf2 100644
--- 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
index b3548a133..2db21ab4f 100644
--- 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),
@@ -436,6 +449,12 @@ impl ChainSpec {
 
             clear_authorities(&pointer, &mut chain_spec_json);
 
+            let key_type_to_use = if para.is_evm_based {
+                SessionKeyType::Evm
+            } else {
+                SessionKeyType::Default
+            };
+
             // Get validators to add as authorities
             let validators: Vec<&NodeSpec> = para
                 .collators
@@ -448,7 +467,7 @@ 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, key_type_to_use);
             } else if chain_spec_json
                 .pointer(&format!("{}/aura", pointer))
                 .is_some()
@@ -466,7 +485,12 @@ impl ChainSpec {
                 .filter(|node| node.is_invulnerable)
                 .collect();
 
-            add_collator_selection(&pointer, &mut chain_spec_json, &invulnerables);
+            add_collator_selection(
+                &pointer,
+                &mut chain_spec_json,
+                &invulnerables,
+                key_type_to_use,
+            );
 
             // override `parachainInfo/parachainId`
             override_parachain_info(&pointer, &mut chain_spec_json, para.id);
@@ -547,7 +571,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 +909,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 +940,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 = if use_stash { sr_stash } else { sr_account };
-    (
-        account_to_use.address.clone(),
-        account_to_use.address.clone(),
-        keys,
-    )
+    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,
     nodes: &[&NodeSpec],
-    use_stash: bool,
+    session_key: SessionKeyType,
 ) {
     let asset_hub_polkadot = chain_spec_json
         .get("id")
@@ -928,7 +964,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 {
@@ -1014,14 +1050,20 @@ fn add_collator_selection(
     runtime_config_ptr: &str,
     chain_spec_json: &mut serde_json::Value,
     nodes: &[&NodeSpec],
+    session_key: SessionKeyType,
 ) {
     if let Some(val) = chain_spec_json.pointer_mut(runtime_config_ptr) {
+        let key_type = if let SessionKeyType::Evm = session_key {
+            "eth"
+        } else {
+            "sr"
+        };
         let keys: Vec<String> = nodes
             .iter()
             .map(|node| {
                 node.accounts
                     .accounts
-                    .get("sr")
+                    .get(key_type)
                     .expect(&format!(
                         "'sr' account should be set at spec computation {THIS_IS_A_BUG}"
                     ))
@@ -1255,17 +1297,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 +1328,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
index c7e758207..8fc4e747c 100644
--- 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/parachain.rs b/crates/orchestrator/src/network_spec/parachain.rs
index 4bea46b7f..ef4240e5c 100644
--- 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,