Add toggle to update nonce only on first operation to avoid nonce mis…

…match with the network

crates/cheatcodes/src/inspector.rs

@@ -1206,9 +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.should_update_nonce.unwrap_or_default(),
+            should_update_nonce: self.broadcast.is_some() ||
+                self.should_update_nonce.unwrap_or_default(),
         };
-        
+
         println!("should_update_nonce??: {:?}", ccx.should_update_nonce);
 
         let zk_create = foundry_zksync_core::vm::ZkCreateInputs {
@@ -1481,9 +1482,12 @@ where {
                 }
             };
             let prev = account.info.nonce;
-            account.info.nonce = prev.saturating_sub(1);
+            let nonce = prev.saturating_sub(1);
+            account.info.nonce = nonce;
+            // 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 {
@@ -1837,7 +1841,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.should_update_nonce.unwrap_or_default(),
+            should_update_nonce: self.broadcast.is_some() ||
+                self.should_update_nonce.unwrap_or_default(),
         };
 
         let mut gas = Gas::new(call.gas_limit);

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 testNonceMismatch() 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

@@ -9,26 +9,25 @@ contract ScriptSetupNonce is Script {
 
     function setUp() public {
         // Perform transactions and deploy contracts in setup to increment nonce and verify broadcast nonce matches onchain
-        // new Greeter();
-        // new Greeter();
         new Greeter();
-        // new Greeter();
+        new Greeter();
+        new Greeter();
+        new Greeter();
     }
 
     function run() public {
-        // NonceChecker checker1 = new NonceChecker();
-        // checker1.checkNonce();
-        // checker1.checkNonce();
-        // vm.getNonce(address(tx.origin));
-        // checkNonce(tx.origin);
-        // vm.getNonce(address(tx.origin));
+        NonceChecker checker1 = new NonceChecker();
+        checker1.checkNonce();
+        checker1.checkNonce();
+        vm.getNonce(address(tx.origin));
+        checkNonce(tx.origin);
+        vm.getNonce(address(tx.origin));
         vm.startBroadcast();
         Greeter greeter = new Greeter();
-        // greeter.greeting("john");
-        // NonceChecker checker = new NonceChecker();
-        // NonceChecker checker2 = new NonceChecker();
-        // checker.assertNonce(vm.getNonce(address(tx.origin)) + 1);
-        // checker.assertDeployNonce(2);
+        greeter.greeting("john");
+        NonceChecker checker = new NonceChecker();
+        NonceChecker checker2 = new NonceChecker();
+        checker.assertNonce(vm.getNonce(address(tx.origin)) + 1);
         vm.stopBroadcast();
     }
 
@@ -58,10 +57,4 @@ contract NonceChecker {
         console.log("real_nonce", real_nonce);
         require(real_nonce == expected, "Nonce mismatch");
     }
-
-    function assertDeployNonce(uint256 expected) public {
-        uint256 real_deploy_nonce = checkDeployNonce();
-        console.log("real_deploy_nonce", real_deploy_nonce);
-        require(real_deploy_nonce == expected, "Deploy nonce mismatch");
-    }
 }

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

@@ -16,5 +16,5 @@ mod ownership;
 mod paymaster;
 mod proxy;
 mod repros;
+mod setup_script;
 mod traces;
-mod setup_script;

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

@@ -1,4 +1,5 @@
-use foundry_test_utils::{forgetest_async, util, TestProject};
+use foundry_config::fs_permissions::PathPermission;
+use foundry_test_utils::{forgetest_async, util, util::OutputExt, TestProject};
 
 use crate::test_helpers::run_zk_script_test;
 
@@ -10,15 +11,38 @@ forgetest_async!(setup_block_on_script_test, |prj, cmd| {
         "./script/ScriptSetup.s.sol",
         "ScriptSetupNonce",
         None,
-        2,
+        5,
         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,7 +27,8 @@ use zksync_multivm::{
 };
 use zksync_state::interface::{ReadStorage, StoragePtr, WriteStorage};
 use zksync_types::{
-    get_nonce_key, l2::L2Tx, PackedEthSignature, StorageKey, Transaction, ACCOUNT_CODE_STORAGE_ADDRESS, CONTRACT_DEPLOYER_ADDRESS, NONCE_HOLDER_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};
 
@@ -331,22 +332,19 @@ 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);
         let previous = era_db.sload(address, index);
         let entry = storage.entry(address).or_default();
         entry.insert(index, StorageSlot::new_changed(previous, v.to_ru256()));
 
-        if k.address() == &NONCE_HOLDER_ADDRESS && get_nonce_key(&initiator_address) == *k {
-            if !should_update_nonce {
-                println!("Ignoring nonce update for {:?}", initiator_address);
-                continue;
-            }
-            println!("Updating nonce for {:?}", initiator_address);
-        }
-
         if k.address() == &ACCOUNT_CODE_STORAGE_ADDRESS {
             if let Some(bytecode) = bytecodes.get(&h256_to_u256(*v)) {
                 let bytecode =

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, .. }) => {