From 76bdf23523774b64c830752cf94f8cecccd8675d Mon Sep 17 00:00:00 2001
From: Dustin Brickwood <dustinbrickwood204@gmail.com>
Date: Mon, 25 Nov 2024 18:51:47 -0600
Subject: [PATCH] feat: adds `--auto-impersonate` cli cmd (#430)

* feat: adds auto-impersonate cli cmd

* chore: add to config setup
---
 src/config/cli.rs     | 10 ++++++++++
 src/config/mod.rs     | 10 ++++++++++
 src/main.rs           |  5 +----
 src/node/in_memory.rs | 15 +++++++++++++--
 4 files changed, 34 insertions(+), 6 deletions(-)

diff --git a/src/config/cli.rs b/src/config/cli.rs
index f43d6069..384941e8 100644
--- a/src/config/cli.rs
+++ b/src/config/cli.rs
@@ -193,6 +193,15 @@ pub struct Cli {
     /// [default: m/44'/60'/0'/0/]
     #[arg(long, help_heading = "Account Configuration")]
     pub derivation_path: Option<String>,
+
+    /// Enables automatic impersonation on startup. This allows any transaction sender to be
+    /// simulated as different accounts, which is useful for testing contract behavior.
+    #[arg(
+        long,
+        visible_alias = "auto-unlock",
+        help_heading = "Account Configuration"
+    )]
+    pub auto_impersonate: bool,
 }
 
 #[derive(Debug, Subcommand, Clone)]
@@ -295,6 +304,7 @@ impl Cli {
             .with_log_level(self.log)
             .with_log_file_path(self.log_file_path.clone())
             .with_account_generator(self.account_generator())
+            .with_auto_impersonate(self.auto_impersonate)
             .with_genesis_balance(genesis_balance)
             .with_cache_config(self.cache.map(|cache_type| {
                 match cache_type {
diff --git a/src/config/mod.rs b/src/config/mod.rs
index 97e97bc5..3d15a637 100644
--- a/src/config/mod.rs
+++ b/src/config/mod.rs
@@ -99,6 +99,8 @@ pub struct TestNodeConfig {
     pub account_generator: Option<AccountGenerator>,
     /// Signer accounts that can sign messages/transactions
     pub signer_accounts: Vec<PrivateKeySigner>,
+    /// Enable auto impersonation of accounts on startup
+    pub enable_auto_impersonate: bool,
     /// Whether the node operates in offline mode
     pub offline: bool,
     /// The host the server will listen on
@@ -144,6 +146,7 @@ impl Default for TestNodeConfig {
             account_generator: None,
             genesis_accounts: genesis_accounts.clone(),
             signer_accounts: genesis_accounts,
+            enable_auto_impersonate: false,
             // 100ETH default balance
             genesis_balance: U256::from(100u128 * 10u128.pow(18)),
 
@@ -628,6 +631,13 @@ impl TestNodeConfig {
             .with_genesis_accounts(accounts)
     }
 
+    /// Sets whether to enable autoImpersonate
+    #[must_use]
+    pub fn with_auto_impersonate(mut self, enable_auto_impersonate: bool) -> Self {
+        self.enable_auto_impersonate = enable_auto_impersonate;
+        self
+    }
+
     /// Set the offline mode
     #[must_use]
     pub fn with_offline(mut self, offline: Option<bool>) -> Self {
diff --git a/src/main.rs b/src/main.rs
index 7840545d..186ab7a3 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -124,10 +124,7 @@ async fn main() -> anyhow::Result<()> {
     let fork_details = match command {
         Command::Run => {
             if config.offline {
-                tracing::warn!(
-                    "Running in offline mode: default fee parameters will be used. \
-        To override, specify values in `config.toml` and use the `--config` flag."
-                );
+                tracing::warn!("Running in offline mode: default fee parameters will be used.");
                 None
             } else {
                 // Initialize the client to get the fee params
diff --git a/src/node/in_memory.rs b/src/node/in_memory.rs
index 4c5952d4..80ce0896 100644
--- a/src/node/in_memory.rs
+++ b/src/node/in_memory.rs
@@ -263,6 +263,11 @@ impl<S: std::fmt::Debug + ForkSource> InMemoryNodeInner<S> {
                     f.estimate_gas_scale_factor,
                 )
             };
+            let impersonation_manager = ImpersonationManager::default();
+            if config.enable_auto_impersonate {
+                // Enable auto impersonation if configured
+                impersonation_manager.set_auto_impersonation(true);
+            }
 
             InMemoryNodeInner {
                 time: TimestampManager::new(f.block_timestamp),
@@ -286,7 +291,7 @@ impl<S: std::fmt::Debug + ForkSource> InMemoryNodeInner<S> {
                     &updated_config.system_contracts_options,
                     updated_config.use_evm_emulator,
                 ),
-                impersonation: Default::default(),
+                impersonation: impersonation_manager,
                 rich_accounts: HashSet::new(),
                 previous_states: Default::default(),
                 observability,
@@ -299,6 +304,12 @@ impl<S: std::fmt::Debug + ForkSource> InMemoryNodeInner<S> {
             blocks.insert(block_hash, create_genesis(NON_FORK_FIRST_BLOCK_TIMESTAMP));
             let fee_input_provider = TestNodeFeeInputProvider::default();
 
+            let impersonation_manager = ImpersonationManager::default();
+            if config.enable_auto_impersonate {
+                // Enable auto impersonation if configured
+                impersonation_manager.set_auto_impersonation(true);
+            }
+
             InMemoryNodeInner {
                 time: TimestampManager::new(NON_FORK_FIRST_BLOCK_TIMESTAMP),
                 current_batch: 0,
@@ -321,7 +332,7 @@ impl<S: std::fmt::Debug + ForkSource> InMemoryNodeInner<S> {
                     &config.system_contracts_options,
                     config.use_evm_emulator,
                 ),
-                impersonation: Default::default(),
+                impersonation: impersonation_manager,
                 rich_accounts: HashSet::new(),
                 previous_states: Default::default(),
                 observability,