refactor: move impersonation management into its own component

Makefile

@@ -4,8 +4,8 @@ fetch-contracts:
 
 # Build the system contracts
 build-contracts:
-	cd contracts/system-contracts && yarn; yarn install; yarn build;
-	cd contracts/l2-contracts && yarn; yarn install; yarn build;
+	cd contracts/system-contracts && yarn install --frozen-lockfile; yarn build;
+	cd contracts/l2-contracts && yarn install --frozen-lockfile; yarn build;
 	./scripts/refresh_contracts.sh
 
 # Clean the system contracts

src/node/eth.rs

@@ -171,22 +171,16 @@ impl<S: ForkSource + std::fmt::Debug + Clone + Send + Sync + 'static> InMemoryNo
 
         l2_tx.set_input(bytes, hash);
 
+        if !self
+            .impersonation
+            .is_impersonating(&l2_tx.common_data.initiator_address)
         {
-            let reader = self
-                .inner
-                .read()
-                .map_err(|_| anyhow::anyhow!("Failed to acquire read lock for accounts."))?;
-            if !reader
-                .impersonated_accounts
-                .contains(&l2_tx.common_data.initiator_address)
-            {
-                let err = format!(
-                    "Initiator address {:?} is not allowed to perform transactions",
-                    l2_tx.common_data.initiator_address
-                );
-                tracing::error!("{err}");
-                return Err(TransparentError(err).into());
-            }
+            let err = format!(
+                "Initiator address {:?} is not allowed to perform transactions",
+                l2_tx.common_data.initiator_address
+            );
+            tracing::error!("{err}");
+            return Err(TransparentError(err).into());
         }
 
         self.seal_block(vec![l2_tx], system_contracts)
@@ -2949,7 +2943,7 @@ mod tests {
             .filters
             .add_block_filter()
             .expect("failed adding block filter");
-        inner.impersonated_accounts.insert(H160::repeat_byte(0x1));
+        inner.impersonation.impersonate(H160::repeat_byte(0x1));
         inner.rich_accounts.insert(H160::repeat_byte(0x1));
         inner
             .previous_states
@@ -2970,7 +2964,7 @@ mod tests {
             blocks: inner.blocks.clone(),
             block_hashes: inner.block_hashes.clone(),
             filters: inner.filters.clone(),
-            impersonated_accounts: inner.impersonated_accounts.clone(),
+            impersonated_accounts: inner.impersonation.impersonated_accounts(),
             rich_accounts: inner.rich_accounts.clone(),
             previous_states: inner.previous_states.clone(),
             raw_storage: storage.raw_storage.clone(),
@@ -3055,7 +3049,7 @@ mod tests {
             .filters
             .add_block_filter()
             .expect("failed adding block filter");
-        inner.impersonated_accounts.insert(H160::repeat_byte(0x1));
+        inner.impersonation.impersonate(H160::repeat_byte(0x1));
         inner.rich_accounts.insert(H160::repeat_byte(0x1));
         inner
             .previous_states
@@ -3077,7 +3071,7 @@ mod tests {
                 blocks: inner.blocks.clone(),
                 block_hashes: inner.block_hashes.clone(),
                 filters: inner.filters.clone(),
-                impersonated_accounts: inner.impersonated_accounts.clone(),
+                impersonated_accounts: inner.impersonation.impersonated_accounts(),
                 rich_accounts: inner.rich_accounts.clone(),
                 previous_states: inner.previous_states.clone(),
                 raw_storage: storage.raw_storage.clone(),
@@ -3108,7 +3102,7 @@ mod tests {
             .filters
             .add_pending_transaction_filter()
             .expect("failed adding pending transaction filter");
-        inner.impersonated_accounts.insert(H160::repeat_byte(0x2));
+        inner.impersonation.impersonate(H160::repeat_byte(0x2));
         inner.rich_accounts.insert(H160::repeat_byte(0x2));
         inner
             .previous_states
@@ -3148,7 +3142,7 @@ mod tests {
         assert_eq!(expected_snapshot.filters, inner.filters);
         assert_eq!(
             expected_snapshot.impersonated_accounts,
-            inner.impersonated_accounts
+            inner.impersonation.impersonated_accounts()
         );
         assert_eq!(expected_snapshot.rich_accounts, inner.rich_accounts);
         assert_eq!(expected_snapshot.previous_states, inner.previous_states);

src/node/impersonate.rs

@@ -0,0 +1,61 @@
+use std::collections::HashSet;
+use std::sync::{Arc, RwLock};
+use zksync_types::Address;
+
+/// Manages impersonated accounts across the system.
+///
+/// Clones always agree on the set of impersonated accounts and updating one affects all other
+/// instances.
+#[derive(Clone, Debug, Default)]
+pub struct ImpersonationManager {
+    state: Arc<RwLock<HashSet<Address>>>,
+}
+
+impl ImpersonationManager {
+    /// Starts impersonation for the provided account.
+    ///
+    /// Returns `true` if the account was not impersonated before.
+    pub fn impersonate(&self, addr: Address) -> bool {
+        tracing::trace!(?addr, "start impersonation");
+        let mut state = self
+            .state
+            .write()
+            .expect("ImpersonationManager lock is poisoned");
+        state.insert(addr)
+    }
+
+    /// Stops impersonation for the provided account.
+    ///
+    /// Returns `true` if the account was impersonated before.
+    pub fn stop_impersonating(&self, addr: &Address) -> bool {
+        tracing::trace!(?addr, "stop impersonation");
+        self.state
+            .write()
+            .expect("ImpersonationManager lock is poisoned")
+            .remove(addr)
+    }
+
+    /// Returns whether the provided account is currently impersonated.
+    pub fn is_impersonating(&self, addr: &Address) -> bool {
+        self.state
+            .read()
+            .expect("ImpersonationManager lock is poisoned")
+            .contains(addr)
+    }
+
+    /// Returns all accounts that are currently being impersonated.
+    pub fn impersonated_accounts(&self) -> HashSet<Address> {
+        self.state
+            .read()
+            .expect("ImpersonationManager lock is poisoned")
+            .clone()
+    }
+
+    /// Overrides currently impersonated accounts with the provided value.
+    pub fn set_impersonated_accounts(&self, accounts: HashSet<Address>) {
+        *self
+            .state
+            .write()
+            .expect("ImpersonationManager lock is poisoned") = accounts;
+    }
+}

src/node/in_memory.rs

@@ -46,6 +46,7 @@ use zksync_types::{
 use zksync_utils::{bytecode::hash_bytecode, h256_to_account_address, h256_to_u256, u256_to_h256};
 use zksync_web3_decl::error::Web3Error;
 
+use crate::node::impersonate::ImpersonationManager;
 use crate::node::time::TimestampManager;
 use crate::{
     bootloader_debug::{BootloaderDebug, BootloaderDebugTracer},
@@ -220,7 +221,7 @@ pub struct InMemoryNodeInner<S> {
     pub config: TestNodeConfig,
     pub console_log_handler: ConsoleLogHandler,
     pub system_contracts: SystemContracts,
-    pub impersonated_accounts: HashSet<Address>,
+    pub impersonation: ImpersonationManager,
     pub rich_accounts: HashSet<H160>,
     /// Keeps track of historical states indexed via block hash. Limited to [MAX_PREVIOUS_STATES].
     pub previous_states: IndexMap<H256, HashMap<StorageKey, StorageValue>>,
@@ -294,7 +295,7 @@ impl<S: std::fmt::Debug + ForkSource> InMemoryNodeInner<S> {
                     &updated_config.system_contracts_options,
                     updated_config.use_evm_emulator,
                 ),
-                impersonated_accounts: Default::default(),
+                impersonation: Default::default(),
                 rich_accounts: HashSet::new(),
                 previous_states: Default::default(),
                 observability,
@@ -329,7 +330,7 @@ impl<S: std::fmt::Debug + ForkSource> InMemoryNodeInner<S> {
                     &config.system_contracts_options,
                     config.use_evm_emulator,
                 ),
-                impersonated_accounts: Default::default(),
+                impersonation: Default::default(),
                 rich_accounts: HashSet::new(),
                 previous_states: Default::default(),
                 observability,
@@ -457,7 +458,7 @@ impl<S: std::fmt::Debug + ForkSource> InMemoryNodeInner<S> {
         let initiator_address = request_with_gas_per_pubdata_overridden
             .from
             .unwrap_or_default();
-        let impersonating = self.impersonated_accounts.contains(&initiator_address);
+        let impersonating = self.impersonation.is_impersonating(&initiator_address);
         let system_contracts = self
             .system_contracts
             .contracts_for_fee_estimate(impersonating)
@@ -769,17 +770,6 @@ impl<S: std::fmt::Debug + ForkSource> InMemoryNodeInner<S> {
         vm.execute(InspectExecutionMode::OneTx)
     }
 
-    /// 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) -> bool {
-        self.impersonated_accounts.insert(address)
-    }
-
-    /// Clears the `impersonated_account` field of the node.
-    pub fn stop_impersonating_account(&mut self, address: Address) -> bool {
-        self.impersonated_accounts.remove(&address)
-    }
-
     /// Archives the current state for later queries.
     pub fn archive_state(&mut self) -> Result<(), String> {
         if self.previous_states.len() > MAX_PREVIOUS_STATES as usize {
@@ -824,7 +814,7 @@ impl<S: std::fmt::Debug + ForkSource> InMemoryNodeInner<S> {
             blocks: self.blocks.clone(),
             block_hashes: self.block_hashes.clone(),
             filters: self.filters.clone(),
-            impersonated_accounts: self.impersonated_accounts.clone(),
+            impersonated_accounts: self.impersonation.impersonated_accounts(),
             rich_accounts: self.rich_accounts.clone(),
             previous_states: self.previous_states.clone(),
             raw_storage: storage.raw_storage.clone(),
@@ -851,7 +841,8 @@ impl<S: std::fmt::Debug + ForkSource> InMemoryNodeInner<S> {
         self.blocks = snapshot.blocks;
         self.block_hashes = snapshot.block_hashes;
         self.filters = snapshot.filters;
-        self.impersonated_accounts = snapshot.impersonated_accounts;
+        self.impersonation
+            .set_impersonated_accounts(snapshot.impersonated_accounts);
         self.rich_accounts = snapshot.rich_accounts;
         self.previous_states = snapshot.previous_states;
         storage.raw_storage = snapshot.raw_storage;
@@ -898,6 +889,7 @@ pub struct InMemoryNode<S: Clone> {
     #[allow(dead_code)]
     pub(crate) system_contracts_options: system_contracts::Options,
     pub(crate) time: TimestampManager,
+    pub(crate) impersonation: ImpersonationManager,
 }
 
 fn contract_address_from_tx_result(execution_result: &VmExecutionResultAndLogs) -> Option<H160> {
@@ -924,11 +916,13 @@ impl<S: ForkSource + std::fmt::Debug + Clone> InMemoryNode<S> {
         let system_contracts_options = config.system_contracts_options;
         let inner = InMemoryNodeInner::new(fork, observability, config);
         let time = inner.time.clone();
+        let impersonation = inner.impersonation.clone();
         InMemoryNode {
             inner: Arc::new(RwLock::new(inner)),
             snapshots: Default::default(),
             system_contracts_options,
             time,
+            impersonation,
         }
     }
 
@@ -1051,7 +1045,7 @@ impl<S: ForkSource + std::fmt::Debug + Clone> InMemoryNode<S> {
             .inner
             .read()
             .map_err(|_| anyhow::anyhow!("Failed to acquire read lock"))?;
-        Ok(if inner.impersonated_accounts.contains(&tx_initiator) {
+        Ok(if inner.impersonation.is_impersonating(&tx_initiator) {
             tracing::info!("🕵️ Executing tx from impersonated account {tx_initiator:?}");
             inner
                 .system_contracts

src/node/in_memory_ext.rs

@@ -303,36 +303,26 @@ impl<S: ForkSource + std::fmt::Debug + Clone + Send + Sync + 'static> InMemoryNo
     }
 
     pub fn impersonate_account(&self, address: Address) -> Result<bool> {
-        self.get_inner()
-            .write()
-            .map_err(|err| anyhow!("failed acquiring lock: {:?}", err))
-            .map(|mut writer| {
-                if writer.set_impersonated_account(address) {
-                    tracing::info!("🕵️ Account {:?} has been impersonated", address);
-                    true
-                } else {
-                    tracing::info!("🕵️ Account {:?} was already impersonated", address);
-                    false
-                }
-            })
+        if self.impersonation.impersonate(address) {
+            tracing::info!("🕵️ Account {:?} has been impersonated", address);
+            Ok(true)
+        } else {
+            tracing::info!("🕵️ Account {:?} was already impersonated", address);
+            Ok(false)
+        }
     }
 
     pub fn stop_impersonating_account(&self, address: Address) -> Result<bool> {
-        self.get_inner()
-            .write()
-            .map_err(|err| anyhow!("failed acquiring lock: {:?}", err))
-            .map(|mut writer| {
-                if writer.stop_impersonating_account(address) {
-                    tracing::info!("🕵️ Stopped impersonating account {:?}", address);
-                    true
-                } else {
-                    tracing::info!(
-                        "🕵️ Account {:?} was not impersonated, nothing to stop",
-                        address
-                    );
-                    false
-                }
-            })
+        if self.impersonation.stop_impersonating(&address) {
+            tracing::info!("🕵️ Stopped impersonating account {:?}", address);
+            Ok(true)
+        } else {
+            tracing::info!(
+                "🕵️ Account {:?} was not impersonated, nothing to stop",
+                address
+            );
+            Ok(false)
+        }
     }
 
     pub fn set_code(&self, address: Address, code: String) -> Result<()> {
@@ -502,18 +492,20 @@ mod tests {
             config: Default::default(),
             console_log_handler: Default::default(),
             system_contracts: Default::default(),
-            impersonated_accounts: Default::default(),
+            impersonation: Default::default(),
             rich_accounts: Default::default(),
             previous_states: Default::default(),
             observability: None,
         };
         let time = old_inner.time.clone();
+        let impersonation = old_inner.impersonation.clone();
 
         let node = InMemoryNode::<HttpForkSource> {
             inner: Arc::new(RwLock::new(old_inner)),
             snapshots: old_snapshots,
             system_contracts_options: old_system_contracts_options,
             time,
+            impersonation,
         };
 
         let address = Address::from_str("0x36615Cf349d7F6344891B1e7CA7C72883F5dc049").unwrap();

src/node/mod.rs

@@ -8,6 +8,7 @@ mod eth;
 mod evm;
 mod fee_model;
 mod hardhat;
+mod impersonate;
 mod in_memory;
 mod in_memory_ext;
 mod net;