foundry_cheatcodes/evm/
prank.rs

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

foundry_cheatcodes/evm/
prank.rs
