From 4a58732b8aff01f0aaa680724f127f5f411e9870 Mon Sep 17 00:00:00 2001 From: grw-ms Date: Thu, 21 Sep 2023 02:42:10 -0700 Subject: [PATCH] fix: make impersonated_accounts a hashset instead of option --- src/hardhat.rs | 38 ++++++++++++++++++++++++++------------ src/node.rs | 26 ++++++++++++++------------ 2 files changed, 40 insertions(+), 24 deletions(-) diff --git a/src/hardhat.rs b/src/hardhat.rs index c0a40133..31623022 100644 --- a/src/hardhat.rs +++ b/src/hardhat.rs @@ -73,7 +73,10 @@ pub trait HardhatNamespaceT { ) -> BoxFuture>; /// Hardhat Network allows you to send transactions impersonating specific account and contract addresses. - /// To impersonate an account use this method, passing the address to impersonate as its parameter: + /// To impersonate an account use this method, passing the address to impersonate as its parameter. + /// After calling this method, any transactions with this sender will be executed without verification. + /// Multiple addresses can be impersonated at once. + /// /// # Arguments /// /// * `address` - The address to impersonate @@ -85,17 +88,17 @@ pub trait HardhatNamespaceT { fn impersonate_account(&self, address: Address) -> BoxFuture>; /// Use this method to stop impersonating an account after having previously used `hardhat_impersonateAccount` - /// The address parameter is included for compatibility- since we only support impersonating one account at a time, it is ignored. + /// The method returns `true` if the account was being impersonated and `false` otherwise. /// /// # Arguments /// - /// * `address` (Optional) - The address to stop impersonating. + /// * `address` - The address to stop impersonating. /// /// # Returns /// /// A `BoxFuture` containing a `Result` with a `bool` representing the success of the operation. #[rpc(name = "hardhat_stopImpersonatingAccount")] - fn stop_impersonating_account(&self, address: Option
) -> BoxFuture>; + fn stop_impersonating_account(&self, address: Address) -> BoxFuture>; } impl HardhatNamespaceT @@ -205,23 +208,34 @@ impl HardhatNamespaceT Box::pin(async move { match inner.write() { Ok(mut inner) => { - inner.set_impersonated_account(address); - log::info!("🕵️ Account {:?} has been impersonated", address); - Ok(true) + if inner.set_impersonated_account(address) { + log::info!("🕵️ Account {:?} has been impersonated", address); + Ok(true) + } else { + log::info!("🕵️ Account {:?} was already impersonated", address); + Ok(false) + } } Err(_) => Err(into_jsrpc_error(Web3Error::InternalError)), } }) } - fn stop_impersonating_account(&self, _address: Option
) -> BoxFuture> { + fn stop_impersonating_account(&self, address: Address) -> BoxFuture> { let inner = Arc::clone(&self.node); Box::pin(async move { match inner.write() { Ok(mut inner) => { - inner.stop_impersonating_account(); - log::info!("🕵️ Stopped impersonating account"); - Ok(true) + if inner.stop_impersonating_account(address) { + log::info!("🕵️ Stopped impersonating account {:?}", address); + Ok(true) + } else { + log::info!( + "🕵️ Account {:?} was not impersonated, nothing to stop", + address + ); + Ok(false) + } } Err(_) => Err(into_jsrpc_error(Web3Error::InternalError)), } @@ -394,7 +408,7 @@ mod tests { // stop impersonating the account let result = hardhat - .stop_impersonating_account(None) + .stop_impersonating_account(to_impersonate) .await .expect("stop_impersonating_account"); assert!(result); diff --git a/src/node.rs b/src/node.rs index 3c5dad40..f87fc1d2 100644 --- a/src/node.rs +++ b/src/node.rs @@ -16,7 +16,7 @@ use futures::FutureExt; use jsonrpc_core::BoxFuture; use std::{ cmp::{self}, - collections::HashMap, + collections::{HashMap, HashSet}, str::FromStr, sync::{Arc, RwLock}, }; @@ -246,7 +246,7 @@ pub struct InMemoryNodeInner { pub resolve_hashes: bool, pub console_log_handler: ConsoleLogHandler, pub system_contracts: SystemContracts, - pub impersonated_account: Option
, + pub impersonated_accounts: HashSet
, } type L2TxResult = ( @@ -561,16 +561,16 @@ impl InMemoryNodeInner { /// Sets the `impersonated_account` field of the node. /// This field is used to override the `tx.initiator_account` field of the transaction in the `run_l2_tx` method. - pub fn set_impersonated_account(&mut self, address: Address) { - self.impersonated_account = Some(address); + pub fn set_impersonated_account(&mut self, address: Address) -> bool { + return self.impersonated_accounts.insert(address); // Use SystemContracts without signature verification self.system_contracts = SystemContracts::from_options(&Options::BuiltInWithoutSecurity); } /// Clears the `impersonated_account` field of the node. - pub fn stop_impersonating_account(&mut self) { - self.impersonated_account = None; + pub fn stop_impersonating_account(&mut self, address: Address) -> bool { + return self.impersonated_accounts.remove(&address); // Restore previous SystemContracts self.system_contracts = SystemContracts::from_options(&Options::BuiltIn); @@ -653,7 +653,7 @@ impl InMemoryNode { resolve_hashes, console_log_handler: ConsoleLogHandler::default(), system_contracts: SystemContracts::from_options(system_contracts_options), - impersonated_account: None, + impersonated_accounts: Default::default(), } } else { let mut block_hashes = HashMap::::new(); @@ -683,7 +683,7 @@ impl InMemoryNode { resolve_hashes, console_log_handler: ConsoleLogHandler::default(), system_contracts: SystemContracts::from_options(system_contracts_options), - impersonated_account: None, + impersonated_accounts: Default::default(), } }; @@ -1126,19 +1126,21 @@ impl InMemoryNode { } /// Runs L2 transaction and commits it to a new block. - fn run_l2_tx(&self, mut l2_tx: L2Tx, execution_mode: TxExecutionMode) -> Result<(), String> { + fn run_l2_tx(&self, l2_tx: L2Tx, execution_mode: TxExecutionMode) -> Result<(), String> { let tx_hash = l2_tx.hash(); { let inner = self .inner .read() .map_err(|e| format!("Failed to acquire read lock: {}", e))?; - if let Some(impersonated_account) = inner.impersonated_account { + if inner + .impersonated_accounts + .contains(&l2_tx.common_data.initiator_address) + { tracing::info!( "🕵️ Executing tx from impersonated account {:?}", - impersonated_account + l2_tx.common_data.initiator_address ); - l2_tx.common_data.initiator_address = impersonated_account; } } log::info!("");