foundry_cheatcodes/evm/
prank.rs

1use crate::{Cheatcode, CheatsCtxt, Result, Vm::*, evm::journaled_account};
2use alloy_primitives::Address;
3use revm::context::JournalTr;
4
5/// Prank information.
6#[derive(Clone, Copy, Debug, Default)]
7pub struct Prank {
8    /// Address of the contract that initiated the prank
9    pub prank_caller: Address,
10    /// Address of `tx.origin` when the prank was initiated
11    pub prank_origin: Address,
12    /// The address to assign to `msg.sender`
13    pub new_caller: Address,
14    /// The address to assign to `tx.origin`
15    pub new_origin: Option<Address>,
16    /// The depth at which the prank was called
17    pub depth: usize,
18    /// Whether the prank stops by itself after the next call
19    pub single_call: bool,
20    /// Whether the prank should be applied to delegate call
21    pub delegate_call: bool,
22    /// Whether the prank has been used yet (false if unused)
23    pub used: bool,
24}
25
26impl Prank {
27    /// Create a new prank.
28    pub fn new(
29        prank_caller: Address,
30        prank_origin: Address,
31        new_caller: Address,
32        new_origin: Option<Address>,
33        depth: usize,
34        single_call: bool,
35        delegate_call: bool,
36    ) -> Self {
37        Self {
38            prank_caller,
39            prank_origin,
40            new_caller,
41            new_origin,
42            depth,
43            single_call,
44            delegate_call,
45            used: false,
46        }
47    }
48
49    /// Apply the prank by setting `used` to true if it is false
50    /// Only returns self in the case it is updated (first application)
51    pub fn first_time_applied(&self) -> Option<Self> {
52        if self.used { None } else { Some(Self { used: true, ..*self }) }
53    }
54}
55
56impl Cheatcode for prank_0Call {
57    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
58        let Self { msgSender } = self;
59        prank(ccx, msgSender, None, true, false)
60    }
61}
62
63impl Cheatcode for startPrank_0Call {
64    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
65        let Self { msgSender } = self;
66        prank(ccx, msgSender, None, false, false)
67    }
68}
69
70impl Cheatcode for prank_1Call {
71    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
72        let Self { msgSender, txOrigin } = self;
73        prank(ccx, msgSender, Some(txOrigin), true, false)
74    }
75}
76
77impl Cheatcode for startPrank_1Call {
78    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
79        let Self { msgSender, txOrigin } = self;
80        prank(ccx, msgSender, Some(txOrigin), false, false)
81    }
82}
83
84impl Cheatcode for prank_2Call {
85    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
86        let Self { msgSender, delegateCall } = self;
87        prank(ccx, msgSender, None, true, *delegateCall)
88    }
89}
90
91impl Cheatcode for startPrank_2Call {
92    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
93        let Self { msgSender, delegateCall } = self;
94        prank(ccx, msgSender, None, false, *delegateCall)
95    }
96}
97
98impl Cheatcode for prank_3Call {
99    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
100        let Self { msgSender, txOrigin, delegateCall } = self;
101        prank(ccx, msgSender, Some(txOrigin), true, *delegateCall)
102    }
103}
104
105impl Cheatcode for startPrank_3Call {
106    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
107        let Self { msgSender, txOrigin, delegateCall } = self;
108        prank(ccx, msgSender, Some(txOrigin), false, *delegateCall)
109    }
110}
111
112impl Cheatcode for stopPrankCall {
113    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
114        let Self {} = self;
115        ccx.state.pranks.remove(&ccx.ecx.journaled_state.depth());
116        Ok(Default::default())
117    }
118}
119
120fn prank(
121    ccx: &mut CheatsCtxt,
122    new_caller: &Address,
123    new_origin: Option<&Address>,
124    single_call: bool,
125    delegate_call: bool,
126) -> Result {
127    // Ensure that we load the account of the pranked address and mark it as touched.
128    // This is necessary to ensure that account state changes (such as the account's `nonce`) are
129    // properly tracked.
130    let account = journaled_account(ccx.ecx, *new_caller)?;
131
132    // Ensure that code exists at `msg.sender` if delegate calling.
133    if delegate_call {
134        ensure!(
135            account.info.clone().code.is_some_and(|code| !code.is_empty()),
136            "cannot `prank` delegate call from an EOA"
137        );
138    }
139
140    let depth = ccx.ecx.journaled_state.depth();
141    if let Some(Prank { used, single_call: current_single_call, .. }) = ccx.state.get_prank(depth) {
142        ensure!(used, "cannot overwrite a prank until it is applied at least once");
143        // This case can only fail if the user calls `vm.startPrank` and then `vm.prank` later on.
144        // This should not be possible without first calling `stopPrank`
145        ensure!(
146            single_call == *current_single_call,
147            "cannot override an ongoing prank with a single vm.prank; \
148             use vm.startPrank to override the current prank"
149        );
150    }
151
152    let prank = Prank::new(
153        ccx.caller,
154        ccx.ecx.tx.caller,
155        *new_caller,
156        new_origin.copied(),
157        depth,
158        single_call,
159        delegate_call,
160    );
161
162    ensure!(
163        ccx.state.broadcast.is_none(),
164        "cannot `prank` for a broadcasted transaction; \
165         pass the desired `tx.origin` into the `broadcast` cheatcode call"
166    );
167
168    ccx.state.pranks.insert(prank.depth, prank);
169    Ok(Default::default())
170}