feat(cheatcodes): Make expectCall only work for the next call's sub…

…calls (#5032) * chore: make expect call only work for the next call * chore: make expectCall actually check only the next call's subcalls * chore: fmt * chore: introduce checks at the main call level, not at the subcall level * chore: handle dangling expected calls gracefully * chore: fix tests * chore: fmt * chore: forge fmt * chore: actually exclude depth the cheatcode was called from * chore: tests * chore: better docs * chore: comment out impossible to check condition on expectCall * chore: remove unused check
foundry-rs · May 25, 2023 · 84b7cbe · 84b7cbe
1 parent 1e78cab
commit 84b7cbe
Show file tree

Hide file tree

Showing 4 changed files with 142 additions and 45 deletions.
diff --git a/evm/src/executor/inspector/cheatcodes/expect.rs b/evm/src/executor/inspector/cheatcodes/expect.rs
@@ -213,6 +213,8 @@ pub struct ExpectedCallData {
     pub count: u64,
     /// The type of call
     pub call_type: ExpectedCallType,
+    /// The depth at which this call must be checked
+    pub depth: u64,
 }
 
 #[derive(Clone, Debug, Default, PartialEq, Eq)]
@@ -289,6 +291,7 @@ fn expect_call(
     min_gas: Option<u64>,
     count: u64,
     call_type: ExpectedCallType,
+    depth: u64,
 ) -> Result {
     match call_type {
         ExpectedCallType::Count => {
@@ -300,8 +303,10 @@ fn expect_call(
                 !expecteds.contains_key(&calldata),
                 "Counted expected calls can only bet set once."
             );
-            expecteds
-                .insert(calldata, (ExpectedCallData { value, gas, min_gas, count, call_type }, 0));
+            expecteds.insert(
+                calldata,
+                (ExpectedCallData { value, gas, min_gas, count, call_type, depth }, 0),
+            );
             Ok(Bytes::new())
         }
         ExpectedCallType::NonCount => {
@@ -319,7 +324,7 @@ fn expect_call(
                 // If it does not exist, then create it.
                 expecteds.insert(
                     calldata,
-                    (ExpectedCallData { value, gas, min_gas, count, call_type }, 0),
+                    (ExpectedCallData { value, gas, min_gas, count, call_type, depth }, 0),
                 );
             }
             Ok(Bytes::new())
@@ -384,6 +389,7 @@ pub fn apply<DB: DatabaseExt>(
             None,
             1,
             ExpectedCallType::NonCount,
+            data.journaled_state.depth(),
         ),
         HEVMCalls::ExpectCall1(inner) => expect_call(
             state,
@@ -394,6 +400,7 @@ pub fn apply<DB: DatabaseExt>(
             None,
             inner.2,
             ExpectedCallType::Count,
+            data.journaled_state.depth(),
         ),
         HEVMCalls::ExpectCall2(inner) => expect_call(
             state,
@@ -404,6 +411,7 @@ pub fn apply<DB: DatabaseExt>(
             None,
             1,
             ExpectedCallType::NonCount,
+            data.journaled_state.depth(),
         ),
         HEVMCalls::ExpectCall3(inner) => expect_call(
             state,
@@ -414,6 +422,7 @@ pub fn apply<DB: DatabaseExt>(
             None,
             inner.3,
             ExpectedCallType::Count,
+            data.journaled_state.depth(),
         ),
         HEVMCalls::ExpectCall4(inner) => {
             let value = inner.1;
@@ -430,6 +439,7 @@ pub fn apply<DB: DatabaseExt>(
                 None,
                 1,
                 ExpectedCallType::NonCount,
+                data.journaled_state.depth(),
             )
         }
         HEVMCalls::ExpectCall5(inner) => {
@@ -447,6 +457,7 @@ pub fn apply<DB: DatabaseExt>(
                 None,
                 inner.4,
                 ExpectedCallType::Count,
+                data.journaled_state.depth(),
             )
         }
         HEVMCalls::ExpectCallMinGas0(inner) => {
@@ -464,6 +475,7 @@ pub fn apply<DB: DatabaseExt>(
                 Some(inner.2 + positive_value_cost_stipend),
                 1,
                 ExpectedCallType::NonCount,
+                data.journaled_state.depth(),
             )
         }
         HEVMCalls::ExpectCallMinGas1(inner) => {
@@ -481,6 +493,7 @@ pub fn apply<DB: DatabaseExt>(
                 Some(inner.2 + positive_value_cost_stipend),
                 inner.4,
                 ExpectedCallType::Count,
+                data.journaled_state.depth(),
             )
         }
         HEVMCalls::MockCall0(inner) => {

diff --git a/evm/src/executor/inspector/cheatcodes/mod.rs b/evm/src/executor/inspector/cheatcodes/mod.rs
@@ -591,7 +591,10 @@ where
                         // The gas matches, if provided
                         expected.gas.map_or(true, |gas| gas == call.gas_limit) &&
                         // The minimum gas matches, if provided
-                        expected.min_gas.map_or(true, |min_gas| min_gas <= call.gas_limit)
+                        expected.min_gas.map_or(true, |min_gas| min_gas <= call.gas_limit) &&
+                        // The expected depth is smaller than the actual depth,
+                        // which means we're in the subcalls of the call were we expect to find the matches.
+                        expected.depth < data.journaled_state.depth()
                     {
                         *actual_count += 1;
                     }
@@ -818,14 +821,18 @@ where
             }
         }
 
-        // If the depth is 0, then this is the root call terminating
-        if data.journaled_state.depth() == 0 {
-            // Match expected calls
-            for (address, calldatas) in &self.expected_calls {
-                // Loop over each address, and for each address, loop over each calldata it expects.
-                for (calldata, (expected, actual_count)) in calldatas {
-                    // Grab the values we expect to see
-                    let ExpectedCallData { gas, min_gas, value, count, call_type } = expected;
+        // Match expected calls
+        for (address, calldatas) in &self.expected_calls {
+            // Loop over each address, and for each address, loop over each calldata it expects.
+            for (calldata, (expected, actual_count)) in calldatas {
+                // Grab the values we expect to see
+                let ExpectedCallData { gas, min_gas, value, count, call_type, depth } = expected;
+                // Only check calls in the corresponding depth,
+                // or if the expected depth is higher than the current depth. This is correct, as
+                // the expected depth can only be bigger than the current depth if
+                // we're either terminating the root call (the test itself), or exiting the intended
+                // call that contained the calls we expected to see.
+                if depth >= &data.journaled_state.depth() {
                     let calldata = Bytes::from(calldata.clone());
 
                     // We must match differently depending on the type of call we expect.
@@ -884,7 +891,10 @@ where
                     }
                 }
             }
+        }
 
+        // If the depth is 0, then this is the root call terminating
+        if data.journaled_state.depth() == 0 {
             // Check if we have any leftover expected emits
             // First, if any emits were found at the root call, then we its ok and we remove them.
             self.expected_emits.retain(|expected| !expected.found);