foundry_cheatcodes/evm/
prank.rs

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

foundry_cheatcodes/evm/
prank.rs
