foundry_cheatcodes/evm/
prank.rs

Help
1use crate::{Cheatcode, CheatsCtxt, Result, Vm::*};
2use alloy_primitives::Address;
3use revm::{context::JournalTr, interpreter::Host};
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 {
53            None
54        } else {
55            Some(Self { used: true, ..*self })
56        }
57    }
58}
59
60impl Cheatcode for prank_0Call {
61    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
62        let Self { msgSender } = self;
63        prank(ccx, msgSender, None, true, false)
64    }
65}
66
67impl Cheatcode for startPrank_0Call {
68    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
69        let Self { msgSender } = self;
70        prank(ccx, msgSender, None, false, false)
71    }
72}
73
74impl Cheatcode for prank_1Call {
75    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
76        let Self { msgSender, txOrigin } = self;
77        prank(ccx, msgSender, Some(txOrigin), true, false)
78    }
79}
80
81impl Cheatcode for startPrank_1Call {
82    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
83        let Self { msgSender, txOrigin } = self;
84        prank(ccx, msgSender, Some(txOrigin), false, false)
85    }
86}
87
88impl Cheatcode for prank_2Call {
89    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
90        let Self { msgSender, delegateCall } = self;
91        prank(ccx, msgSender, None, true, *delegateCall)
92    }
93}
94
95impl Cheatcode for startPrank_2Call {
96    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
97        let Self { msgSender, delegateCall } = self;
98        prank(ccx, msgSender, None, false, *delegateCall)
99    }
100}
101
102impl Cheatcode for prank_3Call {
103    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
104        let Self { msgSender, txOrigin, delegateCall } = self;
105        prank(ccx, msgSender, Some(txOrigin), true, *delegateCall)
106    }
107}
108
109impl Cheatcode for startPrank_3Call {
110    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
111        let Self { msgSender, txOrigin, delegateCall } = self;
112        prank(ccx, msgSender, Some(txOrigin), false, *delegateCall)
113    }
114}
115
116impl Cheatcode for stopPrankCall {
117    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
118        let Self {} = self;
119        ccx.state.pranks.remove(&ccx.ecx.journaled_state.depth());
120        Ok(Default::default())
121    }
122}
123
124fn prank(
125    ccx: &mut CheatsCtxt,
126    new_caller: &Address,
127    new_origin: Option<&Address>,
128    single_call: bool,
129    delegate_call: bool,
130) -> Result {
131    // Ensure that code exists at `msg.sender` if delegate calling.
132    if delegate_call {
133        let code = ccx
134            .load_account_code(*new_caller)
135            .ok_or_else(|| eyre::eyre!("cannot `prank` delegate call from an EOA"))?;
136
137        ensure!(!code.data.is_empty(), "cannot `prank` delegate call from an EOA");
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}
foundry_cheatcodes/evm/prank.rs

foundry_cheatcodes/evm/
prank.rs
