[WIP] Support substrate node in zombienet network

crates/node/Cargo.toml

@@ -32,6 +32,7 @@ zombienet-sdk = { workspace = true }
 [dev-dependencies]
 temp-dir = { workspace = true }
 tokio = { workspace = true }
+tracing-subscriber = { workspace = true }
 
 [lints]
 workspace = true

crates/node/src/substrate.rs

@@ -466,8 +466,8 @@ impl EthereumNode for SubstrateNode {
 }
 
 pub struct SubstrateNodeResolver<F: TxFiller<ReviveNetwork>, P: Provider<ReviveNetwork>> {
-    id: u32,
-    provider: FillProvider<F, P, ReviveNetwork>,
+    pub(crate) id: u32,
+    pub(crate) provider: FillProvider<F, P, ReviveNetwork>,
 }
 
 impl<F: TxFiller<ReviveNetwork>, P: Provider<ReviveNetwork>> ResolverApi

crates/node/src/zombie.rs

@@ -1,3 +1,4 @@
+use core::net;
 use std::{
     fs::{File, create_dir_all},
     path::PathBuf,
@@ -10,12 +11,15 @@ use std::{
 };
 
 use alloy::{
-    genesis::Genesis,
+    consensus::{BlockHeader, TxEnvelope},
+    eips::BlockNumberOrTag,
+    genesis::{Genesis, GenesisAccount},
     network::{
-        EthereumWallet, TransactionBuilder,
+        self, Ethereum, EthereumWallet, Network, NetworkWallet, TransactionBuilder,
+        TransactionBuilderError, UnbuiltTransactionError,
     },
     primitives::{
-        Address, StorageKey,
+        Address, B64, B256, BlockHash, BlockNumber, BlockTimestamp, Bloom, Bytes, StorageKey,
         TxHash, U256,
     },
     providers::{
@@ -25,22 +29,27 @@ use alloy::{
     },
     rpc::types::{
         EIP1186AccountProofResponse, TransactionReceipt,
-        trace::geth::{DiffMode, GethDebugTracingOptions},
+        eth::{Block, Header, Transaction},
+        trace::geth::{DiffMode, GethDebugTracingOptions, PreStateConfig, PreStateFrame},
     },
 };
 use anyhow::Context as _;
 use revive_common::EVMVersion;
 use revive_dt_common::fs::clear_directory;
 use revive_dt_format::traits::ResolverApi;
+use serde_json::{Value as JsonValue, json};
+use sp_core::crypto::Ss58Codec;
+use sp_runtime::AccountId32;
 
 use revive_dt_config::*;
 use revive_dt_node_interaction::EthereumNode;
-use zombienet_sdk::{
-    LocalFileSystem, NetworkConfigBuilder, NetworkConfigExt,
-};
+use tracing::info;
+use zombienet_sdk::{LocalFileSystem, NetworkConfigBuilder, NetworkConfigExt, subxt};
 
 use crate::{
-    common::FallbackGasFiller, substrate::ReviveNetwork,
+    common::FallbackGasFiller,
+    constants::INITIAL_BALANCE,
+    substrate::{ReviveNetwork, SubstrateNodeResolver},
 };
 
 static NODE_COUNT: AtomicU32 = AtomicU32::new(0);
@@ -75,7 +84,10 @@ impl ZombieNode {
     const ZOMBIENET_STDOUT_LOG_FILE_NAME: &str = "node_stdout.log";
     const ZOMBIENET_STDERR_LOG_FILE_NAME: &str = "node_stderr.log";
 
+    pub const SUBSTRATE_EXPORT_CHAINSPEC_COMMAND: &str = "export-chain-spec";
+
     pub fn new(
+        node_path: PathBuf,
         context: impl AsRef<WorkingDirectoryConfiguration>
         + AsRef<EthRpcConfiguration>
         + AsRef<WalletConfiguration>,
@@ -95,11 +107,12 @@ impl ZombieNode {
             logs_directory,
             wallet,
             logs_file_to_flush: Vec::with_capacity(2),
+            node_binary: node_path,
             ..Default::default()
         }
     }
 
-    fn init(&mut self, genesis: Genesis) -> anyhow::Result<&mut Self> {
+    fn init(&mut self, mut genesis: Genesis) -> anyhow::Result<&mut Self> {
         let _ = clear_directory(&self.base_directory);
         let _ = clear_directory(&self.logs_directory);
 
@@ -108,24 +121,25 @@ impl ZombieNode {
         create_dir_all(&self.logs_directory)
             .context("Failed to create logs directory for zombie node")?;
 
-        let _genesis = serde_json::to_value(genesis)?; // TODO: validate this
+        let template_chainspec_path = self.base_directory.join(Self::CHAIN_SPEC_JSON_FILE);
+        self.prepare_chainspec(template_chainspec_path.clone(), genesis)?;
+
+        let node_binary = self.node_binary.to_str().unwrap_or_default();
 
         let network_config = NetworkConfigBuilder::new()
             .with_relaychain(|r| {
                 r.with_chain("rococo-local")
                     .with_default_command("polkadot")
-                    //.with_genesis_overrides(zombie_genesis)
-                    //.with_chain_spec_path(self.base_directory.join(Self::CHAIN_SPEC_JSON_FILE))
                     .with_node(|node| node.with_name("alice"))
             })
             .with_global_settings(|g| g.with_base_dir(&self.base_directory))
             .with_parachain(|p| {
                 p.with_id(Self::PARACHAIN_ID)
-                    .evm_based(true)
+                    .with_chain_spec_path(template_chainspec_path.to_str().unwrap())
                     .with_collator(|n| {
                         n.with_name("collator")
-                            .with_command("polkadot-parachain")
-                            //.with_args(vec!["--rpc-methods=Unsafe".into(), "--rpc-cors=all".into()])
+                            .with_command(node_binary)
+                            .with_args(vec!["--dev".into()])
                             .with_ws_port(Self::BASE_RPC_PORT + self.id as u16)
                     })
             })
@@ -137,6 +151,104 @@ impl ZombieNode {
         Ok(self)
     }
 
+    fn prepare_chainspec(
+        &mut self,
+        template_chainspec_path: PathBuf,
+        mut genesis: Genesis,
+    ) -> anyhow::Result<()> {
+        let output = std::process::Command::new(&self.node_binary)
+            .arg(Self::SUBSTRATE_EXPORT_CHAINSPEC_COMMAND)
+            .arg("--chain")
+            .arg("dev")
+            .output()
+            .context("Failed to export the chain-spec")?;
+
+        if !output.status.success() {
+            anyhow::bail!(
+                "Substrate-node export-chain-spec failed: {}",
+                String::from_utf8_lossy(&output.stderr)
+            );
+        }
+
+        let content = String::from_utf8(output.stdout)
+            .context("Failed to decode Substrate export-chain-spec output as UTF-8")?;
+        let mut chainspec_json: JsonValue =
+            serde_json::from_str(&content).context("Failed to parse Substrate chain spec JSON")?;
+
+        let existing_chainspec_balances =
+            chainspec_json["genesis"]["runtimeGenesis"]["patch"]["balances"]["balances"]
+                .as_array()
+                .cloned()
+                .unwrap_or_default();
+
+        let mut merged_balances: Vec<(String, u128)> = existing_chainspec_balances
+            .into_iter()
+            .filter_map(|val| {
+                if let Some(arr) = val.as_array() {
+                    if arr.len() == 2 {
+                        let account = arr[0].as_str()?.to_string();
+                        let balance = arr[1].as_f64()? as u128;
+                        return Some((account, balance));
+                    }
+                }
+                None
+            })
+            .collect();
+
+        let mut eth_balances = {
+            for signer_address in
+                <EthereumWallet as NetworkWallet<Ethereum>>::signer_addresses(&self.wallet)
+            {
+                // Note, the use of the entry API here means that we only modify the entries for any
+                // account that is not in the `alloc` field of the genesis state.
+                genesis
+                    .alloc
+                    .entry(signer_address)
+                    .or_insert(GenesisAccount::default().with_balance(U256::from(INITIAL_BALANCE)));
+            }
+            self.extract_balance_from_genesis_file(&genesis)
+                .context("Failed to extract balances from EVM genesis JSON")?
+        };
+
+        merged_balances.append(&mut eth_balances);
+
+        chainspec_json["genesis"]["runtimeGenesis"]["patch"]["balances"]["balances"] =
+            json!(merged_balances);
+
+        let writer = std::fs::File::create(&template_chainspec_path)
+            .context("Failed to create substrate template chainspec file")?;
+
+        serde_json::to_writer_pretty(writer, &chainspec_json)
+            .context("Failed to write substrate template chainspec JSON")?;
+
+        Ok(())
+    }
+
+    fn extract_balance_from_genesis_file(
+        &self,
+        genesis: &Genesis,
+    ) -> anyhow::Result<Vec<(String, u128)>> {
+        genesis
+            .alloc
+            .iter()
+            .try_fold(Vec::new(), |mut vec, (address, acc)| {
+                let substrate_address = Self::eth_to_substrate_address(address);
+                let balance = acc.balance.try_into()?;
+                vec.push((substrate_address, balance));
+                Ok(vec)
+            })
+    }
+
+    fn eth_to_substrate_address(address: &Address) -> String {
+        let eth_bytes = address.0.0;
+
+        let mut padded = [0xEEu8; 32];
+        padded[..20].copy_from_slice(&eth_bytes);
+
+        let account_id = AccountId32::from(padded);
+        account_id.to_ss58check()
+    }
+
     async fn spawn_process(&mut self) -> anyhow::Result<&mut Self> {
         let Some(network_config) = self.network_config.clone() else {
             anyhow::bail!("Node not initialized, call init() first");
@@ -145,11 +257,9 @@ impl ZombieNode {
         let network = network_config
             .spawn_native()
             .await
-            .context("Failed to spawn zombienet network")?;
-        network
-            .wait_until_is_up(60)
-            .await
-            .context("Network failed to start within timeout")?;
+            .map_err(|e| anyhow::anyhow!("Failed to spawn zombienet network: {e:?}"))?;
+
+        tracing::info!("Zombienet network is up");
 
         let ws_uri = network
             .parachain(Self::PARACHAIN_ID)
@@ -239,14 +349,35 @@ impl EthereumNode for ZombieNode {
         &self,
         tx_hash: TxHash,
     ) -> Pin<Box<dyn Future<Output = anyhow::Result<DiffMode>> + '_>> {
-        todo!()
+        Box::pin(async move {
+            let trace_options = GethDebugTracingOptions::prestate_tracer(PreStateConfig {
+                diff_mode: Some(true),
+                disable_code: None,
+                disable_storage: None,
+            });
+            match self
+                .trace_transaction(tx_hash, trace_options)
+                .await?
+                .try_into_pre_state_frame()?
+            {
+                PreStateFrame::Diff(diff) => Ok(diff),
+                _ => anyhow::bail!("expected a diff mode trace"),
+            }
+        })
     }
 
     fn balance_of(
         &self,
         address: Address,
     ) -> Pin<Box<dyn Future<Output = anyhow::Result<U256>> + '_>> {
-        todo!()
+        Box::pin(async move {
+            self.provider()
+                .await
+                .context("Failed to get the zombie provider")?
+                .get_balance(address)
+                .await
+                .map_err(Into::into)
+        })
     }
 
     fn latest_state_proof(
@@ -254,27 +385,37 @@ impl EthereumNode for ZombieNode {
         address: Address,
         keys: Vec<StorageKey>,
     ) -> Pin<Box<dyn Future<Output = anyhow::Result<EIP1186AccountProofResponse>> + '_>> {
-        todo!()
+        Box::pin(async move {
+            self.provider()
+                .await
+                .context("Failed to get the zombie provider")?
+                .get_proof(address, keys)
+                .latest()
+                .await
+                .map_err(Into::into)
+        })
     }
 
     fn resolver(
         &self,
     ) -> Pin<Box<dyn Future<Output = anyhow::Result<Arc<dyn ResolverApi + '_>>> + '_>> {
-        todo!()
+        Box::pin(async move {
+            let id = self.id;
+            let provider = self.provider().await?;
+            Ok(Arc::new(SubstrateNodeResolver { id, provider }) as Arc<dyn ResolverApi>)
+        })
     }
 
     fn evm_version(&self) -> EVMVersion {
-        todo!()
+        EVMVersion::Cancun
     }
 }
 
 #[cfg(test)]
 mod tests {
     use alloy::rpc::types::TransactionRequest;
-    use std::sync::{LazyLock, Mutex};
-
-    use std::fs;
-
+    use std::{sync,fs};
+    
     use super::*;
     use crate::Node;
 
@@ -286,7 +427,7 @@ mod tests {
 
     fn new_node() -> (TestExecutionContext, ZombieNode) {
         let context = test_config();
-        let mut node = ZombieNode::new(&context);
+        let mut node = ZombieNode::new(context.zombienet_configuration.path.clone(), &context);
         let genesis = context.genesis_configuration.genesis().unwrap().clone();
         node.init(genesis).unwrap();
         (context, node)
@@ -297,7 +438,7 @@ mod tests {
         let context = test_config();
         let mut ids = Vec::new();
         for _ in 0..5 {
-            let node = ZombieNode::new(&context);
+            let (_, node) = new_node();
             ids.push(node.id);
         }
         // Check uniqueness
@@ -325,6 +466,16 @@ mod tests {
 
     #[tokio::test]
     async fn test_transfer_transaction_should_return_receipt() {
+        use tracing_subscriber::filter::LevelFilter;
+        use tracing_subscriber::{EnvFilter, FmtSubscriber};
+        tracing_subscriber::fmt()
+            .with_env_filter(
+                EnvFilter::builder()
+                    .with_default_directive(LevelFilter::DEBUG.into())
+                    .from_env_lossy(),
+            )
+            .init();
+
         let (context, mut node) = new_node();
 
         let node = node.spawn_process().await.unwrap();
@@ -338,9 +489,8 @@ mod tests {
         let transaction = TransactionRequest::default()
             .to(account_address)
             .value(U256::from(100_000_000_000_000u128));
-        let receipt = provider.send_transaction(transaction).await;
-        tracing::info!("Sending transaction");
 
+        let receipt = provider.send_transaction(transaction).await;
         let _ = receipt
             .expect("Failed to send the transfer transaction")
             .get_receipt()