fix: Nonce mismatches what network expects (#726)

* Nonce mismatch on setup with script

* Add toggle to update nonce only on first operation to avoid nonce mismatch with the network

* Remove prints and format test

* Update crates/cheatcodes/src/inspector.rs

Co-authored-by: Federico Rodríguez <[email protected]>

* Update crates/cheatcodes/src/inspector.rs

Co-authored-by: Federico Rodríguez <[email protected]>

* Update crates/forge/tests/fixtures/zk/ScriptSetup.s.sol

Co-authored-by: Federico Rodríguez <[email protected]>

* Update crates/cheatcodes/src/inspector.rs

Co-authored-by: Federico Rodríguez <[email protected]>

* change variable and test name

* Change test name of noncemismatch

* Add description, filter before iteration for loop and take() in cheatcode context creation

* Remove taking in inspector of zk should update nonce

* chore: explicit behavior for nonce update persistence (#729)

* explicit behavior for nonce update persistence

* prevent short-circuit bug

* enable persistance only once per test contract deployment

* Format and change test

* Forge fmt

* improve comments

* Add comment to test

---------

Co-authored-by: Federico Rodríguez <[email protected]>
Co-authored-by: Nisheeth Barthwal <[email protected]>

crates/cheatcodes/src/inspector.rs

@@ -410,6 +410,39 @@ impl ArbitraryStorage {
 /// List of transactions that can be broadcasted.
 pub type BroadcastableTransactions = VecDeque<BroadcastableTransaction>;
 
+/// Allows overriding nonce update behavior for the tx caller in the zkEVM.
+///
+/// Since each CREATE or CALL is executed as a separate transaction within zkEVM, we currently skip
+/// persisting nonce updates as it erroneously increments the tx nonce. However, under certain
+/// situations, e.g. deploying contracts, transacts, etc. the nonce updates must be persisted.
+#[derive(Default, Debug, Clone)]
+pub enum ZkPersistNonceUpdate {
+    /// Never update the nonce. This is currently the default behavior.
+    #[default]
+    Never,
+    /// Override the default behavior, and persist nonce update for tx caller for the next
+    /// zkEVM execution _only_.
+    PersistNext,
+}
+
+impl ZkPersistNonceUpdate {
+    /// Persist nonce update for the tx caller for next execution.
+    pub fn persist_next(&mut self) {
+        *self = Self::PersistNext;
+    }
+
+    /// Retrieve if a nonce update must be persisted, or not. Resets the state to default.
+    pub fn check(&mut self) -> bool {
+        let persist_nonce_update = match self {
+            Self::Never => false,
+            Self::PersistNext => true,
+        };
+        *self = Default::default();
+
+        persist_nonce_update
+    }
+}
+
 /// Setting for migrating the database to zkEVM storage when starting in ZKsync mode.
 /// The migration is performed on the DB via the inspector so must only be performed once.
 #[derive(Debug, Default, Clone)]
@@ -608,6 +641,9 @@ pub struct Cheatcodes {
     /// This can be done as each test runs with its own [Cheatcodes] instance, thereby
     /// providing the necessary level of isolation.
     pub persisted_factory_deps: HashMap<H256, Vec<u8>>,
+
+    /// Nonce update persistence behavior in zkEVM for the tx caller.
+    pub zk_persist_nonce_update: ZkPersistNonceUpdate,
 }
 
 // This is not derived because calling this in `fn new` with `..Default::default()` creates a second
@@ -705,6 +741,7 @@ impl Cheatcodes {
             persisted_factory_deps: Default::default(),
             paymaster_params: None,
             zk_use_factory_deps: Default::default(),
+            zk_persist_nonce_update: Default::default(),
         }
     }
 
@@ -1231,13 +1268,16 @@ impl Cheatcodes {
         let factory_deps = self.dual_compiled_contracts.fetch_all_factory_deps(contract);
         tracing::debug!(contract = contract.name, "using dual compiled contract");
 
+        let zk_persist_nonce_update = self.zk_persist_nonce_update.check();
         let ccx = foundry_zksync_core::vm::CheatcodeTracerContext {
             mocked_calls: self.mocked_calls.clone(),
             expected_calls: Some(&mut self.expected_calls),
             accesses: self.accesses.as_mut(),
             persisted_factory_deps: Some(&mut self.persisted_factory_deps),
             paymaster_data: self.paymaster_params.take(),
+            persist_nonce_update: self.broadcast.is_some() || zk_persist_nonce_update,
         };
+
         let zk_create = foundry_zksync_core::vm::ZkCreateInputs {
             value: input.value().to_u256(),
             msg_sender: input.caller(),
@@ -1508,9 +1548,12 @@ where {
                 }
             };
             let prev = account.info.nonce;
-            account.info.nonce = prev.saturating_sub(1);
+            let nonce = prev.saturating_sub(1);
+            account.info.nonce = nonce;
+            // NOTE(zk): We sync with the nonce changes to ensure that the nonce matches
+            foundry_zksync_core::cheatcodes::set_nonce(sender, U256::from(nonce), ecx_inner);
 
-            trace!(target: "cheatcodes", %sender, nonce=account.info.nonce, prev, "corrected nonce");
+            trace!(target: "cheatcodes", %sender, nonce, prev, "corrected nonce");
         }
 
         if call.target_address == CHEATCODE_ADDRESS {
@@ -1855,13 +1898,14 @@ where {
         }
 
         info!("running call in zkEVM {:#?}", call);
-
+        let zk_persist_nonce_update = self.zk_persist_nonce_update.check();
         let ccx = foundry_zksync_core::vm::CheatcodeTracerContext {
             mocked_calls: self.mocked_calls.clone(),
             expected_calls: Some(&mut self.expected_calls),
             accesses: self.accesses.as_mut(),
             persisted_factory_deps: Some(&mut self.persisted_factory_deps),
             paymaster_data: self.paymaster_params.take(),
+            persist_nonce_update: self.broadcast.is_some() || zk_persist_nonce_update,
         };
 
         let mut gas = Gas::new(call.gas_limit);

crates/forge/src/runner.rs

@@ -148,12 +148,16 @@ impl ContractRunner<'_> {
         self.executor.deploy_create2_deployer()?;
 
         // Test contract has already been deployed so we can migrate the database to zkEVM storage
-        // in the next runner execution.
+        // in the next runner execution. Additionally we can allow persisting the next nonce update
+        // to simulate EVM behavior where only the tx that deploys the test contract increments the
+        // nonce.
         if let Some(cheatcodes) = &mut self.executor.inspector.cheatcodes {
             if let Some(zk_startup_migration) = &mut cheatcodes.zk_startup_migration {
                 debug!("test contract deployed, allowing startup storage migration");
                 zk_startup_migration.allow();
             }
+            debug!("test contract deployed, allowing persisting next nonce update");
+            cheatcodes.zk_persist_nonce_update.persist_next();
         }
 
         // Optionally call the `setUp` function

crates/forge/tests/fixtures/zk/ScriptSetup.s.sol

@@ -0,0 +1,70 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.13;
+
+import {Script} from "forge-std/Script.sol";
+import {Greeter} from "../src/Greeter.sol";
+
+contract ScriptSetupNonce is Script {
+    function setUp() public {
+        uint256 initial_nonce = checkNonce(address(tx.origin));
+        // Perform transactions and deploy contracts in setup to increment nonce and verify broadcast nonce matches onchain
+        new Greeter();
+        new Greeter();
+        new Greeter();
+        new Greeter();
+        assert(checkNonce(address(tx.origin)) == initial_nonce);
+    }
+
+    function run() public {
+        // Get initial nonce
+        uint256 initial_nonce = checkNonce(address(tx.origin));
+        assert(initial_nonce == vm.getNonce(address(tx.origin)));
+
+        // Create and interact with non-broadcasted contract to verify nonce is not incremented
+        Greeter notBroadcastGreeter = new Greeter();
+        notBroadcastGreeter.greeting("john");
+        assert(checkNonce(address(tx.origin)) == initial_nonce);
+
+        // Start broadcasting transactions
+        vm.startBroadcast();
+        // Deploy and interact with broadcasted contracts
+        Greeter greeter = new Greeter();
+        greeter.greeting("john");
+
+        // Deploy checker and verify nonce
+        NonceChecker checker = new NonceChecker();
+        // We expect the nonce to be incremented by 1 because the check is done in an external
+        // call
+        checker.assertNonce(vm.getNonce(address(tx.origin)) + 1);
+        vm.stopBroadcast();
+    }
+
+    function checkNonce(address addr) public returns (uint256) {
+        // We prank here to avoid accidentally "polluting" the nonce of `addr` during the call
+        // for example when `addr` is `tx.origin`
+        vm.prank(address(this), address(this));
+        return NonceLib.getNonce(addr);
+    }
+}
+
+contract NonceChecker {
+    function checkNonce() public returns (uint256) {
+        return NonceLib.getNonce(address(tx.origin));
+    }
+
+    function assertNonce(uint256 expected) public {
+        uint256 real_nonce = checkNonce();
+        require(real_nonce == expected, "Nonce mismatch");
+    }
+}
+
+library NonceLib {
+    address constant NONCE_HOLDER = address(0x8003);
+
+    /// Retrieve tx nonce for `addr` from the NONCE_HOLDER system contract
+    function getNonce(address addr) internal returns (uint256) {
+        (bool success, bytes memory data) = NONCE_HOLDER.call(abi.encodeWithSignature("getMinNonce(address)", addr));
+        require(success, "Failed to get nonce");
+        return abi.decode(data, (uint256));
+    }
+}

crates/forge/tests/it/zk/mod.rs

@@ -13,6 +13,7 @@ mod invariant;
 mod linking;
 mod logs;
 mod nft;
+mod nonce;
 mod ownership;
 mod paymaster;
 mod proxy;

crates/forge/tests/it/zk/nonce.rs

@@ -0,0 +1,34 @@
+use crate::{
+    config::TestConfig,
+    test_helpers::{run_zk_script_test, TEST_DATA_DEFAULT},
+};
+use forge::revm::primitives::SpecId;
+use foundry_test_utils::{forgetest_async, util, Filter, TestProject};
+
+forgetest_async!(setup_block_on_script_test, |prj, cmd| {
+    setup_deploy_prj(&mut prj);
+    run_zk_script_test(
+        prj.root(),
+        &mut cmd,
+        "./script/ScriptSetup.s.sol",
+        "ScriptSetupNonce",
+        None,
+        4,
+        Some(&["-vvvvv"]),
+    );
+});
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_zk_contract_nonce_mismatch() {
+    let runner = TEST_DATA_DEFAULT.runner_zksync();
+    let filter = Filter::new("testTxOriginNonceDoesNotUpdate", "NonceMismatchTest", ".*");
+
+    TestConfig::with_filter(runner, filter).evm_spec(SpecId::SHANGHAI).run().await;
+}
+
+fn setup_deploy_prj(prj: &mut TestProject) {
+    util::initialize(prj.root());
+    prj.add_script("ScriptSetup.s.sol", include_str!("../../fixtures/zk/ScriptSetup.s.sol"))
+        .unwrap();
+    prj.add_source("Greeter.sol", include_str!("../../../../../testdata/zk/Greeter.sol")).unwrap();
+}

crates/script/src/runner.rs

@@ -171,12 +171,16 @@ impl ScriptRunner {
         traces.extend(constructor_traces.map(|traces| (TraceKind::Deployment, traces)));
 
         // Script has already been deployed so we can migrate the database to zkEVM storage
-        // in the next runner execution.
+        // in the next runner execution. Additionally we can allow persisting the next nonce update
+        // to simulate EVM behavior where only the tx that deploys the test contract increments the
+        // nonce.
         if let Some(cheatcodes) = &mut self.executor.inspector.cheatcodes {
             if let Some(zk_startup_migration) = &mut cheatcodes.zk_startup_migration {
                 debug!("script deployed, allowing startup storage migration");
                 zk_startup_migration.allow();
             }
+            debug!("script deployed, allowing persisting next nonce update");
+            cheatcodes.zk_persist_nonce_update.persist_next();
         }
 
         // Optionally call the `setUp` function

crates/zksync/core/src/cheatcodes.rs

@@ -86,15 +86,7 @@ where
     <DB as Database>::Error: Debug,
 {
     info!(?address, ?nonce, "cheatcode setNonce");
-    //ensure nonce is _only_ tx nonce
-    let (tx_nonce, _deploy_nonce) = decompose_full_nonce(nonce.to_u256());
-
-    let nonce_addr = NONCE_HOLDER_ADDRESS.to_address();
-    ecx.load_account(nonce_addr).expect("account could not be loaded");
-    let zk_address = address.to_h160();
-    let nonce_key = get_nonce_key(&zk_address).key().to_ru256();
-    ecx.touch(&nonce_addr);
-    ecx.sstore(nonce_addr, nonce_key, tx_nonce.to_ru256()).expect("failed storing value");
+    crate::set_tx_nonce(address, nonce, ecx);
 }
 
 /// Gets nonce for a specific address.

crates/zksync/core/src/lib.rs

@@ -27,10 +27,12 @@ use alloy_signer::Signature;
 use alloy_transport::Transport;
 use convert::{
     ConvertAddress, ConvertBytes, ConvertH160, ConvertH256, ConvertRU256, ConvertSignature,
-    ToSignable,
+    ConvertU256, ToSignable,
 };
 use eyre::{eyre, OptionExt};
+use revm::{Database, InnerEvmContext};
 use serde::{Deserialize, Serialize};
+use std::fmt::Debug;
 
 pub use utils::{fix_l2_gas_limit, fix_l2_gas_price};
 pub use vm::{balance, encode_create_params, nonce};
@@ -42,7 +44,10 @@ pub use zksync_types::{
     IMMUTABLE_SIMULATOR_STORAGE_ADDRESS, KNOWN_CODES_STORAGE_ADDRESS, L2_BASE_TOKEN_ADDRESS,
     NONCE_HOLDER_ADDRESS,
 };
-use zksync_types::{utils::storage_key_for_eth_balance, U256};
+use zksync_types::{
+    utils::{decompose_full_nonce, nonces_to_full_nonce, storage_key_for_eth_balance},
+    U256,
+};
 pub use zksync_utils::bytecode::hash_bytecode;
 pub use zksync_web3_rs::{
     eip712::{Eip712Meta, Eip712Transaction, Eip712TransactionRequest, PaymasterParams},
@@ -324,6 +329,50 @@ pub fn get_immutable_slot_key(address: Address, slot_index: rU256) -> H256 {
     H256(*immutable_value_key)
 }
 
+/// Sets transaction nonce for a specific address.
+pub fn set_tx_nonce<DB>(address: Address, nonce: rU256, ecx: &mut InnerEvmContext<DB>)
+where
+    DB: Database,
+    DB::Error: Debug,
+{
+    //ensure nonce is _only_ tx nonce
+    let (tx_nonce, _deploy_nonce) = decompose_full_nonce(nonce.to_u256());
+
+    let nonce_addr = NONCE_HOLDER_ADDRESS.to_address();
+    ecx.load_account(nonce_addr).expect("account could not be loaded");
+    let nonce_key = get_nonce_key(address);
+    ecx.touch(&nonce_addr);
+    // We make sure to keep the old deployment nonce
+    let old_deploy_nonce = ecx
+        .sload(nonce_addr, nonce_key)
+        .map(|v| decompose_full_nonce(v.to_u256()).1)
+        .unwrap_or_default();
+    let updated_nonce = nonces_to_full_nonce(tx_nonce, old_deploy_nonce);
+    ecx.sstore(nonce_addr, nonce_key, updated_nonce.to_ru256()).expect("failed storing value");
+}
+
+/// Sets deployment nonce for a specific address.
+pub fn set_deployment_nonce<DB>(address: Address, nonce: rU256, ecx: &mut InnerEvmContext<DB>)
+where
+    DB: Database,
+    DB::Error: Debug,
+{
+    //ensure nonce is _only_ deployment nonce
+    let (_tx_nonce, deploy_nonce) = decompose_full_nonce(nonce.to_u256());
+
+    let nonce_addr = NONCE_HOLDER_ADDRESS.to_address();
+    ecx.load_account(nonce_addr).expect("account could not be loaded");
+    let nonce_key = get_nonce_key(address);
+    ecx.touch(&nonce_addr);
+    // We make sure to keep the old transaction nonce
+    let old_tx_nonce = ecx
+        .sload(nonce_addr, nonce_key)
+        .map(|v| decompose_full_nonce(v.to_u256()).0)
+        .unwrap_or_default();
+    let updated_nonce = nonces_to_full_nonce(old_tx_nonce, deploy_nonce);
+    ecx.sstore(nonce_addr, nonce_key, updated_nonce.to_ru256()).expect("failed storing value");
+}
+
 #[cfg(test)]
 mod tests {
     use std::str::FromStr;

crates/zksync/core/src/vm/inspect.rs

@@ -27,8 +27,8 @@ use zksync_multivm::{
 };
 use zksync_state::interface::{ReadStorage, StoragePtr, WriteStorage};
 use zksync_types::{
-    l2::L2Tx, PackedEthSignature, StorageKey, Transaction, ACCOUNT_CODE_STORAGE_ADDRESS,
-    CONTRACT_DEPLOYER_ADDRESS,
+    get_nonce_key, l2::L2Tx, PackedEthSignature, StorageKey, Transaction,
+    ACCOUNT_CODE_STORAGE_ADDRESS, CONTRACT_DEPLOYER_ADDRESS, NONCE_HOLDER_ADDRESS,
 };
 use zksync_utils::{be_words_to_bytes, h256_to_account_address, h256_to_u256, u256_to_h256};
 
@@ -190,6 +190,9 @@ where
     let is_create = call_ctx.is_create;
     info!(?call_ctx, "executing transaction in zk vm");
 
+    let initiator_address = tx.common_data.initiator_address;
+    let persist_nonce_update = ccx.persist_nonce_update;
+
     if tx.common_data.signature.is_empty() {
         // FIXME: This is a hack to make sure that the signature is not empty.
         // Fails without a signature here: https://github.com/matter-labs/zksync-era/blob/73a1e8ff564025d06e02c2689da238ae47bb10c3/core/lib/types/src/transaction_request.rs#L381
@@ -329,7 +332,14 @@ where
 
     let mut storage: rHashMap<Address, rHashMap<rU256, StorageSlot>> = Default::default();
     let mut codes: rHashMap<Address, (B256, Bytecode)> = Default::default();
-    for (k, v) in &modified_storage {
+    // We skip nonce updates when should_update_nonce is false to avoid nonce mismatch
+    let filtered = modified_storage.iter().filter(|(k, _)| {
+        !(k.address() == &NONCE_HOLDER_ADDRESS &&
+            get_nonce_key(&initiator_address) == **k &&
+            !persist_nonce_update)
+    });
+
+    for (k, v) in filtered {
         let address = k.address().to_address();
         let index = k.key().to_ru256();
         era_db.load_account(address);

crates/zksync/core/src/vm/runner.rs

@@ -84,7 +84,11 @@ where
         is_static: false,
     };
 
-    let mut ccx = CheatcodeTracerContext { persisted_factory_deps, ..Default::default() };
+    let mut ccx = CheatcodeTracerContext {
+        persisted_factory_deps,
+        persist_nonce_update: true,
+        ..Default::default()
+    };
 
     match inspect::<_, DB::Error>(tx, &mut ecx, &mut ccx, call_ctx) {
         Ok(ZKVMExecutionResult { execution_result: result, .. }) => {