feat(cheatcodes): add delegatecall to pranking

crates/cheatcodes/spec/src/vm.rs

@@ -491,6 +491,21 @@ 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;
+
+    #[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/src/evm/prank.rs

@@ -1,6 +1,19 @@
 use crate::{Cheatcode, Cheatcodes, CheatsCtxt, DatabaseExt, Result, Vm::*};
 use alloy_primitives::Address;
 
+// Update prank so that you can use it for delegatecalling from a test contract, but throw an error
+// if the address passed to vm.prank(addr) before a delegatecall has no code (to ensure you can't
+// delegatecall from an EOA). Is there any related change to how pranking tx.origin is impacted? I
+// don't think anything around tx.origin needs to change, but just making sure
+
+// Cheat codes work by capture the transaction context and manipulating the environment based on
+// cheatcodes
+
+// Notes:
+// https://github.com/EdwardJES/foundry/blob/cb109b1699f82d009574d13aa59f1585a3fbfdb2/crates/cheatcodes/src/inspector.rs#L723
+// Call with executor: here we could intercept the the call with delegate call
+// Possibly intercept here https://github.com/EdwardJES/foundry/blob/cb109b1699f82d009574d13aa59f1585a3fbfdb2/crates/cheatcodes/src/inspector.rs#L834
+
 /// Prank information.
 #[derive(Clone, Debug, Default)]
 pub struct Prank {
@@ -16,6 +29,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 +44,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 +72,56 @@ impl Prank {
 impl Cheatcode for prank_0Call {
     fn apply_stateful<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> Result {
         let Self { msgSender } = self;
-        prank(ccx, msgSender, None, true)
+        prank(ccx, msgSender, None, true, false)
     }
 }
 
 impl Cheatcode for startPrank_0Call {
     fn apply_stateful<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> Result {
         let Self { msgSender } = self;
-        prank(ccx, msgSender, None, false)
+        prank(ccx, msgSender, None, false, false)
     }
 }
 
 impl Cheatcode for prank_1Call {
     fn apply_stateful<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> 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<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> 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<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> Result {
+        let Self { msgSender, delegateCall } = self;
+        prank(ccx, msgSender, None, true, *delegateCall)
+    }
+}
+
+impl Cheatcode for startPrank_2Call {
+    fn apply_stateful<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> Result {
+        let Self { msgSender, delegateCall } = self;
+        prank(ccx, msgSender, None, false, *delegateCall)
+    }
+}
+
+impl Cheatcode for prank_3Call {
+    fn apply_stateful<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> Result {
+        let Self { msgSender, txOrigin, delegateCall } = self;
+        prank(ccx, msgSender, Some(txOrigin), true, *delegateCall)
+    }
+}
+
+impl Cheatcode for startPrank_3Call {
+    fn apply_stateful<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> Result {
+        let Self { msgSender, txOrigin, delegateCall } = self;
+        prank(ccx, msgSender, Some(txOrigin), false, *delegateCall)
     }
 }
 
@@ -85,6 +138,7 @@ fn prank<DB: DatabaseExt>(
     new_caller: &Address,
     new_origin: Option<&Address>,
     single_call: bool,
+    delegate_call: bool,
 ) -> Result {
     let prank = Prank::new(
         ccx.caller,
@@ -93,8 +147,15 @@ fn prank<DB: DatabaseExt>(
         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/inspector.rs

@@ -37,7 +37,7 @@ use itertools::Itertools;
 use rand::{rngs::StdRng, Rng, SeedableRng};
 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,
     },
@@ -833,6 +833,19 @@ impl Cheatcodes {
 
         // Apply our prank
         if let Some(prank) = &self.prank {
+            // Apply delegate call. call.caller will not equal prank.prank_caller
+            if let CallScheme::DelegateCall = 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;
 
@@ -946,7 +959,7 @@ impl Cheatcodes {
                 initialized = false;
                 old_balance = U256::ZERO;
             }
-            let kind = match call.scheme {
+            let kind: Vm::AccountAccessKind = match call.scheme {
                 CallScheme::Call => crate::Vm::AccountAccessKind::Call,
                 CallScheme::CallCode => crate::Vm::AccountAccessKind::CallCode,
                 CallScheme::DelegateCall => crate::Vm::AccountAccessKind::DelegateCall,