foundry_cheatcodes/evm/
prank.rs

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

foundry_cheatcodes/evm/
prank.rs
