Support different callers

crates/node/src/common.rs

@@ -1,9 +1,14 @@
+use ::core::pin::Pin;
+
 use alloy::{
-    network::{Network, TransactionBuilder},
+    consensus::SignableTransaction,
+    network::{Network, TransactionBuilder, TxSigner},
+    primitives::Address,
     providers::{
         Provider, SendableTx,
         fillers::{GasFiller, TxFiller},
     },
+    signers::{Signature, local::PrivateKeySigner},
     transports::TransportResult,
 };
 
@@ -76,3 +81,28 @@ where
         }
     }
 }
+
+/// This is a signer that is able to sign transactions for a specific address with another private
+/// key.
+pub struct AddressSigner {
+    pub private_key: PrivateKeySigner,
+    pub address: Address,
+}
+
+impl TxSigner<Signature> for AddressSigner {
+    fn address(&self) -> Address {
+        self.address
+    }
+
+    fn sign_transaction<'a, 'b, 'c>(
+        &'a self,
+        tx: &'b mut dyn SignableTransaction<Signature>,
+    ) -> Pin<Box<dyn Future<Output = Result<Signature, alloy::signers::Error>> + Send + 'c>>
+    where
+        'a: 'c,
+        'b: 'c,
+        Self: 'c,
+    {
+        <PrivateKeySigner as TxSigner<Signature>>::sign_transaction(&self.private_key, tx)
+    }
+}

crates/node/src/geth.rs

@@ -11,6 +11,7 @@ use std::{
 
 use alloy::{
     eips::BlockNumberOrTag,
+    genesis::Genesis,
     network::{Ethereum, EthereumWallet},
     primitives::{Address, BlockHash, BlockNumber, BlockTimestamp, U256},
     providers::{
@@ -22,12 +23,16 @@ use alloy::{
         TransactionReceipt, TransactionRequest,
         trace::geth::{DiffMode, GethDebugTracingOptions, PreStateConfig, PreStateFrame},
     },
+    signers::local::PrivateKeySigner,
 };
 use revive_dt_config::Arguments;
 use revive_dt_node_interaction::{BlockingExecutor, EthereumNode};
 use tracing::Level;
 
-use crate::{Node, common::FallbackGasFiller};
+use crate::{
+    Node,
+    common::{AddressSigner, FallbackGasFiller},
+};
 
 static NODE_COUNT: AtomicU32 = AtomicU32::new(0);
 
@@ -50,7 +55,10 @@ pub struct Instance {
     network_id: u64,
     start_timeout: u64,
     wallet: EthereumWallet,
+    private_key: PrivateKeySigner,
     nonce_manager: CachedNonceManager,
+    additional_callers: Vec<Address>,
+
     /// This vector stores [`File`] objects that we use for logging which we want to flush when the
     /// node object is dropped. We do not store them in a structured fashion at the moment (in
     /// separate fields) as the logic that we need to apply to them is all the same regardless of
@@ -78,8 +86,24 @@ impl Instance {
         create_dir_all(&self.base_directory)?;
         create_dir_all(&self.logs_directory)?;
 
+        // Modifying the genesis file and adding all of the additional callers as contract accounts.
+        let mut genesis = serde_json::from_str::<Genesis>(&genesis)?;
+        for additional_caller in self.additional_callers.iter() {
+            let account = genesis.alloc.entry(*additional_caller).or_default();
+            account.private_key = Some(self.private_key.to_bytes());
+            *account = account
+                .clone()
+                .with_balance("1000000000000000000".parse().expect("Can't fail"));
+        }
         let genesis_path = self.base_directory.join(Self::GENESIS_JSON_FILE);
-        File::create(&genesis_path)?.write_all(genesis.as_bytes())?;
+        serde_json::to_writer(File::create(&genesis_path)?, &genesis)?;
+
+        for additional_caller in self.additional_callers.iter() {
+            self.wallet.register_signer(AddressSigner {
+                private_key: self.private_key.clone(),
+                address: *additional_caller,
+            });
+        }
 
         let mut child = Command::new(&self.geth)
             .arg("init")
@@ -424,7 +448,7 @@ impl EthereumNode for Instance {
 }
 
 impl Node for Instance {
-    fn new(config: &Arguments) -> Self {
+    fn new(config: &Arguments, additional_callers: &[Address]) -> Self {
         let geth_directory = config.directory().join(Self::BASE_DIRECTORY);
         let id = NODE_COUNT.fetch_add(1, Ordering::SeqCst);
         let base_directory = geth_directory.join(id.to_string());
@@ -444,6 +468,8 @@ impl Node for Instance {
             // the vector. It's the stdout and stderr of the geth node.
             logs_file_to_flush: Vec::with_capacity(2),
             nonce_manager: Default::default(),
+            additional_callers: additional_callers.to_vec(),
+            private_key: config.signer(),
         }
     }
 
@@ -520,7 +546,7 @@ mod tests {
 
     fn new_node() -> (Instance, TempDir) {
         let (args, temp_dir) = test_config();
-        let mut node = Instance::new(&args);
+        let mut node = Instance::new(&args, &[]);
         node.init(GENESIS_JSON.to_owned())
             .expect("Failed to initialize the node")
             .spawn_process()
@@ -530,21 +556,21 @@ mod tests {
 
     #[test]
     fn init_works() {
-        Instance::new(&test_config().0)
+        Instance::new(&test_config().0, &[])
             .init(GENESIS_JSON.to_string())
             .unwrap();
     }
 
     #[test]
     fn spawn_works() {
-        Instance::new(&test_config().0)
+        Instance::new(&test_config().0, &[])
             .spawn(GENESIS_JSON.to_string())
             .unwrap();
     }
 
     #[test]
     fn version_works() {
-        let version = Instance::new(&test_config().0).version().unwrap();
+        let version = Instance::new(&test_config().0, &[]).version().unwrap();
         assert!(
             version.starts_with("geth version"),
             "expected version string, got: '{version}'"

crates/node/src/kitchensink.rs

@@ -507,7 +507,7 @@ impl EthereumNode for KitchensinkNode {
 }
 
 impl Node for KitchensinkNode {
-    fn new(config: &Arguments) -> Self {
+    fn new(config: &Arguments, _additional_callers: &[Address]) -> Self {
         let kitchensink_directory = config.directory().join(Self::BASE_DIRECTORY);
         let id = NODE_COUNT.fetch_add(1, Ordering::SeqCst);
         let base_directory = kitchensink_directory.join(id.to_string());
@@ -814,6 +814,12 @@ impl TransactionBuilder<KitchenSinkNetwork> for <Ethereum as Network>::Transacti
     > {
         Ok(wallet.sign_request(self).await?)
     }
+
+    fn take_nonce(&mut self) -> Option<u64> {
+        <<Ethereum as Network>::TransactionRequest as TransactionBuilder<Ethereum>>::take_nonce(
+            self,
+        )
+    }
 }
 
 #[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
@@ -1054,7 +1060,7 @@ mod tests {
         let _guard = NODE_START_MUTEX.lock().unwrap();
 
         let (args, temp_dir) = test_config();
-        let mut node = KitchensinkNode::new(&args);
+        let mut node = KitchensinkNode::new(&args, &[]);
         node.init(GENESIS_JSON)
             .expect("Failed to initialize the node")
             .spawn_process()
@@ -1109,7 +1115,7 @@ mod tests {
         }
         "#;
 
-        let mut dummy_node = KitchensinkNode::new(&test_config().0);
+        let mut dummy_node = KitchensinkNode::new(&test_config().0, &[]);
 
         // Call `init()`
         dummy_node.init(genesis_content).expect("init failed");
@@ -1153,7 +1159,7 @@ mod tests {
         }
         "#;
 
-        let node = KitchensinkNode::new(&test_config().0);
+        let node = KitchensinkNode::new(&test_config().0, &[]);
 
         let result = node
             .extract_balance_from_genesis_file(genesis_json)
@@ -1226,7 +1232,7 @@ mod tests {
     fn spawn_works() {
         let (config, _temp_dir) = test_config();
 
-        let mut node = KitchensinkNode::new(&config);
+        let mut node = KitchensinkNode::new(&config, &[]);
         node.spawn(GENESIS_JSON.to_string()).unwrap();
     }
 
@@ -1234,7 +1240,7 @@ mod tests {
     fn version_works() {
         let (config, _temp_dir) = test_config();
 
-        let node = KitchensinkNode::new(&config);
+        let node = KitchensinkNode::new(&config, &[]);
         let version = node.version().unwrap();
 
         assert!(
@@ -1247,7 +1253,7 @@ mod tests {
     fn eth_rpc_version_works() {
         let (config, _temp_dir) = test_config();
 
-        let node = KitchensinkNode::new(&config);
+        let node = KitchensinkNode::new(&config, &[]);
         let version = node.eth_rpc_version().unwrap();
 
         assert!(

crates/node/src/lib.rs

@@ -1,5 +1,6 @@
 //! This crate implements the testing nodes.
 
+use alloy::primitives::Address;
 use revive_dt_config::Arguments;
 use revive_dt_node_interaction::EthereumNode;
 
@@ -14,7 +15,7 @@ pub const GENESIS_JSON: &str = include_str!("../../../genesis.json");
 /// An abstract interface for testing nodes.
 pub trait Node: EthereumNode {
     /// Create a new uninitialized instance.
-    fn new(config: &Arguments) -> Self;
+    fn new(config: &Arguments, additional_callers: &[Address]) -> Self;
 
     /// Spawns a node configured according to the genesis json.
     ///

crates/node/src/pool.rs

@@ -6,6 +6,7 @@ use std::{
     thread,
 };
 
+use alloy::primitives::Address;
 use anyhow::Context;
 use revive_dt_config::Arguments;
 
@@ -23,7 +24,7 @@ where
     T: Node + Send + 'static,
 {
     /// Create a new Pool. This will start as many nodes as there are workers in `config`.
-    pub fn new(config: &Arguments) -> anyhow::Result<Self> {
+    pub fn new(config: &Arguments, additional_callers: &[Address]) -> anyhow::Result<Self> {
         let nodes = config.workers;
         let genesis = read_to_string(&config.genesis_file).context(format!(
             "can not read genesis file: {}",
@@ -34,7 +35,10 @@ where
         for _ in 0..nodes {
             let config = config.clone();
             let genesis = genesis.clone();
-            handles.push(thread::spawn(move || spawn_node::<T>(&config, genesis)));
+            let additional_callers = additional_callers.to_vec();
+            handles.push(thread::spawn(move || {
+                spawn_node::<T>(&config, genesis, &additional_callers)
+            }));
         }
 
         let mut nodes = Vec::with_capacity(nodes);
@@ -60,8 +64,12 @@ where
     }
 }
 
-fn spawn_node<T: Node + Send>(args: &Arguments, genesis: String) -> anyhow::Result<T> {
-    let mut node = T::new(args);
+fn spawn_node<T: Node + Send>(
+    args: &Arguments,
+    genesis: String,
+    additional_callers: &[Address],
+) -> anyhow::Result<T> {
+    let mut node = T::new(args, additional_callers);
     tracing::info!("starting node: {}", node.connection_string());
     node.spawn(genesis)?;
     Ok(node)