fix nonce on batch and script

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

@@ -4,26 +4,42 @@ pragma solidity ^0.8.13;
 import {Script} from "forge-std/Script.sol";
 import {Greeter} from "../src/Greeter.sol";
 
+interface VmExt {
+    function zkGetTransactionNonce(
+        address account
+    ) external view returns (uint64 nonce);
+    function zkGetDeploymentNonce(
+        address account
+    ) external view returns (uint64 nonce);
+}
+
 contract ScriptSetupNonce is Script {
+    VmExt internal constant vmExt = VmExt(VM_ADDRESS);
+
     function setUp() public {
-        uint256 initial_nonce = checkNonce(address(tx.origin));
+        uint256 initialNonceTx = checkTxNonce(address(tx.origin));
+        uint256 initialNonceDeploy = checkDeployNonce(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);
+        assert(checkTxNonce(address(tx.origin)) == initialNonceTx);
+        assert(checkDeployNonce(address(tx.origin)) == initialNonceDeploy);
     }
 
     function run() public {
         // Get initial nonce
-        uint256 initial_nonce = checkNonce(address(tx.origin));
-        assert(initial_nonce == vm.getNonce(address(tx.origin)));
+        uint256 initialNonceTx = checkTxNonce(tx.origin);
+        uint256 initialNonceDeploy = checkDeployNonce(tx.origin);
+        assert(initialNonceTx == vmExt.zkGetTransactionNonce(tx.origin));
+        assert(initialNonceDeploy == vmExt.zkGetDeploymentNonce(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);
+        assert(checkTxNonce(tx.origin) == initialNonceTx);
+        assert(checkDeployNonce(tx.origin) == initialNonceDeploy);
 
         // Start broadcasting transactions
         vm.startBroadcast();
@@ -33,37 +49,71 @@ contract ScriptSetupNonce is Script {
 
         // Deploy checker and verify nonce
         NonceChecker checker = new NonceChecker();
+
+        vm.stopBroadcast();
+
         // 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();
+        checker.assertTxNonce(
+            vmExt.zkGetTransactionNonce(address(tx.origin)) + 1
+        );
+        checker.assertDeployNonce(
+            vmExt.zkGetDeploymentNonce(address(tx.origin))
+        );
     }
 
-    function checkNonce(address addr) public returns (uint256) {
+    function checkTxNonce(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);
+        return NonceLib.getTxNonce(addr);
+    }
+
+    function checkDeployNonce(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.getDeployNonce(addr);
     }
 }
 
 contract NonceChecker {
-    function checkNonce() public returns (uint256) {
-        return NonceLib.getNonce(address(tx.origin));
+    function checkTxNonce() public returns (uint256) {
+        return NonceLib.getTxNonce(address(tx.origin));
+    }
+
+    function checkDeployNonce() public returns (uint256) {
+        return NonceLib.getDeployNonce(address(tx.origin));
+    }
+
+    function assertTxNonce(uint256 expected) public {
+        uint256 real_nonce = checkTxNonce();
+        require(real_nonce == expected, "tx nonce mismatch");
     }
 
-    function assertNonce(uint256 expected) public {
-        uint256 real_nonce = checkNonce();
-        require(real_nonce == expected, "Nonce mismatch");
+    function assertDeployNonce(uint256 expected) public {
+        uint256 real_nonce = checkDeployNonce();
+        require(real_nonce == expected, "deploy 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));
+    function getTxNonce(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));
+    }
+
+    /// Retrieve tx nonce for `addr` from the NONCE_HOLDER system contract
+    function getDeployNonce(address addr) internal returns (uint256) {
+        (bool success, bytes memory data) = NONCE_HOLDER.call(
+            abi.encodeWithSignature("getDeploymentNonce(address)", addr)
+        );
         require(success, "Failed to get nonce");
         return abi.decode(data, (uint256));
     }

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

@@ -13,7 +13,7 @@ forgetest_async!(setup_block_on_script_test, |prj, cmd| {
         "./script/ScriptSetup.s.sol",
         "ScriptSetupNonce",
         None,
-        4,
+        3,
         Some(&["-vvvvv"]),
     )
     .await;

crates/strategy/zksync/src/cheatcode/runner/mod.rs

@@ -194,6 +194,18 @@ impl CheatcodeInspectorStrategyRunner for ZksyncCheatcodeInspectorStrategyRunner
             rpc,
             transaction: TransactionMaybeSigned::Unsigned(tx),
         });
+
+        // Explicitly increment tx nonce if calls are not isolated.
+        // This isn't needed in EVM, but required in zkEVM as the nonces are split.
+        if !config.evm_opts.isolate {
+            let tx_nonce = nonce;
+            foundry_zksync_core::set_tx_nonce(
+                broadcast.new_origin,
+                U256::from(tx_nonce.saturating_add(1)),
+                ecx_inner,
+            );
+            debug!(target: "cheatcodes", address=%broadcast.new_origin, new_tx_nonce=tx_nonce+1, tx_nonce, "incremented zksync tx nonce");
+        }
     }
 
     fn record_broadcastable_call_transactions(
@@ -292,15 +304,14 @@ impl CheatcodeInspectorStrategyRunner for ZksyncCheatcodeInspectorStrategyRunner
         });
         debug!(target: "cheatcodes", tx=?broadcastable_transactions.back().unwrap(), "broadcastable call");
 
-        // Explicitly increment nonce if calls are not isolated.
+        // Explicitly increment tx nonce if calls are not isolated.
         if !config.evm_opts.isolate {
             foundry_zksync_core::set_tx_nonce(
                 broadcast.new_origin,
                 U256::from(tx_nonce.saturating_add(1)),
                 ecx_inner,
             );
-            debug!(target: "cheatcodes", address=%broadcast.new_origin, nonce=tx_nonce+1,
-        tx_nonce, "incremented zksync tx nonce");
+            debug!(target: "cheatcodes", address=%broadcast.new_origin, new_tx_nonce=tx_nonce+1, tx_nonce, "incremented zksync tx nonce");
         }
     }
 

crates/zksync/core/src/lib.rs

@@ -253,7 +253,7 @@ where
     DB: Database,
     DB::Error: Debug,
 {
-    //ensure nonce is _only_ tx nonce
+    // 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();
@@ -269,6 +269,26 @@ where
     ecx.sstore(nonce_addr, nonce_key, updated_nonce.to_ru256()).expect("failed storing value");
 }
 
+/// Increment transaction nonce for a specific address.
+pub fn increment_tx_nonce<DB>(address: Address, ecx: &mut InnerEvmContext<DB>)
+where
+    DB: Database,
+    DB::Error: Debug,
+{
+    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 (tx_nonce, deploy_nonce) = ecx
+        .sload(nonce_addr, nonce_key)
+        .map(|v| decompose_full_nonce(v.to_u256()))
+        .unwrap_or_default();
+    let updated_nonce = nonces_to_full_nonce(tx_nonce.saturating_add(1u32.into()), 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

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

@@ -25,9 +25,13 @@ use zksync_multivm::{
     vm_latest::{HistoryDisabled, ToTracerPointer, Vm},
 };
 use zksync_types::{
-    get_nonce_key, h256_to_address, h256_to_u256, l2::L2Tx, transaction_request::PaymasterParams,
-    u256_to_h256, utils::nonces_to_full_nonce, PackedEthSignature, StorageKey, Transaction,
-    ACCOUNT_CODE_STORAGE_ADDRESS, CONTRACT_DEPLOYER_ADDRESS,
+    get_nonce_key, h256_to_address, h256_to_u256,
+    l2::L2Tx,
+    transaction_request::PaymasterParams,
+    u256_to_h256,
+    utils::{decompose_full_nonce, nonces_to_full_nonce},
+    PackedEthSignature, StorageKey, Transaction, ACCOUNT_CODE_STORAGE_ADDRESS,
+    CONTRACT_DEPLOYER_ADDRESS,
 };
 use zksync_vm_interface::storage::{ReadStorage, StoragePtr, WriteStorage};
 
@@ -41,7 +45,7 @@ use std::{
 use crate::{
     be_words_to_bytes,
     convert::{ConvertAddress, ConvertH160, ConvertH256, ConvertRU256, ConvertU256},
-    fix_l2_gas_limit, fix_l2_gas_price, is_system_address,
+    fix_l2_gas_limit, fix_l2_gas_price, increment_tx_nonce, is_system_address, set_tx_nonce,
     state::{new_full_nonce, parse_full_nonce, FullNonce},
     vm::{
         db::{ZKVMData, DEFAULT_CHAIN_ID},
@@ -104,6 +108,7 @@ where
         );
         tx.common_data.fee.gas_limit = new_gas_limit;
 
+        let initiator_address = tx.initiator_account();
         info!("executing batched tx ({}/{})", idx + 1, total_txns);
         let mut result = inspect(tx, ecx, ccx, call_ctx.clone())?;
 
@@ -151,6 +156,11 @@ where
             }
             _ => unreachable!("aggregated result must only contain success"),
         }
+
+        // Increment the nonce manually if there are multiple batches
+        if total_txns > 1 {
+            increment_tx_nonce(initiator_address.to_address(), ecx);
+        }
     }
 
     Ok(aggregated_result.expect("must have result"))
@@ -331,81 +341,15 @@ where
     let mut codes: rHashMap<Address, (B256, Bytecode)> = Default::default();
 
     let initiator_nonce_key = get_nonce_key(&initiator_address);
-    let caller_nonce_key = get_nonce_key(&call_ctx.msg_sender.to_h160());
 
     // NOTE(zk) We need to revert the tx nonce of the initiator as we intercept and dispatch
     // CALLs and CREATEs to the zkEVM. The CREATEs always increment the deployment nonce
-    // which must be persisted, but the CALLs are not real transactions and hence must
-    // be reverted.
-    if !call_ctx.is_create {
-        if let Some(initiator_nonce) = modified_storage.get_mut(&initiator_nonce_key) {
-            let FullNonce { tx_nonce, deploy_nonce } = parse_full_nonce(initiator_nonce.to_ru256());
-            let new_tx_nonce = tx_nonce.saturating_sub(1);
-            info!(address=?initiator_address, from=?tx_nonce, to=?new_tx_nonce, "reverting initiator tx nonce for CALL");
-            *initiator_nonce = new_full_nonce(new_tx_nonce, deploy_nonce).to_h256();
-        }
-    }
-
-    // NOTE(zk) We apply the nonce updates from the `tx_caller` to `msg_sender` if they differ.
-    // This is because zkEVM does not like being called from a non-EOA and it's simpler
-    // to call it with tx signer and adapt the nonce changes here. The nonce changes to `tx_caller`
-    // are thus reverted.
-    // We skip the balance changes as we override `msg.sender` in `executeTransaction` which takes
-    // care of it. The same cannot be done for `validateTransaction` due to the many safeguards
-    // around correct nonce update in the bootloader.
-    if initiator_address.to_address() != call_ctx.msg_sender {
-        if let (Some(initiator_nonce), Some(caller_nonce)) =
-            (modified_storage.get(&initiator_nonce_key), modified_storage.get(&caller_nonce_key))
-        {
-            let new_initiator_nonce = parse_full_nonce(initiator_nonce.to_ru256());
-            let new_caller_nonce = parse_full_nonce(caller_nonce.to_ru256());
-
-            let old_initiator_nonce = era_db.get_full_nonce(initiator_address.to_address());
-
-            let tx_nonce_inc =
-                new_initiator_nonce.tx_nonce.saturating_sub(old_initiator_nonce.tx_nonce);
-            let deploy_nonce_inc =
-                new_initiator_nonce.deploy_nonce.saturating_sub(old_initiator_nonce.deploy_nonce);
-
-            let corrected_initiator_nonce = FullNonce {
-                tx_nonce: old_initiator_nonce.tx_nonce,
-                deploy_nonce: old_initiator_nonce.deploy_nonce,
-            };
-            let corrected_caller_nonce = FullNonce {
-                tx_nonce: new_caller_nonce.tx_nonce.saturating_add(tx_nonce_inc),
-                deploy_nonce: new_caller_nonce.deploy_nonce.saturating_add(deploy_nonce_inc),
-            };
-
-            modified_storage.insert(
-                initiator_nonce_key,
-                nonces_to_full_nonce(
-                    old_initiator_nonce.tx_nonce.into(),
-                    old_initiator_nonce.deploy_nonce.into(),
-                )
-                .to_h256(),
-            );
-            modified_storage.insert(
-                caller_nonce_key,
-                nonces_to_full_nonce(
-                    corrected_caller_nonce.tx_nonce.into(),
-                    corrected_caller_nonce.deploy_nonce.into(),
-                )
-                .to_h256(),
-            );
-
-            info!(
-                address=?initiator_address,
-                from=?new_initiator_nonce,
-                to=?corrected_initiator_nonce,
-                "correcting initiator nonce",
-            );
-            info!(
-                address=?call_ctx.msg_sender,
-                from=?new_caller_nonce,
-                to=?corrected_caller_nonce,
-                "correcting caller nonce",
-            );
-        }
+    // which must be persisted, but the tx nonce increase must be reverted.
+    if let Some(initiator_nonce) = modified_storage.get_mut(&initiator_nonce_key) {
+        let FullNonce { tx_nonce, deploy_nonce } = parse_full_nonce(initiator_nonce.to_ru256());
+        let new_tx_nonce = tx_nonce.saturating_sub(1);
+        error!(address=?initiator_address, from=?tx_nonce, to=?new_tx_nonce, deploy_nonce, "reverting initiator tx nonce for CALL");
+        *initiator_nonce = new_full_nonce(new_tx_nonce, deploy_nonce).to_h256();
     }
 
     for (k, v) in modified_storage {