feat(cheatcodes): add delegatecall to pranking (foundry-rs#8863)

* begin api and rough comments

* impl cheatcode

* add check for eoa

* fix eoa check on each prank call

* add to assets

* prank compiling

* delegate call working, storage not upating

* delegate call working, some tidy up

* add prank2 calls

* impl remaining tests

* formatting

* forge fmt

* add pranks to cheatcodes.json

* run cargo cheats

* If verbosity level is 1 or higher, it shows dirty files.

* Fix, add EOA prank test

* Revert "If verbosity level is 1 or higher, it shows dirty files."

This reverts commit d03ac1d.

* Fix test

* apply on extdelegatecall

---------

Co-authored-by: mgiagante <[email protected]>
Co-authored-by: grandizzy <[email protected]>

Author: 3 people

crates/cheatcodes/assets/cheatcodes.json

@@ -600,6 +600,22 @@ interface Vm {
     #[cheatcode(group = Evm, safety = Unsafe)]
     function startPrank(address msgSender, address txOrigin) external;
 
+    /// Sets the *next* delegate call's `msg.sender` to be the input address.
+    #[cheatcode(group = Evm, safety = Unsafe)]
+    function prank(address msgSender, bool delegateCall) external;
+
+    /// Sets all subsequent delegate calls' `msg.sender` to be the input address until `stopPrank` is called.
+    #[cheatcode(group = Evm, safety = Unsafe)]
+    function startPrank(address msgSender, bool delegateCall) external;
+
+    /// Sets the *next* delegate call's `msg.sender` to be the input address, and the `tx.origin` to be the second input.
+    #[cheatcode(group = Evm, safety = Unsafe)]
+    function prank(address msgSender, address txOrigin, bool delegateCall) external;
+
+    /// Sets all subsequent delegate calls' `msg.sender` to be the input address until `stopPrank` is called, and the `tx.origin` to be the second input.
+    #[cheatcode(group = Evm, safety = Unsafe)]
+    function startPrank(address msgSender, address txOrigin, bool delegateCall) external;
+
     /// Resets subsequent calls' `msg.sender` to be `address(this)`.
     #[cheatcode(group = Evm, safety = Unsafe)]
     function stopPrank() external;

crates/cheatcodes/spec/src/vm.rs

@@ -16,6 +16,8 @@ pub struct Prank {
     pub depth: u64,
     /// Whether the prank stops by itself after the next call
     pub single_call: bool,
+    /// Whether the prank should be be applied to delegate call
+    pub delegate_call: bool,
     /// Whether the prank has been used yet (false if unused)
     pub used: bool,
 }
@@ -29,8 +31,18 @@ impl Prank {
         new_origin: Option<Address>,
         depth: u64,
         single_call: bool,
+        delegate_call: bool,
     ) -> Self {
-        Self { prank_caller, prank_origin, new_caller, new_origin, depth, single_call, used: false }
+        Self {
+            prank_caller,
+            prank_origin,
+            new_caller,
+            new_origin,
+            depth,
+            single_call,
+            delegate_call,
+            used: false,
+        }
     }
 
     /// Apply the prank by setting `used` to true iff it is false
@@ -47,28 +59,56 @@ impl Prank {
 impl Cheatcode for prank_0Call {
     fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
         let Self { msgSender } = self;
-        prank(ccx, msgSender, None, true)
+        prank(ccx, msgSender, None, true, false)
     }
 }
 
 impl Cheatcode for startPrank_0Call {
     fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
         let Self { msgSender } = self;
-        prank(ccx, msgSender, None, false)
+        prank(ccx, msgSender, None, false, false)
     }
 }
 
 impl Cheatcode for prank_1Call {
     fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
         let Self { msgSender, txOrigin } = self;
-        prank(ccx, msgSender, Some(txOrigin), true)
+        prank(ccx, msgSender, Some(txOrigin), true, false)
     }
 }
 
 impl Cheatcode for startPrank_1Call {
     fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
         let Self { msgSender, txOrigin } = self;
-        prank(ccx, msgSender, Some(txOrigin), false)
+        prank(ccx, msgSender, Some(txOrigin), false, false)
+    }
+}
+
+impl Cheatcode for prank_2Call {
+    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
+        let Self { msgSender, delegateCall } = self;
+        prank(ccx, msgSender, None, true, *delegateCall)
+    }
+}
+
+impl Cheatcode for startPrank_2Call {
+    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
+        let Self { msgSender, delegateCall } = self;
+        prank(ccx, msgSender, None, false, *delegateCall)
+    }
+}
+
+impl Cheatcode for prank_3Call {
+    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
+        let Self { msgSender, txOrigin, delegateCall } = self;
+        prank(ccx, msgSender, Some(txOrigin), true, *delegateCall)
+    }
+}
+
+impl Cheatcode for startPrank_3Call {
+    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
+        let Self { msgSender, txOrigin, delegateCall } = self;
+        prank(ccx, msgSender, Some(txOrigin), false, *delegateCall)
     }
 }
 
@@ -85,6 +125,7 @@ fn prank(
     new_caller: &Address,
     new_origin: Option<&Address>,
     single_call: bool,
+    delegate_call: bool,
 ) -> Result {
     let prank = Prank::new(
         ccx.caller,
@@ -93,8 +134,15 @@ fn prank(
         new_origin.copied(),
         ccx.ecx.journaled_state.depth(),
         single_call,
+        delegate_call,
     );
 
+    // Ensure that code exists at `msg.sender` if delegate calling.
+    if delegate_call {
+        let code = ccx.code(*new_caller)?;
+        ensure!(!code.is_empty(), "cannot `prank` delegate call from an EOA");
+    }
+
     if let Some(Prank { used, single_call: current_single_call, .. }) = ccx.state.prank {
         ensure!(used, "cannot overwrite a prank until it is applied at least once");
         // This case can only fail if the user calls `vm.startPrank` and then `vm.prank` later on.

crates/cheatcodes/src/evm/prank.rs

@@ -42,7 +42,7 @@ use proptest::test_runner::{RngAlgorithm, TestRng, TestRunner};
 use rand::Rng;
 use revm::{
     interpreter::{
-        opcode as op, CallInputs, CallOutcome, CallScheme, CreateInputs, CreateOutcome,
+        opcode as op, CallInputs, CallOutcome, CallScheme, CallValue, CreateInputs, CreateOutcome,
         EOFCreateInputs, EOFCreateKind, Gas, InstructionResult, Interpreter, InterpreterAction,
         InterpreterResult,
     },
@@ -941,6 +941,19 @@ where {
 
         // Apply our prank
         if let Some(prank) = &self.prank {
+            // Apply delegate call, `call.caller`` will not equal `prank.prank_caller`
+            if let CallScheme::DelegateCall | CallScheme::ExtDelegateCall = call.scheme {
+                if prank.delegate_call {
+                    call.target_address = prank.new_caller;
+                    call.caller = prank.new_caller;
+                    let acc = ecx.journaled_state.account(prank.new_caller);
+                    call.value = CallValue::Apparent(acc.info.balance);
+                    if let Some(new_origin) = prank.new_origin {
+                        ecx.env.tx.caller = new_origin;
+                    }
+                }
+            }
+
             if ecx.journaled_state.depth() >= prank.depth && call.caller == prank.prank_caller {
                 let mut prank_applied = false;