foundry_cheatcodes/evm/
prank.rs

1use crate::{Cheatcode, CheatsCtxt, Result, Vm::*, evm::journaled_account};
2use alloy_primitives::Address;
3use foundry_evm_core::backend::DatabaseExt;
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 Cheatcode for prank_0Call {
61    fn apply_stateful<CTX: ContextTr<Journal: JournalExt, Db: DatabaseExt>>(
62        &self,
63        ccx: &mut CheatsCtxt<'_, CTX>,
64    ) -> Result {
65        let Self { msgSender } = self;
66        prank(ccx, msgSender, None, true, false)
67    }
68}
69
70impl Cheatcode for startPrank_0Call {
71    fn apply_stateful<CTX: ContextTr<Journal: JournalExt, Db: DatabaseExt>>(
72        &self,
73        ccx: &mut CheatsCtxt<'_, CTX>,
74    ) -> Result {
75        let Self { msgSender } = self;
76        prank(ccx, msgSender, None, false, false)
77    }
78}
79
80impl Cheatcode for prank_1Call {
81    fn apply_stateful<CTX: ContextTr<Journal: JournalExt, Db: DatabaseExt>>(
82        &self,
83        ccx: &mut CheatsCtxt<'_, CTX>,
84    ) -> Result {
85        let Self { msgSender, txOrigin } = self;
86        prank(ccx, msgSender, Some(txOrigin), true, false)
87    }
88}
89
90impl Cheatcode for startPrank_1Call {
91    fn apply_stateful<CTX: ContextTr<Journal: JournalExt, Db: DatabaseExt>>(
92        &self,
93        ccx: &mut CheatsCtxt<'_, CTX>,
94    ) -> Result {
95        let Self { msgSender, txOrigin } = self;
96        prank(ccx, msgSender, Some(txOrigin), false, false)
97    }
98}
99
100impl Cheatcode for prank_2Call {
101    fn apply_stateful<CTX: ContextTr<Journal: JournalExt, Db: DatabaseExt>>(
102        &self,
103        ccx: &mut CheatsCtxt<'_, CTX>,
104    ) -> Result {
105        let Self { msgSender, delegateCall } = self;
106        prank(ccx, msgSender, None, true, *delegateCall)
107    }
108}
109
110impl Cheatcode for startPrank_2Call {
111    fn apply_stateful<CTX: ContextTr<Journal: JournalExt, Db: DatabaseExt>>(
112        &self,
113        ccx: &mut CheatsCtxt<'_, CTX>,
114    ) -> Result {
115        let Self { msgSender, delegateCall } = self;
116        prank(ccx, msgSender, None, false, *delegateCall)
117    }
118}
119
120impl Cheatcode for prank_3Call {
121    fn apply_stateful<CTX: ContextTr<Journal: JournalExt, Db: DatabaseExt>>(
122        &self,
123        ccx: &mut CheatsCtxt<'_, CTX>,
124    ) -> Result {
125        let Self { msgSender, txOrigin, delegateCall } = self;
126        prank(ccx, msgSender, Some(txOrigin), true, *delegateCall)
127    }
128}
129
130impl Cheatcode for startPrank_3Call {
131    fn apply_stateful<CTX: ContextTr<Journal: JournalExt, Db: DatabaseExt>>(
132        &self,
133        ccx: &mut CheatsCtxt<'_, CTX>,
134    ) -> Result {
135        let Self { msgSender, txOrigin, delegateCall } = self;
136        prank(ccx, msgSender, Some(txOrigin), false, *delegateCall)
137    }
138}
139
140impl Cheatcode for stopPrankCall {
141    fn apply_stateful<CTX: ContextTr<Journal: JournalExt, Db: DatabaseExt>>(
142        &self,
143        ccx: &mut CheatsCtxt<'_, CTX>,
144    ) -> Result {
145        let Self {} = self;
146        ccx.state.pranks.remove(&ccx.ecx.journal().depth());
147        Ok(Default::default())
148    }
149}
150
151fn prank<CTX: ContextTr<Journal: JournalExt, Db: DatabaseExt>>(
152    ccx: &mut CheatsCtxt<'_, CTX>,
153    new_caller: &Address,
154    new_origin: Option<&Address>,
155    single_call: bool,
156    delegate_call: bool,
157) -> Result {
158    // Ensure that we load the account of the pranked address and mark it as touched.
159    // This is necessary to ensure that account state changes (such as the account's `nonce`) are
160    // properly tracked.
161    let account = journaled_account(ccx.ecx, *new_caller)?;
162
163    // Ensure that code exists at `msg.sender` if delegate calling.
164    if delegate_call {
165        ensure!(
166            account.info.code.as_ref().is_some_and(|code| !code.is_empty()),
167            "cannot `prank` delegate call from an EOA"
168        );
169    }
170
171    let depth = ccx.ecx.journal().depth();
172    if let Some(Prank { used, single_call: current_single_call, .. }) = ccx.state.get_prank(depth) {
173        ensure!(used, "cannot overwrite a prank until it is applied at least once");
174        // This case can only fail if the user calls `vm.startPrank` and then `vm.prank` later on.
175        // This should not be possible without first calling `stopPrank`
176        ensure!(
177            single_call == *current_single_call,
178            "cannot override an ongoing prank with a single vm.prank; \
179             use vm.startPrank to override the current prank"
180        );
181    }
182
183    let prank = Prank::new(
184        ccx.caller,
185        ccx.ecx.tx().caller(),
186        *new_caller,
187        new_origin.copied(),
188        depth,
189        single_call,
190        delegate_call,
191    );
192
193    ensure!(
194        ccx.state.broadcast.is_none(),
195        "cannot `prank` for a broadcasted transaction; \
196         pass the desired `tx.origin` into the `broadcast` cheatcode call"
197    );
198
199    ccx.state.pranks.insert(prank.depth, prank);
200    Ok(Default::default())
201}
foundry_cheatcodes/evm/prank.rs

foundry_cheatcodes/evm/
prank.rs
