fix: do not handle expectCall in zkEVM (#807)

* skip base call * add test --------- Co-authored-by: Federico Rodríguez <[email protected]>
matter-labs · Jan 7, 2025 · 9be4421 · 9be4421
1 parent 8c10790
commit 9be4421
Show file tree

Hide file tree

Showing 4 changed files with 79 additions and 14 deletions.
diff --git a/crates/forge/tests/it/zk/cheats.rs b/crates/forge/tests/it/zk/cheats.rs
@@ -79,6 +79,14 @@ async fn test_zk_cheat_expect_emit_works() {
     TestConfig::with_filter(runner, filter).spec_id(SpecId::SHANGHAI).run().await;
 }
 
+#[tokio::test(flavor = "multi_thread")]
+async fn test_zk_cheat_expect_call_works() {
+    let runner = TEST_DATA_DEFAULT.runner_zksync();
+    let filter = Filter::new("testExpectCall", "ZkCheatcodesTest", ".*");
+
+    TestConfig::with_filter(runner, filter).spec_id(SpecId::SHANGHAI).run().await;
+}
+
 #[tokio::test(flavor = "multi_thread")]
 async fn test_zk_cheat_mock_with_value_function() {
     let runner = TEST_DATA_DEFAULT.runner_zksync();

diff --git a/crates/zksync/core/src/vm/runner.rs b/crates/zksync/core/src/vm/runner.rs
@@ -77,6 +77,7 @@ where
         tx_caller: env.tx.caller,
         msg_sender: env.tx.caller,
         contract: transact_to.to_address(),
+        input: if is_create { None } else { Some(env.tx.data.clone()) },
         delegate_as: None,
         block_number: env.block.number,
         block_timestamp: env.block.timestamp,
@@ -192,6 +193,7 @@ where
         tx_caller: ecx.env.tx.caller,
         msg_sender,
         contract: CONTRACT_DEPLOYER_ADDRESS.to_address(),
+        input: None,
         delegate_as: None,
         block_number: ecx.env.block.number,
         block_timestamp: ecx.env.block.timestamp,
@@ -258,6 +260,7 @@ where
         tx_caller: ecx.env.tx.caller,
         msg_sender: call.caller,
         contract: call.bytecode_address,
+        input: Some(call.input.clone()),
         delegate_as: match call.scheme {
             CallScheme::DelegateCall => Some(call.target_address),
             _ => None,

diff --git a/crates/zksync/core/src/vm/tracers/cheatcode.rs b/crates/zksync/core/src/vm/tracers/cheatcode.rs
@@ -10,6 +10,7 @@ use foundry_cheatcodes_common::{
     mock::{MockCallDataContext, MockCallReturnData},
     record::RecordAccess,
 };
+use tracing::debug;
 use zksync_multivm::{
     interface::tracer::TracerExecutionStatus,
     tracers::dynamic::vm_1_5_0::DynTracer,
@@ -117,6 +118,8 @@ pub struct CallContext {
     pub msg_sender: Address,
     /// Target contract's address.
     pub contract: Address,
+    /// Target contract's input (if CALL).
+    pub input: Option<Bytes>,
     /// Delegated contract's address. This is used
     /// to override `address(this)` for delegate calls.
     pub delegate_as: Option<Address>,
@@ -213,21 +216,36 @@ impl<S: ReadStorage, H: HistoryMode> DynTracer<S, SimpleMemory<H>> for Cheatcode
                 self.expected_calls.get_mut(&current.code_address.to_address())
             {
                 let calldata = get_calldata(&state, memory);
-                // Match every partial/full calldata
-                for (expected_calldata, (expected, actual_count)) in expected_calls_for_target {
-                    // Increment actual times seen if...
-                    // The calldata is at most, as big as this call's input, and
-                    if expected_calldata.len() <= calldata.len() &&
-                    // Both calldata match, taking the length of the assumed smaller one (which will have at least the selector), and
-                    *expected_calldata == calldata[..expected_calldata.len()] &&
-                    // The value matches, if provided
-                    expected
-                        .value
-                        .map_or(true, |value|{
-                             value == rU256::from(current.context_u128_value)})
-                    {
-                        *actual_count += 1;
+
+                // We skip recording the base call for `expectCall` cheatcode that initiated this
+                // transaction. The initial call is recorded in revm when it was
+                // made, and before being dispatched to zkEVM.
+                let is_base_call = current.code_address.to_address() == self.call_context.contract &&
+                    self.call_context
+                        .input
+                        .as_ref()
+                        .map(|input| input.0.to_vec() == calldata)
+                        .unwrap_or_default();
+
+                if !is_base_call {
+                    // Match every partial/full calldata
+                    for (expected_calldata, (expected, actual_count)) in expected_calls_for_target {
+                        // Increment actual times seen if...
+                        // The calldata is at most, as big as this call's input, and
+                        if expected_calldata.len() <= calldata.len() &&
+                        // Both calldata match, taking the length of the assumed smaller one (which will have at least the selector), and
+                        *expected_calldata == calldata[..expected_calldata.len()] &&
+                        // The value matches, if provided
+                        expected
+                            .value
+                            .map_or(true, |value|{
+                                value == rU256::from(current.context_u128_value)})
+                        {
+                            *actual_count += 1;
+                        }
                     }
+                } else {
+                    debug!("skip recording base call in zkEVM for expectCall cheatcode");
                 }
             }
         }

diff --git a/testdata/zk/Cheatcodes.t.sol b/testdata/zk/Cheatcodes.t.sol
@@ -49,6 +49,32 @@ contract MyProxyCaller {
     }
 }
 
+contract Proxy {
+    InnerProxy innerProxy;
+
+    constructor(InnerProxy _innerProxy) {
+        innerProxy = _innerProxy;
+    }
+
+    function proxyCall(ICaller caller) public {
+        innerProxy.proxyCall(caller);
+    }
+}
+
+contract InnerProxy {
+    function proxyCall(ICaller caller) public {
+        caller.call(10);
+    }
+}
+
+interface ICaller {
+    function call(uint8 _value) external;
+}
+
+contract Caller is ICaller {
+    function call(uint8 _value) external {}
+}
+
 contract Emitter {
     event EventConstructor(string message);
     event EventFunction(string message);
@@ -269,6 +295,16 @@ contract ZkCheatcodesTest is DSTest {
         assertEq(zkvmEntries.length, evmEntries.length);
     }
 
+    function testExpectCallCountsTopMostCallOnce() public {
+        InnerProxy innerProxy = new InnerProxy();
+        Proxy proxy = new Proxy(innerProxy);
+        Caller caller = new Caller();
+        vm.expectCall(address(proxy), abi.encodeCall(proxy.proxyCall, (caller)), 1);
+        vm.expectCall(address(innerProxy), abi.encodeCall(innerProxy.proxyCall, (caller)), 1);
+        vm.expectCall(address(caller), abi.encodeCall(caller.call, (10)), 1);
+        proxy.proxyCall(caller);
+    }
+
     // Utility function
     function getCodeCheck(string memory contractName, string memory outDir) internal {
         bytes memory bytecode = vm.getCode(contractName);