fix: Nonce mismatches what network expects

crates/cheatcodes/src/inspector.rs

@@ -575,6 +575,8 @@ 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>>,
+
+    pub zk_should_update_nonce: Option<bool>,
 }
 
 // This is not derived because calling this in `fn new` with `..Default::default()` creates a second
@@ -671,6 +673,7 @@ impl Cheatcodes {
             persisted_factory_deps: Default::default(),
             paymaster_params: None,
             zk_use_factory_deps: Default::default(),
+            zk_should_update_nonce: None,
         }
     }
 
@@ -1132,6 +1135,8 @@ impl Cheatcodes {
             }
         }
 
+        self.zk_should_update_nonce.take();
+
         None
     }
 
@@ -1201,7 +1206,10 @@ impl Cheatcodes {
             accesses: self.accesses.as_mut(),
             persisted_factory_deps: Some(&mut self.persisted_factory_deps),
             paymaster_data: self.paymaster_params.take(),
+            should_update_nonce: self.broadcast.is_some() ||
+                self.zk_should_update_nonce.unwrap_or_default(),
         };
+
         let zk_create = foundry_zksync_core::vm::ZkCreateInputs {
             value: input.value().to_u256(),
             msg_sender: input.caller(),
@@ -1472,9 +1480,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 {
@@ -1783,6 +1794,8 @@ where {
             }
         }
 
+        self.zk_should_update_nonce.take();
+
         None
     }
 
@@ -1826,6 +1839,8 @@ where {
             accesses: self.accesses.as_mut(),
             persisted_factory_deps: Some(&mut self.persisted_factory_deps),
             paymaster_data: self.paymaster_params.take(),
+            should_update_nonce: self.broadcast.is_some() ||
+                self.zk_should_update_nonce.unwrap_or_default(),
         };
 
         let mut gas = Gas::new(call.gas_limit);

crates/evm/evm/src/executors/mod.rs

@@ -428,6 +428,9 @@ impl Executor {
     #[instrument(name = "call", level = "debug", skip_all)]
     pub fn call_with_env(&self, mut env: EnvWithHandlerCfg) -> eyre::Result<RawCallResult> {
         let mut inspector = self.inspector().clone();
+        if let Some(cheatcodes) = inspector.cheatcodes.as_mut() {
+            cheatcodes.zk_should_update_nonce = Some(true);
+        }
         let mut backend = CowBackend::new_borrowed(self.backend());
         let result = match &self.zk_tx {
             None => backend.inspect(&mut env, &mut inspector)?,
@@ -450,6 +453,9 @@ impl Executor {
     #[instrument(name = "transact", level = "debug", skip_all)]
     pub fn transact_with_env(&mut self, mut env: EnvWithHandlerCfg) -> eyre::Result<RawCallResult> {
         let mut inspector = self.inspector.clone();
+        if let Some(cheatcodes) = inspector.cheatcodes.as_mut() {
+            cheatcodes.zk_should_update_nonce = Some(true);
+        }
         let backend = &mut self.backend;
         let result_and_state = match self.zk_tx.take() {
             None => backend.inspect(&mut env, &mut inspector)?,

crates/forge/tests/fixtures/zk/NonceMismatch.t.sol

@@ -0,0 +1,26 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.18;
+
+import "forge-std/Test.sol";
+import {Greeter} from "../src/Greeter.sol";
+
+contract NonceMismatchTest is Test {
+    function setUp() public {
+        // Deploy contracts in setup to increment nonce
+        new Greeter();
+        new Greeter();
+        new Greeter();
+        new Greeter();
+    }
+
+    function testTxOriginNonceDoesNotUpdate() public {
+        uint256 nonce = vm.getNonce(address(tx.origin));
+        assertEq(nonce, 2);
+
+        // Deploy another contract
+        new Greeter();
+
+        uint256 newNonce = vm.getNonce(address(tx.origin));
+        assertEq(newNonce, 2);
+    }
+}

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

@@ -0,0 +1,55 @@
+// 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 {
+        uint256 initial_nonce = checkNonce(address(tx.origin));
+        assert(initial_nonce == vm.getNonce(address(tx.origin)));
+        Greeter notBroadcastGreeter = new Greeter();
+        notBroadcastGreeter.greeting("john");
+        assert(checkNonce(address(tx.origin)) == initial_nonce);
+        vm.startBroadcast();
+        Greeter greeter = new Greeter();
+        greeter.greeting("john");
+        NonceChecker checker = new NonceChecker();
+        checker.assertNonce(vm.getNonce(address(tx.origin)) + 1);
+        vm.stopBroadcast();
+    }
+
+    function checkNonce(address addr) public returns (uint256) {
+        vm.prank(address(this), address(this));
+        (bool success, bytes memory data) = address(0x000000000000000000000000000000000000008003).call(
+            abi.encodeWithSignature("getMinNonce(address)", addr)
+        );
+        require(success, "Failed to get nonce");
+        return abi.decode(data, (uint256));
+    }
+}
+
+contract NonceChecker {
+    function checkNonce() public returns (uint256) {
+        (bool success, bytes memory data) = address(0x000000000000000000000000000000000000008003).call(
+            abi.encodeWithSignature("getMinNonce(address)", address(tx.origin))
+        );
+        require(success, "Failed to get nonce");
+        return abi.decode(data, (uint256));
+    }
+
+    function assertNonce(uint256 expected) public {
+        uint256 real_nonce = checkNonce();
+        require(real_nonce == expected, "Nonce mismatch");
+    }
+}

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

@@ -12,6 +12,7 @@ mod fuzz;
 mod invariant;
 mod logs;
 mod nft;
+mod nonce;
 mod ownership;
 mod paymaster;
 mod proxy;

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

@@ -0,0 +1,48 @@
+use foundry_config::fs_permissions::PathPermission;
+use foundry_test_utils::{forgetest_async, util, util::OutputExt, TestProject};
+
+use crate::test_helpers::run_zk_script_test;
+
+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 (prj, mut cmd) = util::setup_forge(
+        "test_zk_contract_nonce_mismatch",
+        foundry_test_utils::foundry_compilers::PathStyle::Dapptools,
+    );
+    util::initialize(prj.root());
+
+    cmd.args(["install", "matter-labs/era-contracts", "--no-commit", "--shallow"]).assert_success();
+    cmd.forge_fuse();
+
+    let mut config = cmd.config();
+    config.fs_permissions.add(PathPermission::read("./zkout"));
+    prj.write_config(config);
+
+    prj.add_source("Greeter.sol", include_str!("../../../../../testdata/zk/Greeter.sol")).unwrap();
+
+    prj.add_test("NonceMismatch.t.sol", include_str!("../../fixtures/zk/NonceMismatch.t.sol"))
+        .unwrap();
+
+    cmd.args(["test", "--evm-version", "shanghai", "--mc", "NonceMismatchTest"]);
+    cmd.assert_success().get_output().stdout_lossy().contains("Suite result: ok");
+}
+
+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/zksync/core/src/cheatcodes.rs

@@ -9,7 +9,7 @@ use tracing::info;
 use zksync_types::{
     block::{pack_block_info, unpack_block_info},
     get_nonce_key,
-    utils::{decompose_full_nonce, storage_key_for_eth_balance},
+    utils::{decompose_full_nonce, nonces_to_full_nonce, storage_key_for_eth_balance},
     ACCOUNT_CODE_STORAGE_ADDRESS, CURRENT_VIRTUAL_BLOCK_INFO_POSITION, KNOWN_CODES_STORAGE_ADDRESS,
     L2_BASE_TOKEN_ADDRESS, NONCE_HOLDER_ADDRESS, SYSTEM_CONTEXT_ADDRESS,
 };
@@ -94,7 +94,13 @@ where
     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");
+    // 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");
 }
 
 /// Gets nonce for a specific address.

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 should_update_nonce = ccx.should_update_nonce;
+
     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,12 @@ 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
+    for (k, v) in modified_storage.iter().filter(|(k, _)| {
+        !(k.address() == &NONCE_HOLDER_ADDRESS &&
+            get_nonce_key(&initiator_address) == **k &&
+            !should_update_nonce)
+    }) {
         let address = k.address().to_address();
         let index = k.key().to_ru256();
         era_db.load_account(address);

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

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

crates/zksync/core/src/vm/tracers/cheatcode.rs

@@ -90,6 +90,8 @@ pub struct CheatcodeTracerContext<'a> {
     pub persisted_factory_deps: Option<&'a mut HashMap<H256, Vec<u8>>>,
     /// Paymaster data
     pub paymaster_data: Option<ZkPaymasterData>,
+    /// Whether to persist the nonce updates or not
+    pub should_update_nonce: bool,
 }
 
 /// Tracer result to return back to foundry.