diff --git a/evm/src/executor/inspector/cheatcodes/env.rs b/evm/src/executor/inspector/cheatcodes/env.rs index 058e4a63d18b..7c70f2910cbf 100644 --- a/evm/src/executor/inspector/cheatcodes/env.rs +++ b/evm/src/executor/inspector/cheatcodes/env.rs @@ -52,6 +52,33 @@ pub struct Prank { pub depth: u64, /// Whether the prank stops by itself after the next call pub single_call: bool, + /// Whether the prank has been used yet (false if unused) + pub used: bool, +} + +impl Prank { + pub fn new( + prank_caller: Address, + prank_origin: Address, + new_caller: Address, + new_origin: Option
, + depth: u64, + single_call: bool, + ) -> Prank { + Prank { prank_caller, prank_origin, new_caller, new_origin, depth, single_call, used: false } + } + + /// Apply the prank by setting `used` to true iff it is false + pub fn apply(&self) -> Self { + if self.used { + self.clone() + } else { + Prank { + used: true, + ..self.clone() + } + } + } } /// Sets up broadcasting from a script using `origin` as the sender @@ -118,7 +145,13 @@ fn prank( depth: u64, single_call: bool, ) -> Result { - let prank = Prank { prank_caller, prank_origin, new_caller, new_origin, depth, single_call }; + let prank = Prank::new(prank_caller, prank_origin, new_caller, new_origin, depth, single_call); + + if let Some(Prank { used, ..}) = state.prank { + if !used { + return Err("You cannot overwrite `prank` until it is applied at least once".encode().into()); + } + } if state.broadcast.is_some() { return Err("You cannot `prank` for a broadcasted transaction. Pass the desired tx.origin into the broadcast cheatcode call".encode().into()); diff --git a/evm/src/executor/inspector/cheatcodes/mod.rs b/evm/src/executor/inspector/cheatcodes/mod.rs index 7d1d7d3158d0..ffbf0f2b519c 100644 --- a/evm/src/executor/inspector/cheatcodes/mod.rs +++ b/evm/src/executor/inspector/cheatcodes/mod.rs @@ -583,15 +583,22 @@ where if data.journaled_state.depth() >= prank.depth && call.context.caller == prank.prank_caller { + let mut prank_applied = false; // At the target depth we set `msg.sender` if data.journaled_state.depth() == prank.depth { call.context.caller = prank.new_caller; call.transfer.source = prank.new_caller; + prank_applied = true; } // At the target depth, or deeper, we set `tx.origin` if let Some(new_origin) = prank.new_origin { data.env.tx.caller = new_origin; + prank_applied = true; + } + // If prank applied, then update the prank state accordingly + if prank_applied { + self.prank = Some(prank.apply()); } } }