fix(anvil): also remove code when impersonating (#2696)

Author: mattsse

anvil/src/eth/backend/cheats.rs

@@ -2,6 +2,7 @@
 
 use anvil_core::eth::transaction::TypedTransaction;
 use ethers::types::{Address, Signature, H256, U256};
+use forge::revm::Bytecode;
 use parking_lot::RwLock;
 use std::{collections::HashMap, sync::Arc};
 use tracing::trace;
@@ -32,15 +33,20 @@ impl CheatsManager {
     /// This also accepts the actual code hash if the address is a contract to bypass EIP-3607
     ///
     /// Returns `true` if the account is already impersonated
-    pub fn impersonate(&self, addr: Address, code_hash: Option<H256>) -> bool {
+    pub fn impersonate(
+        &self,
+        addr: Address,
+        code_hash: Option<H256>,
+        code: Option<Bytecode>,
+    ) -> bool {
         trace!(target: "cheats", "Start impersonating {:?}", addr);
-        self.state.write().impersonated_account.insert(addr, code_hash).is_some()
+        self.state.write().impersonated_account.insert(addr, (code_hash, code)).is_some()
     }
 
     /// Removes the account that from the impersonated set
-    pub fn stop_impersonating(&self, addr: &Address) -> Option<H256> {
+    pub fn stop_impersonating(&self, addr: &Address) -> Option<(Option<H256>, Option<Bytecode>)> {
         trace!(target: "cheats", "Stop impersonating {:?}", addr);
-        self.state.write().impersonated_account.remove(addr).flatten()
+        self.state.write().impersonated_account.remove(addr)
     }
 
     /// Returns true if the `addr` is currently impersonated
@@ -62,7 +68,7 @@ pub struct CheatsState {
     /// If the account is a contract it holds the hash of the contracts code that is temporarily
     /// set to `KECCAK_EMPTY` to bypass EIP-3607 which rejects transactions from senders with
     /// deployed code
-    pub impersonated_account: HashMap<Address, Option<H256>>,
+    pub impersonated_account: HashMap<Address, (Option<H256>, Option<Bytecode>)>,
     /// The signature used for the `eth_sendUnsignedTransaction` cheat code
     pub bypass_signature: Signature,
 }

anvil/src/eth/backend/mem/mod.rs

@@ -203,22 +203,28 @@ impl Backend {
             return true
         }
         // need to bypass EIP-3607: Reject transactions from senders with deployed code by setting
-        // the code hash to `KECCAK_EMPTY` temporarily
+        // the code hash to `KECCAK_EMPTY` temporarily and also remove the code itself and add back
+        // when we stop impersonating
         let mut account = self.db.read().await.basic(addr);
         let mut code_hash = None;
+        let mut code = None;
         if account.code_hash != KECCAK_EMPTY {
             code_hash = Some(std::mem::replace(&mut account.code_hash, KECCAK_EMPTY));
+            code = account.code.take();
             self.db.write().await.insert_account(addr, account);
         }
-        self.cheats.impersonate(addr, code_hash)
+        self.cheats.impersonate(addr, code_hash, code)
     }
 
     /// Removes the account that from the impersonated set
+    ///
+    /// If the impersonated `addr` is a contract then we also reset the code here
     pub async fn stop_impersonating(&self, addr: Address) {
-        if let Some(code_hash) = self.cheats.stop_impersonating(&addr) {
+        if let Some((Some(code_hash), code)) = self.cheats.stop_impersonating(&addr) {
             let mut db = self.db.write().await;
             let mut account = db.basic(addr);
             account.code_hash = code_hash;
+            account.code = code;
             db.insert_account(addr, account)
         }
     }
@@ -1119,6 +1125,10 @@ impl Backend {
         self.with_database_at(number, |db| {
             trace!(target: "backend", "get code for {:?}", address);
             let account = db.basic(address);
+            if account.code_hash == KECCAK_EMPTY {
+                // if the code hash is `KECCAK_EMPTY`, we check no further
+                return Ok(Default::default())
+            }
             let code = if let Some(code) = account.code {
                 code
             } else {

anvil/tests/it/anvil_api.rs

@@ -1,5 +1,5 @@
 //! tests for custom anvil endpoints
-use crate::abi::*;
+use crate::{abi::*, fork::fork_config};
 use anvil::{spawn, Hardfork, NodeConfig};
 use anvil_core::eth::EthRequest;
 use ethers::{
@@ -117,6 +117,37 @@ async fn can_impersonate_contract() {
     assert_eq!("Hello World!", greeting);
 }
 
+#[tokio::test(flavor = "multi_thread")]
+async fn can_impersonate_gnosis_safe() {
+    let (api, handle) = spawn(fork_config()).await;
+    let provider = handle.http_provider();
+
+    // <https://help.gnosis-safe.io/en/articles/4971293-i-don-t-remember-my-safe-address-where-can-i-find-it>
+    let safe: Address = "0xA063Cb7CFd8E57c30c788A0572CBbf2129ae56B6".parse().unwrap();
+
+    let code = provider.get_code(safe, None).await.unwrap();
+    assert!(!code.is_empty());
+
+    api.anvil_impersonate_account(safe).await.unwrap();
+
+    let code = provider.get_code(safe, None).await.unwrap();
+    // impersonated contract code is temporarily removed
+    assert!(code.is_empty());
+
+    let balance = U256::from(1e18 as u64);
+    // fund the impersonated account
+    api.anvil_set_balance(safe, balance).await.unwrap();
+
+    let on_chain_balance = provider.get_balance(safe, None).await.unwrap();
+    assert_eq!(on_chain_balance, balance);
+
+    api.anvil_stop_impersonating_account(safe).await.unwrap();
+
+    let code = provider.get_code(safe, None).await.unwrap();
+    // code is added back after stop impersonating
+    assert!(!code.is_empty());
+}
+
 #[tokio::test(flavor = "multi_thread")]
 async fn can_mine_manually() {
     let (api, handle) = spawn(NodeConfig::test()).await;