feat: add anvil_autoImpersonateAccount (#416)

* make `anvil_setNonce` consistent with upstream

* clarify that `anvil_reset` behaves consistently with upstream

There is no expectation that we
should support local block revert

* implement `anvil_autoImpersonateAccount`

* clippy

* update supported apis

* fix tests

SUPPORTED_APIS.md

@@ -14,10 +14,11 @@ The `status` options are:
 
 | Namespace | API | <div style="width:130px">Status</div> | Description |
 | --- | --- | --- | --- |
+| `ANVIL` | `anvil_autoImpersonateAccount` | `SUPPORTED` | Sets auto impersonation status.|
 | `ANVIL` | `anvil_setNonce` | `SUPPORTED` | Sets the nonce of an address.|
 | `ANVIL` | `anvil_impersonateAccount` | `SUPPORTED` | Impersonate an account |
 | `ANVIL` | `anvil_stopImpersonatingAccount` | `SUPPORTED` | Stop impersonating an account after having previously used `anvil_impersonateAccount` |
-| `ANVIL` | `anvil_reset` | `PARTIALLY` | Resets the state of the network; cannot revert to past block numbers, unless they're in a fork |
+| `ANVIL` | `anvil_reset` | `SUPPORTED` | Resets the state of the network; cannot revert to past block numbers, unless they're in a fork |
 | `ANVIL` | `anvil_mine` | `SUPPORTED` | Mine any number of blocks at once, in constant time |
 | `ANVIL` | `anvil_setBalance` | `SUPPORTED` | Modifies the balance of an account |
 | `ANVIL` | `anvil_setCode` | `SUPPORTED` | Sets the bytecode of a given account |

e2e-tests/test/anvil-apis.test.ts

@@ -28,15 +28,36 @@ describe("anvil_setBalance", function () {
 describe("anvil_setNonce", function () {
   it("Should update the nonce of an account", async function () {
     // Arrange
+    const richWallet = new Wallet(RichAccounts[0].PrivateKey).connect(provider);
     const userWallet = Wallet.createRandom().connect(provider);
+
+    // Simply asserts that `richWallet` can still send successful transactions
+    async function assertCanSendTx() {
+      const tx = {
+        to: userWallet.address,
+        value: ethers.utils.parseEther("0.42"),
+      };
+
+      const txResponse = await richWallet.sendTransaction(tx);
+      const txReceipt = await txResponse.wait();
+      expect(txReceipt.status).to.equal(1);
+    }
+
     const newNonce = 42;
 
-    // Act
-    await provider.send("anvil_setNonce", [userWallet.address, ethers.utils.hexlify(newNonce)]);
+    // Advance nonce to 42
+    await provider.send("anvil_setNonce", [richWallet.address, ethers.utils.hexlify(newNonce)]);
 
     // Assert
-    const nonce = await userWallet.getNonce();
-    expect(nonce).to.equal(newNonce);
+    expect(await richWallet.getNonce()).to.equal(newNonce);
+    await assertCanSendTx();
+
+    // Rollback nonce to 0
+    await provider.send("anvil_setNonce", [richWallet.address, ethers.utils.hexlify(0)]);
+
+    // Assert
+    expect(await richWallet.getNonce()).to.equal(0);
+    await assertCanSendTx();
   });
 });
 
@@ -88,6 +109,33 @@ describe("anvil_impersonateAccount & anvil_stopImpersonatingAccount", function (
   });
 });
 
+describe("anvil_autoImpersonateAccount", function () {
+  it("Should allow transfers of funds without knowing the Private Key", async function () {
+    // Arrange
+    const userWallet = Wallet.createRandom().connect(provider);
+    const richAccount = RichAccounts[6].Account;
+    const beforeBalance = await provider.getBalance(richAccount);
+
+    // Act
+    await provider.send("anvil_autoImpersonateAccount", [true]);
+
+    const signer = await ethers.getSigner(richAccount);
+    const tx = {
+      to: userWallet.address,
+      value: ethers.utils.parseEther("0.42"),
+    };
+
+    const recieptTx = await signer.sendTransaction(tx);
+    await recieptTx.wait();
+
+    await provider.send("anvil_autoImpersonateAccount", [false]);
+
+    // Assert
+    expect((await userWallet.getBalance()).eq(ethers.utils.parseEther("0.42"))).to.true;
+    expect((await provider.getBalance(richAccount)).eq(beforeBalance.sub(ethers.utils.parseEther("0.42")))).to.true;
+  });
+});
+
 describe("anvil_setCode", function () {
   it("Should set code at an address", async function () {
     // Arrange

src/namespaces/anvil.rs

@@ -5,6 +5,18 @@ use super::{ResetRequest, RpcResult};
 
 #[rpc]
 pub trait AnvilNamespaceT {
+    /// Sets auto impersonation status.
+    ///
+    /// # Arguments
+    ///
+    /// * `enabled` - `true` makes every account impersonated, `false` disables this behavior
+    ///
+    /// # Returns
+    ///
+    /// A `BoxFuture` containing a `Result` representing the success of the operation.
+    #[rpc(name = "anvil_autoImpersonateAccount")]
+    fn auto_impersonate_account(&self, enabled: bool) -> RpcResult<()>;
+
     /// Sets the balance of the given address to the given balance.
     ///
     /// # Arguments
@@ -29,7 +41,7 @@ pub trait AnvilNamespaceT {
     ///
     /// A `BoxFuture` containing a `Result` with a `bool` representing the success of the operation.
     #[rpc(name = "anvil_setNonce")]
-    fn set_nonce(&self, address: Address, balance: U256) -> RpcResult<bool>;
+    fn set_nonce(&self, address: Address, nonce: U256) -> RpcResult<bool>;
 
     /// Sometimes you may want to advance the latest block number of the network by a large number of blocks.
     /// One way to do this would be to call the evm_mine RPC method multiple times, but this is too slow if you want to mine thousands of blocks.

src/node/anvil.rs

@@ -11,6 +11,11 @@ use crate::{
 impl<S: ForkSource + std::fmt::Debug + Clone + Send + Sync + 'static> AnvilNamespaceT
     for InMemoryNode<S>
 {
+    fn auto_impersonate_account(&self, enabled: bool) -> RpcResult<()> {
+        self.auto_impersonate_account(enabled);
+        Ok(()).into_boxed_future()
+    }
+
     fn set_balance(&self, address: Address, balance: U256) -> RpcResult<bool> {
         self.set_balance(address, balance)
             .map_err(|err| {
@@ -20,8 +25,8 @@ impl<S: ForkSource + std::fmt::Debug + Clone + Send + Sync + 'static> AnvilNames
             .into_boxed_future()
     }
 
-    fn set_nonce(&self, address: Address, balance: U256) -> RpcResult<bool> {
-        self.set_nonce(address, balance)
+    fn set_nonce(&self, address: Address, nonce: U256) -> RpcResult<bool> {
+        self.set_nonce(address, nonce)
             .map_err(|err| {
                 tracing::error!("failed setting nonce: {:?}", err);
                 into_jsrpc_error(Web3Error::InternalError(err))

src/node/eth.rs

@@ -2967,7 +2967,7 @@ mod tests {
             blocks: inner.blocks.clone(),
             block_hashes: inner.block_hashes.clone(),
             filters: inner.filters.clone(),
-            impersonated_accounts: inner.impersonation.impersonated_accounts(),
+            impersonation_state: inner.impersonation.state(),
             rich_accounts: inner.rich_accounts.clone(),
             previous_states: inner.previous_states.clone(),
             raw_storage: storage.raw_storage.clone(),
@@ -3004,8 +3004,8 @@ mod tests {
         assert_eq!(expected_snapshot.block_hashes, actual_snapshot.block_hashes);
         assert_eq!(expected_snapshot.filters, actual_snapshot.filters);
         assert_eq!(
-            expected_snapshot.impersonated_accounts,
-            actual_snapshot.impersonated_accounts
+            expected_snapshot.impersonation_state,
+            actual_snapshot.impersonation_state
         );
         assert_eq!(
             expected_snapshot.rich_accounts,
@@ -3074,7 +3074,7 @@ mod tests {
                 blocks: inner.blocks.clone(),
                 block_hashes: inner.block_hashes.clone(),
                 filters: inner.filters.clone(),
-                impersonated_accounts: inner.impersonation.impersonated_accounts(),
+                impersonation_state: inner.impersonation.state(),
                 rich_accounts: inner.rich_accounts.clone(),
                 previous_states: inner.previous_states.clone(),
                 raw_storage: storage.raw_storage.clone(),
@@ -3144,8 +3144,8 @@ mod tests {
         assert_eq!(expected_snapshot.block_hashes, inner.block_hashes);
         assert_eq!(expected_snapshot.filters, inner.filters);
         assert_eq!(
-            expected_snapshot.impersonated_accounts,
-            inner.impersonation.impersonated_accounts()
+            expected_snapshot.impersonation_state,
+            inner.impersonation.state()
         );
         assert_eq!(expected_snapshot.rich_accounts, inner.rich_accounts);
         assert_eq!(expected_snapshot.previous_states, inner.previous_states);

src/node/impersonate.rs

@@ -8,10 +8,20 @@ use zksync_types::Address;
 /// instances.
 #[derive(Clone, Debug, Default)]
 pub struct ImpersonationManager {
-    state: Arc<RwLock<HashSet<Address>>>,
+    state: Arc<RwLock<ImpersonationState>>,
 }
 
 impl ImpersonationManager {
+    /// Sets the auto impersonation flag, when `true` it makes all accounts impersonated by default.
+    /// Setting to `false` disabled this behavior.
+    pub fn set_auto_impersonation(&self, enabled: bool) {
+        tracing::trace!(enabled, "auto impersonation status set");
+        self.state
+            .write()
+            .expect("ImpersonationManager lock is poisoned")
+            .auto = enabled
+    }
+
     /// Starts impersonation for the provided account.
     ///
     /// Returns `true` if the account was not impersonated before.
@@ -21,7 +31,7 @@ impl ImpersonationManager {
             .state
             .write()
             .expect("ImpersonationManager lock is poisoned");
-        state.insert(addr)
+        state.accounts.insert(addr)
     }
 
     /// Stops impersonation for the provided account.
@@ -32,30 +42,40 @@ impl ImpersonationManager {
         self.state
             .write()
             .expect("ImpersonationManager lock is poisoned")
+            .accounts
             .remove(addr)
     }
 
     /// Returns whether the provided account is currently impersonated.
     pub fn is_impersonating(&self, addr: &Address) -> bool {
-        self.state
+        let state = self
+            .state
             .read()
-            .expect("ImpersonationManager lock is poisoned")
-            .contains(addr)
+            .expect("ImpersonationManager lock is poisoned");
+        state.auto || state.accounts.contains(addr)
     }
 
-    /// Returns all accounts that are currently being impersonated.
-    pub fn impersonated_accounts(&self) -> HashSet<Address> {
+    /// Returns internal state representation.
+    pub fn state(&self) -> ImpersonationState {
         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>) {
+    /// Overrides current internal state with the provided value.
+    pub fn set_state(&self, state: ImpersonationState) {
         *self
             .state
             .write()
-            .expect("ImpersonationManager lock is poisoned") = accounts;
+            .expect("ImpersonationManager lock is poisoned") = state;
     }
 }
+
+#[derive(Clone, Debug, Default, PartialEq)]
+pub struct ImpersonationState {
+    /// If `true` then all accounts are impersonated regardless of `accounts` contents
+    pub auto: bool,
+    /// Accounts that are currently impersonated
+    pub accounts: HashSet<Address>,
+}

src/node/in_memory.rs

@@ -46,7 +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::impersonate::{ImpersonationManager, ImpersonationState};
 use crate::node::time::TimestampManager;
 use crate::{
     bootloader_debug::{BootloaderDebug, BootloaderDebugTracer},
@@ -805,7 +805,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.impersonation.impersonated_accounts(),
+            impersonation_state: self.impersonation.state(),
             rich_accounts: self.rich_accounts.clone(),
             previous_states: self.previous_states.clone(),
             raw_storage: storage.raw_storage.clone(),
@@ -832,8 +832,7 @@ impl<S: std::fmt::Debug + ForkSource> InMemoryNodeInner<S> {
         self.blocks = snapshot.blocks;
         self.block_hashes = snapshot.block_hashes;
         self.filters = snapshot.filters;
-        self.impersonation
-            .set_impersonated_accounts(snapshot.impersonated_accounts);
+        self.impersonation.set_state(snapshot.impersonation_state);
         self.rich_accounts = snapshot.rich_accounts;
         self.previous_states = snapshot.previous_states;
         storage.raw_storage = snapshot.raw_storage;
@@ -859,7 +858,7 @@ pub struct Snapshot {
     pub(crate) blocks: HashMap<H256, Block<TransactionVariant>>,
     pub(crate) block_hashes: HashMap<u64, H256>,
     pub(crate) filters: EthFilters,
-    pub(crate) impersonated_accounts: HashSet<Address>,
+    pub(crate) impersonation_state: ImpersonationState,
     pub(crate) rich_accounts: HashSet<H160>,
     pub(crate) previous_states: IndexMap<H256, HashMap<StorageKey, StorageValue>>,
     pub(crate) raw_storage: InMemoryStorage,

src/node/in_memory_ext.rs

@@ -1,11 +1,11 @@
 use anyhow::anyhow;
 use zksync_types::{
     get_code_key, get_nonce_key,
-    utils::{decompose_full_nonce, nonces_to_full_nonce, storage_key_for_eth_balance},
+    utils::{nonces_to_full_nonce, storage_key_for_eth_balance},
     StorageKey,
 };
 use zksync_types::{AccountTreeId, Address, U256, U64};
-use zksync_utils::{h256_to_u256, u256_to_h256};
+use zksync_utils::u256_to_h256;
 
 use crate::{
     fork::{ForkDetails, ForkSource},
@@ -179,33 +179,9 @@ impl<S: ForkSource + std::fmt::Debug + Clone + Send + Sync + 'static> InMemoryNo
         self.get_inner()
             .write()
             .map_err(|err| anyhow!("failed acquiring lock: {:?}", err))
-            .and_then(|mut writer| {
+            .map(|mut writer| {
                 let nonce_key = get_nonce_key(&address);
-                let full_nonce = match writer.fork_storage.read_value_internal(&nonce_key) {
-                    Ok(full_nonce) => full_nonce,
-                    Err(error) => {
-                        return Err(anyhow!(error.to_string()));
-                    }
-                };
-                let (mut account_nonce, mut deployment_nonce) =
-                    decompose_full_nonce(h256_to_u256(full_nonce));
-                if account_nonce >= nonce {
-                    return Err(anyhow!(
-                        "Account Nonce is already set to a higher value ({}, requested {})",
-                        account_nonce,
-                        nonce
-                    ));
-                }
-                account_nonce = nonce;
-                if deployment_nonce >= nonce {
-                    return Err(anyhow!(
-                        "Deployment Nonce is already set to a higher value ({}, requested {})",
-                        deployment_nonce,
-                        nonce
-                    ));
-                }
-                deployment_nonce = nonce;
-                let enforced_full_nonce = nonces_to_full_nonce(account_nonce, deployment_nonce);
+                let enforced_full_nonce = nonces_to_full_nonce(nonce, nonce);
                 tracing::info!(
                     "👷 Nonces for address {:?} have been set to {}",
                     address,
@@ -214,7 +190,7 @@ impl<S: ForkSource + std::fmt::Debug + Clone + Send + Sync + 'static> InMemoryNo
                 writer
                     .fork_storage
                     .set_value(nonce_key, u256_to_h256(enforced_full_nonce));
-                Ok(true)
+                true
             })
     }
 
@@ -302,6 +278,10 @@ impl<S: ForkSource + std::fmt::Debug + Clone + Send + Sync + 'static> InMemoryNo
         }
     }
 
+    pub fn auto_impersonate_account(&self, enabled: bool) {
+        self.impersonation.set_auto_impersonation(enabled);
+    }
+
     pub fn impersonate_account(&self, address: Address) -> Result<bool> {
         if self.impersonation.impersonate(address) {
             tracing::info!("🕵️ Account {:?} has been impersonated", address);
@@ -378,6 +358,7 @@ mod tests {
     use zksync_multivm::interface::storage::ReadStorage;
     use zksync_types::{api::BlockNumber, fee::Fee, l2::L2Tx, PackedEthSignature};
     use zksync_types::{Nonce, H256};
+    use zksync_utils::h256_to_u256;
 
     #[tokio::test]
     async fn test_set_balance() {
@@ -408,9 +389,11 @@ mod tests {
         assert_eq!(nonce_after, U256::from(1337));
         assert_ne!(nonce_before, nonce_after);
 
-        // setting nonce lower than the current one should fail
-        let result = node.set_nonce(address, U256::from(1336));
-        assert!(result.is_err());
+        let result = node.set_nonce(address, U256::from(1336)).unwrap();
+        assert!(result);
+
+        let nonce_after = node.get_transaction_count(address, None).await.unwrap();
+        assert_eq!(nonce_after, U256::from(1336));
     }
 
     #[tokio::test]