From 2f8977921d7f2138c23d67b26d4a3dc1dc390be9 Mon Sep 17 00:00:00 2001 From: 4meta5 Date: Tue, 25 Apr 2023 22:53:11 -0400 Subject: [PATCH 01/14] init start change prank --- evm/src/executor/abi/mod.rs | 1 + evm/src/executor/inspector/cheatcodes/env.rs | 16 ++++++++++++++++ forge/README.md | 4 ++++ testdata/cheats/Cheats.sol | 3 +++ 4 files changed, 24 insertions(+) diff --git a/evm/src/executor/abi/mod.rs b/evm/src/executor/abi/mod.rs index 3ba07ce1c0e5e..7217d47845a13 100644 --- a/evm/src/executor/abi/mod.rs +++ b/evm/src/executor/abi/mod.rs @@ -76,6 +76,7 @@ abigen!( prank(address,address) startPrank(address) startPrank(address,address) + startChangePrank(address,address) stopPrank() deal(address,uint256) diff --git a/evm/src/executor/inspector/cheatcodes/env.rs b/evm/src/executor/inspector/cheatcodes/env.rs index 3dd35f57bf6c3..a9886463c6485 100644 --- a/evm/src/executor/inspector/cheatcodes/env.rs +++ b/evm/src/executor/inspector/cheatcodes/env.rs @@ -329,6 +329,22 @@ pub fn apply( data.journaled_state.depth(), false, )?, + HEVMCalls::StartChangePrank(inner) => { + if state.prank.is_some() { + // stop existing prank + state.prank = None; + } + // start new prank + prank( + state, + caller, + data.env.tx.caller, + inner.0, + Some(inner.1), + data.journaled_state.depth(), + false, + )? + } HEVMCalls::StopPrank(_) => { state.prank = None; Bytes::new() diff --git a/forge/README.md b/forge/README.md index eb2f068800f16..9e8f5be2594a8 100644 --- a/forge/README.md +++ b/forge/README.md @@ -161,6 +161,8 @@ which implements the following methods: - `function startPrank(address sender, address origin)`: Performs smart contract calls as another address, while also setting `tx.origin`. The account impersonation lasts until the end of the transaction, or until `stopPrank` is called. +- `function startChangePrank(address sender, address origin)`: If `startPrank` not yet called then calls `startPrank`. If `startPrank` has been called, then calls `stopPrank` before calling `startPrank` with new prank arguments. + - `function stopPrank()`: Stop calling smart contracts with the address set at `startPrank` - `function expectRevert( expectedError)`: @@ -284,6 +286,8 @@ interface Hevm { function prank(address,address) external; // Sets all subsequent calls' msg.sender to be the input address until `stopPrank` is called, and the tx.origin to be the second input function startPrank(address,address) external; + // Optimistically sets all subsequent calls' msg.sender to be the input address until `stopPrank` is called, and the tx.origin to be the second input + function startChangePrank(address,address) external; // Resets subsequent calls' msg.sender to be `address(this)` function stopPrank() external; // Sets an address' balance, (who, newBalance) diff --git a/testdata/cheats/Cheats.sol b/testdata/cheats/Cheats.sol index 1983784036580..9e56fd3f81098 100644 --- a/testdata/cheats/Cheats.sol +++ b/testdata/cheats/Cheats.sol @@ -140,6 +140,9 @@ interface Cheats { // Sets all subsequent calls' msg.sender to be the input address until `stopPrank` is called, and the tx.origin to be the second input function startPrank(address, address) external; + // Optimistically sets all subsequent calls' msg.sender to be the input address until `stopPrank` is called, and the tx.origin to be the second input + function startChangePrank(address, address) external; + // Resets subsequent calls' msg.sender to be `address(this)` function stopPrank() external; From a19209a67dd9de2f22532da82d26aa6b3594cad0 Mon Sep 17 00:00:00 2001 From: 4meta5 Date: Tue, 25 Apr 2023 23:36:30 -0400 Subject: [PATCH 02/14] testChangePrank --- testdata/cheats/Prank.t.sol | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/testdata/cheats/Prank.t.sol b/testdata/cheats/Prank.t.sol index 81e24ed5c944f..09b17aa8a34cb 100644 --- a/testdata/cheats/Prank.t.sol +++ b/testdata/cheats/Prank.t.sol @@ -102,6 +102,34 @@ contract PrankTest is DSTest { ); } + function testChangePrank(address sender0, address origin0, address sender1, address origin1) public { + address oldOrigin = tx.origin; + + // Perform first prank + Victim victim = new Victim(); + cheats.prank(sender0, origin0); + victim.assertCallerAndOrigin( + sender0, "msg.sender was not set during prank", origin0, "tx.origin was not set during prank" + ); + + // Change to second prank + cheats.startChangePrank(sender1, origin1); + victim.assertCallerAndOrigin( + sender1, "msg.sender was not set during prank", origin1, "tx.origin was not set during prank" + ); + + // Stop prank + cheats.stopPrank(); + + // Ensure we cleaned up correctly + victim.assertCallerAndOrigin( + address(this), + "msg.sender was not cleaned up in nested prank", + oldOrigin, + "tx.origin was not cleaned up in nested prank" + ); + } + function testPrankOrigin(address sender, address origin) public { address oldOrigin = tx.origin; From 97c99f18c9a06cc7d05b23f77617527e80d3e36c Mon Sep 17 00:00:00 2001 From: 4meta5 Date: Thu, 27 Apr 2023 17:06:26 -0400 Subject: [PATCH 03/14] revert startChangePrank and change startPrank to overwrite existing prank instead of erroring as per review suggestion --- evm/src/executor/abi/mod.rs | 1 - evm/src/executor/inspector/cheatcodes/env.rs | 20 -------------- forge/README.md | 4 --- testdata/cheats/Cheats.sol | 3 --- testdata/cheats/Prank.t.sol | 28 -------------------- 5 files changed, 56 deletions(-) diff --git a/evm/src/executor/abi/mod.rs b/evm/src/executor/abi/mod.rs index 7217d47845a13..3ba07ce1c0e5e 100644 --- a/evm/src/executor/abi/mod.rs +++ b/evm/src/executor/abi/mod.rs @@ -76,7 +76,6 @@ abigen!( prank(address,address) startPrank(address) startPrank(address,address) - startChangePrank(address,address) stopPrank() deal(address,uint256) diff --git a/evm/src/executor/inspector/cheatcodes/env.rs b/evm/src/executor/inspector/cheatcodes/env.rs index a9886463c6485..058e4a63d18b4 100644 --- a/evm/src/executor/inspector/cheatcodes/env.rs +++ b/evm/src/executor/inspector/cheatcodes/env.rs @@ -120,10 +120,6 @@ fn prank( ) -> Result { let prank = Prank { prank_caller, prank_origin, new_caller, new_origin, depth, single_call }; - if state.prank.is_some() { - return Err("You have an active prank already.".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()); } @@ -329,22 +325,6 @@ pub fn apply( data.journaled_state.depth(), false, )?, - HEVMCalls::StartChangePrank(inner) => { - if state.prank.is_some() { - // stop existing prank - state.prank = None; - } - // start new prank - prank( - state, - caller, - data.env.tx.caller, - inner.0, - Some(inner.1), - data.journaled_state.depth(), - false, - )? - } HEVMCalls::StopPrank(_) => { state.prank = None; Bytes::new() diff --git a/forge/README.md b/forge/README.md index 9e8f5be2594a8..eb2f068800f16 100644 --- a/forge/README.md +++ b/forge/README.md @@ -161,8 +161,6 @@ which implements the following methods: - `function startPrank(address sender, address origin)`: Performs smart contract calls as another address, while also setting `tx.origin`. The account impersonation lasts until the end of the transaction, or until `stopPrank` is called. -- `function startChangePrank(address sender, address origin)`: If `startPrank` not yet called then calls `startPrank`. If `startPrank` has been called, then calls `stopPrank` before calling `startPrank` with new prank arguments. - - `function stopPrank()`: Stop calling smart contracts with the address set at `startPrank` - `function expectRevert( expectedError)`: @@ -286,8 +284,6 @@ interface Hevm { function prank(address,address) external; // Sets all subsequent calls' msg.sender to be the input address until `stopPrank` is called, and the tx.origin to be the second input function startPrank(address,address) external; - // Optimistically sets all subsequent calls' msg.sender to be the input address until `stopPrank` is called, and the tx.origin to be the second input - function startChangePrank(address,address) external; // Resets subsequent calls' msg.sender to be `address(this)` function stopPrank() external; // Sets an address' balance, (who, newBalance) diff --git a/testdata/cheats/Cheats.sol b/testdata/cheats/Cheats.sol index 9e56fd3f81098..1983784036580 100644 --- a/testdata/cheats/Cheats.sol +++ b/testdata/cheats/Cheats.sol @@ -140,9 +140,6 @@ interface Cheats { // Sets all subsequent calls' msg.sender to be the input address until `stopPrank` is called, and the tx.origin to be the second input function startPrank(address, address) external; - // Optimistically sets all subsequent calls' msg.sender to be the input address until `stopPrank` is called, and the tx.origin to be the second input - function startChangePrank(address, address) external; - // Resets subsequent calls' msg.sender to be `address(this)` function stopPrank() external; diff --git a/testdata/cheats/Prank.t.sol b/testdata/cheats/Prank.t.sol index 09b17aa8a34cb..81e24ed5c944f 100644 --- a/testdata/cheats/Prank.t.sol +++ b/testdata/cheats/Prank.t.sol @@ -102,34 +102,6 @@ contract PrankTest is DSTest { ); } - function testChangePrank(address sender0, address origin0, address sender1, address origin1) public { - address oldOrigin = tx.origin; - - // Perform first prank - Victim victim = new Victim(); - cheats.prank(sender0, origin0); - victim.assertCallerAndOrigin( - sender0, "msg.sender was not set during prank", origin0, "tx.origin was not set during prank" - ); - - // Change to second prank - cheats.startChangePrank(sender1, origin1); - victim.assertCallerAndOrigin( - sender1, "msg.sender was not set during prank", origin1, "tx.origin was not set during prank" - ); - - // Stop prank - cheats.stopPrank(); - - // Ensure we cleaned up correctly - victim.assertCallerAndOrigin( - address(this), - "msg.sender was not cleaned up in nested prank", - oldOrigin, - "tx.origin was not cleaned up in nested prank" - ); - } - function testPrankOrigin(address sender, address origin) public { address oldOrigin = tx.origin; From 4efb00bbf66b5f8ca567266d2fa711dbfc1acd34 Mon Sep 17 00:00:00 2001 From: 4meta5 Date: Thu, 27 Apr 2023 22:26:47 -0400 Subject: [PATCH 04/14] add tests for prank0 after prank1 and prank1 after prank0 --- testdata/cheats/Prank.t.sol | 52 +++++++++++++++++++++++++++++++++++++ 1 file changed, 52 insertions(+) diff --git a/testdata/cheats/Prank.t.sol b/testdata/cheats/Prank.t.sol index 81e24ed5c944f..ab4469106f326 100644 --- a/testdata/cheats/Prank.t.sol +++ b/testdata/cheats/Prank.t.sol @@ -118,6 +118,58 @@ contract PrankTest is DSTest { ); } + function testPrank1AfterPrank0(address sender, address origin) public { + // Perform the prank + address oldOrigin = tx.origin; + Victim victim = new Victim(); + cheats.prank(sender); + victim.assertCallerAndOrigin( + sender, "msg.sender was not set during prank", oldOrigin, "tx.origin was not set during prank" + ); + + // Ensure we cleaned up correctly + victim.assertCallerAndOrigin( + address(this), "msg.sender was not cleaned up", oldOrigin, "tx.origin invariant failed" + ); + + // Overwrite the prank + cheats.prank(sender, origin); + victim.assertCallerAndOrigin( + sender, "msg.sender was not set during prank", origin, "tx.origin invariant failed" + ); + + // Ensure we cleaned up correctly + victim.assertCallerAndOrigin( + address(this), "msg.sender was not cleaned up", oldOrigin, "tx.origin invariant failed" + ); + } + + function testPrank0AfterPrank1(address sender, address origin) public { + // Perform the prank + address oldOrigin = tx.origin; + Victim victim = new Victim(); + cheats.prank(sender, origin); + victim.assertCallerAndOrigin( + sender, "msg.sender was not set during prank", origin, "tx.origin was not set during prank" + ); + + // Ensure we cleaned up correctly + victim.assertCallerAndOrigin( + address(this), "msg.sender was not cleaned up", oldOrigin, "tx.origin invariant failed" + ); + + // Overwrite the prank + cheats.prank(sender); + victim.assertCallerAndOrigin( + sender, "msg.sender was not set during prank", oldOrigin, "tx.origin invariant failed" + ); + + // Ensure we cleaned up correctly + victim.assertCallerAndOrigin( + address(this), "msg.sender was not cleaned up", oldOrigin, "tx.origin invariant failed" + ); + } + function testPrankConstructorSender(address sender) public { cheats.prank(sender); ConstructorVictim victim = new ConstructorVictim( From 4b476c14fa9292c74dffeeb57505a2ef1cfad1fa Mon Sep 17 00:00:00 2001 From: 4meta5 Date: Fri, 28 Apr 2023 15:27:38 -0400 Subject: [PATCH 05/14] fmt --- testdata/cheats/Prank.t.sol | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/testdata/cheats/Prank.t.sol b/testdata/cheats/Prank.t.sol index ab4469106f326..914729d7beaed 100644 --- a/testdata/cheats/Prank.t.sol +++ b/testdata/cheats/Prank.t.sol @@ -143,7 +143,7 @@ contract PrankTest is DSTest { address(this), "msg.sender was not cleaned up", oldOrigin, "tx.origin invariant failed" ); } - + function testPrank0AfterPrank1(address sender, address origin) public { // Perform the prank address oldOrigin = tx.origin; From 15ac502cca25900a9cac63cbcdccd71fd3c489bb Mon Sep 17 00:00:00 2001 From: 4meta5 Date: Fri, 28 Apr 2023 16:10:36 -0400 Subject: [PATCH 06/14] add error if prank is not used at least once before overwritten as per suggestion --- evm/src/executor/inspector/cheatcodes/env.rs | 35 +++++++++++++++++++- evm/src/executor/inspector/cheatcodes/mod.rs | 7 ++++ 2 files changed, 41 insertions(+), 1 deletion(-) diff --git a/evm/src/executor/inspector/cheatcodes/env.rs b/evm/src/executor/inspector/cheatcodes/env.rs index 058e4a63d18b4..7c70f2910cbfe 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 7d1d7d3158d09..ffbf0f2b519c6 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()); } } } From 853731b15b118380e4418ff5980b3c218f6040b0 Mon Sep 17 00:00:00 2001 From: 4meta5 Date: Fri, 28 Apr 2023 16:15:42 -0400 Subject: [PATCH 07/14] fmt --- evm/src/executor/inspector/cheatcodes/env.rs | 21 +++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/evm/src/executor/inspector/cheatcodes/env.rs b/evm/src/executor/inspector/cheatcodes/env.rs index 7c70f2910cbfe..79437927a5f5e 100644 --- a/evm/src/executor/inspector/cheatcodes/env.rs +++ b/evm/src/executor/inspector/cheatcodes/env.rs @@ -65,7 +65,15 @@ impl Prank { depth: u64, single_call: bool, ) -> Prank { - Prank { prank_caller, prank_origin, new_caller, new_origin, depth, single_call, used: false } + 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 @@ -73,10 +81,7 @@ impl Prank { if self.used { self.clone() } else { - Prank { - used: true, - ..self.clone() - } + Prank { used: true, ..self.clone() } } } } @@ -147,9 +152,11 @@ fn prank( ) -> Result { let prank = Prank::new(prank_caller, prank_origin, new_caller, new_origin, depth, single_call); - if let Some(Prank { used, ..}) = state.prank { + if let Some(Prank { used, .. }) = state.prank { if !used { - return Err("You cannot overwrite `prank` until it is applied at least once".encode().into()); + return Err("You cannot overwrite `prank` until it is applied at least once" + .encode() + .into()) } } From 4aad6f3d7dbc4827fc0d4202b1510c2ea67382ee Mon Sep 17 00:00:00 2001 From: 4meta5 Date: Sat, 29 Apr 2023 17:32:01 -0400 Subject: [PATCH 08/14] unit test for startPrank0 - and startPrank1 - --- startPrank0 | 0 startPrank1 | 0 testdata/cheats/Prank.t.sol | 52 +++++++++++++++++++++++++++++++++++++ 3 files changed, 52 insertions(+) create mode 100644 startPrank0 create mode 100644 startPrank1 diff --git a/startPrank0 b/startPrank0 new file mode 100644 index 0000000000000..e69de29bb2d1d diff --git a/startPrank1 b/startPrank1 new file mode 100644 index 0000000000000..e69de29bb2d1d diff --git a/testdata/cheats/Prank.t.sol b/testdata/cheats/Prank.t.sol index 914729d7beaed..1c7cf7198106b 100644 --- a/testdata/cheats/Prank.t.sol +++ b/testdata/cheats/Prank.t.sol @@ -169,7 +169,59 @@ contract PrankTest is DSTest { address(this), "msg.sender was not cleaned up", oldOrigin, "tx.origin invariant failed" ); } + + function testStartPrank1AfterStartPrank0(address sender, address origin) public { + // Perform the prank + address oldOrigin = tx.origin; + Victim victim = new Victim(); + cheats.startPrank(sender); + victim.assertCallerAndOrigin( + sender, "msg.sender was not set during prank", oldOrigin, "tx.origin was not set during prank" + ); + + // Ensure we cleaned up correctly + victim.assertCallerAndOrigin( + address(this), "msg.sender was not cleaned up", oldOrigin, "tx.origin invariant failed" + ); + + // Overwrite the prank + cheats.startPrank(sender, origin); + victim.assertCallerAndOrigin( + sender, "msg.sender was not set during prank", origin, "tx.origin invariant failed" + ); + + // Ensure we cleaned up correctly + victim.assertCallerAndOrigin( + address(this), "msg.sender was not cleaned up", oldOrigin, "tx.origin invariant failed" + ); + } + function testStartPrank0AfterStartPrank1(address sender, address origin) public { + // Perform the prank + address oldOrigin = tx.origin; + Victim victim = new Victim(); + cheats.startPrank(sender, origin); + victim.assertCallerAndOrigin( + sender, "msg.sender was not set during prank", origin, "tx.origin was not set during prank" + ); + + // Ensure we cleaned up correctly + victim.assertCallerAndOrigin( + address(this), "msg.sender was not cleaned up", oldOrigin, "tx.origin invariant failed" + ); + + // Overwrite the prank + cheats.startPrank(sender); + victim.assertCallerAndOrigin( + sender, "msg.sender was not set during prank", oldOrigin, "tx.origin invariant failed" + ); + + // Ensure we cleaned up correctly + victim.assertCallerAndOrigin( + address(this), "msg.sender was not cleaned up", oldOrigin, "tx.origin invariant failed" + ); + } + function testPrankConstructorSender(address sender) public { cheats.prank(sender); ConstructorVictim victim = new ConstructorVictim( From 56fe7094bf315b13579ae50a82d4ff7ce3d8bb7a Mon Sep 17 00:00:00 2001 From: 4meta5 Date: Sat, 29 Apr 2023 17:34:34 -0400 Subject: [PATCH 09/14] fix --- startPrank0 | 0 startPrank1 | 0 2 files changed, 0 insertions(+), 0 deletions(-) delete mode 100644 startPrank0 delete mode 100644 startPrank1 diff --git a/startPrank0 b/startPrank0 deleted file mode 100644 index e69de29bb2d1d..0000000000000 diff --git a/startPrank1 b/startPrank1 deleted file mode 100644 index e69de29bb2d1d..0000000000000 From a29f3375711fc7f0d8b6a1a37c014371d47dd35b Mon Sep 17 00:00:00 2001 From: 4meta5 Date: Sat, 29 Apr 2023 17:42:45 -0400 Subject: [PATCH 10/14] remove clones by only updating prank after first time applied --- evm/src/executor/inspector/cheatcodes/env.rs | 7 ++++--- evm/src/executor/inspector/cheatcodes/mod.rs | 5 ++++- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/evm/src/executor/inspector/cheatcodes/env.rs b/evm/src/executor/inspector/cheatcodes/env.rs index 04cd04e24f7cd..6cb27902e6c81 100644 --- a/evm/src/executor/inspector/cheatcodes/env.rs +++ b/evm/src/executor/inspector/cheatcodes/env.rs @@ -78,11 +78,12 @@ impl Prank { } /// Apply the prank by setting `used` to true iff it is false - pub fn apply(&self) -> Self { + /// Only returns self in the case it is updated + pub fn apply(&self) -> Option { if self.used { - self.clone() + None } else { - Prank { used: true, ..self.clone() } + Some(Prank { used: true, ..self.clone() }) } } } diff --git a/evm/src/executor/inspector/cheatcodes/mod.rs b/evm/src/executor/inspector/cheatcodes/mod.rs index 657cdeb1866e8..ea40b24856e22 100644 --- a/evm/src/executor/inspector/cheatcodes/mod.rs +++ b/evm/src/executor/inspector/cheatcodes/mod.rs @@ -608,8 +608,11 @@ where prank_applied = true; } // If prank applied, then update the prank state accordingly + // iff it has changed (after the first time applied) if prank_applied { - self.prank = Some(prank.apply()); + if let Some(applied_prank) = prank.apply() { + self.prank = Some(applied_prank); + } } } } From 2a89075f970c49543ce4d3a52f45f49458bbbc10 Mon Sep 17 00:00:00 2001 From: 4meta5 Date: Sat, 29 Apr 2023 17:43:42 -0400 Subject: [PATCH 11/14] fmt --- testdata/cheats/Prank.t.sol | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/testdata/cheats/Prank.t.sol b/testdata/cheats/Prank.t.sol index 1c7cf7198106b..a78ce977d798b 100644 --- a/testdata/cheats/Prank.t.sol +++ b/testdata/cheats/Prank.t.sol @@ -169,7 +169,7 @@ contract PrankTest is DSTest { address(this), "msg.sender was not cleaned up", oldOrigin, "tx.origin invariant failed" ); } - + function testStartPrank1AfterStartPrank0(address sender, address origin) public { // Perform the prank address oldOrigin = tx.origin; @@ -221,7 +221,7 @@ contract PrankTest is DSTest { address(this), "msg.sender was not cleaned up", oldOrigin, "tx.origin invariant failed" ); } - + function testPrankConstructorSender(address sender) public { cheats.prank(sender); ConstructorVictim victim = new ConstructorVictim( From c8d8506cf09d5740e1875b441da6f0ba99d893b0 Mon Sep 17 00:00:00 2001 From: 4meta5 Date: Sat, 29 Apr 2023 19:36:13 -0400 Subject: [PATCH 12/14] more readable names --- evm/src/executor/inspector/cheatcodes/env.rs | 4 ++-- evm/src/executor/inspector/cheatcodes/mod.rs | 6 +++--- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/evm/src/executor/inspector/cheatcodes/env.rs b/evm/src/executor/inspector/cheatcodes/env.rs index 6cb27902e6c81..a36db028ab6f6 100644 --- a/evm/src/executor/inspector/cheatcodes/env.rs +++ b/evm/src/executor/inspector/cheatcodes/env.rs @@ -78,8 +78,8 @@ impl Prank { } /// Apply the prank by setting `used` to true iff it is false - /// Only returns self in the case it is updated - pub fn apply(&self) -> Option { + /// Only returns self in the case it is updated (first application) + pub fn first_time_applied(&self) -> Option { if self.used { None } else { diff --git a/evm/src/executor/inspector/cheatcodes/mod.rs b/evm/src/executor/inspector/cheatcodes/mod.rs index ea40b24856e22..a75e57716d6da 100644 --- a/evm/src/executor/inspector/cheatcodes/mod.rs +++ b/evm/src/executor/inspector/cheatcodes/mod.rs @@ -607,10 +607,10 @@ where data.env.tx.caller = h160_to_b160(new_origin); prank_applied = true; } - // If prank applied, then update the prank state accordingly - // iff it has changed (after the first time applied) + + // If prank applied for first time, then update if prank_applied { - if let Some(applied_prank) = prank.apply() { + if let Some(applied_prank) = prank.first_time_applied() { self.prank = Some(applied_prank); } } From bb779a582023c3cf7c9641e133e0e7dba05104e1 Mon Sep 17 00:00:00 2001 From: Enrique Ortiz Date: Wed, 10 May 2023 15:33:42 -0400 Subject: [PATCH 13/14] chore: fix/add tests, use ensure util --- evm/src/executor/inspector/cheatcodes/env.rs | 10 ++--- testdata/cheats/Prank.t.sol | 45 ++++++++++++++++---- 2 files changed, 40 insertions(+), 15 deletions(-) diff --git a/evm/src/executor/inspector/cheatcodes/env.rs b/evm/src/executor/inspector/cheatcodes/env.rs index 5c1fed8f8b28d..9c9609f724def 100644 --- a/evm/src/executor/inspector/cheatcodes/env.rs +++ b/evm/src/executor/inspector/cheatcodes/env.rs @@ -138,19 +138,15 @@ fn prank( 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()) - } + ensure!(used, "You cannot overwrite `prank` until it is applied at least once"); } - + ensure!( state.broadcast.is_none(), "You cannot `prank` for a broadcasted transaction.\ Pass the desired tx.origin into the broadcast cheatcode call" ); - + state.prank = Some(prank); Ok(Bytes::new()) } diff --git a/testdata/cheats/Prank.t.sol b/testdata/cheats/Prank.t.sol index 475712e6b3c9f..9e3365ebd68aa 100644 --- a/testdata/cheats/Prank.t.sol +++ b/testdata/cheats/Prank.t.sol @@ -176,26 +176,54 @@ contract PrankTest is DSTest { Victim victim = new Victim(); cheats.startPrank(sender); victim.assertCallerAndOrigin( - sender, "msg.sender was not set during prank", oldOrigin, "tx.origin was not set during prank" + sender, "msg.sender was not set during prank", oldOrigin, "tx.origin was set during prank incorrectly" ); - // Ensure we cleaned up correctly + // Ensure prank is still up as startPrank covers multiple calls victim.assertCallerAndOrigin( - address(this), "msg.sender was not cleaned up", oldOrigin, "tx.origin invariant failed" + sender, "msg.sender was cleaned up incorrectly", oldOrigin, "tx.origin invariant failed" ); // Overwrite the prank cheats.startPrank(sender, origin); victim.assertCallerAndOrigin( - sender, "msg.sender was not set during prank", origin, "tx.origin invariant failed" + sender, "msg.sender was not set during prank", origin, "tx.origin was not set" ); - // Ensure we cleaned up correctly + // Ensure prank is still up as startPrank covers multiple calls + victim.assertCallerAndOrigin( + sender, "msg.sender was cleaned up incorrectly", origin, "tx.origin invariant failed" + ); + + cheats.stopPrank(); + // Ensure everything is back to normal after stopPrank victim.assertCallerAndOrigin( address(this), "msg.sender was not cleaned up", oldOrigin, "tx.origin invariant failed" ); } + function testFailOverwriteUnusedPrank(address sender, address origin) public { + // Set the prank, but not use it + address oldOrigin = tx.origin; + Victim victim = new Victim(); + cheats.startPrank(sender, origin); + // try to overwrite the prank. This should fail. + cheats.startPrank(address(this), origin); + } + + function testFailOverwriteUnusedPrankAfterSuccessfulPrank(address sender, address origin) public { + // Set the prank, but not use it + address oldOrigin = tx.origin; + Victim victim = new Victim(); + cheats.startPrank(sender, origin); + victim.assertCallerAndOrigin( + sender, "msg.sender was not set during prank", origin, "tx.origin was set during prank incorrectly" + ); + cheats.startPrank(address(this), origin); + // try to overwrite the prank. This should fail. + cheats.startPrank(sender, origin); + } + function testStartPrank0AfterStartPrank1(address sender, address origin) public { // Perform the prank address oldOrigin = tx.origin; @@ -205,17 +233,18 @@ contract PrankTest is DSTest { sender, "msg.sender was not set during prank", origin, "tx.origin was not set during prank" ); - // Ensure we cleaned up correctly + // Ensure prank is still ongoing as we haven't called stopPrank victim.assertCallerAndOrigin( - address(this), "msg.sender was not cleaned up", oldOrigin, "tx.origin invariant failed" + sender, "msg.sender was cleaned up incorrectly", origin, "tx.origin was cleaned up incorrectly" ); // Overwrite the prank cheats.startPrank(sender); victim.assertCallerAndOrigin( - sender, "msg.sender was not set during prank", oldOrigin, "tx.origin invariant failed" + sender, "msg.sender was not set during prank", oldOrigin, "tx.origin was not reset correctly" ); + cheats.stopPrank(); // Ensure we cleaned up correctly victim.assertCallerAndOrigin( address(this), "msg.sender was not cleaned up", oldOrigin, "tx.origin invariant failed" From a48965a29553effcc8f86247edb529a26a323e0d Mon Sep 17 00:00:00 2001 From: Enrique Ortiz Date: Wed, 10 May 2023 16:06:57 -0400 Subject: [PATCH 14/14] chore: add missing edge case test --- testdata/cheats/Prank.t.sol | 28 ++++++++++++++++++++++++---- 1 file changed, 24 insertions(+), 4 deletions(-) diff --git a/testdata/cheats/Prank.t.sol b/testdata/cheats/Prank.t.sol index 9e3365ebd68aa..0dddb2ffe702d 100644 --- a/testdata/cheats/Prank.t.sol +++ b/testdata/cheats/Prank.t.sol @@ -170,6 +170,28 @@ contract PrankTest is DSTest { ); } + function testStartPrank0AfterPrank1(address sender, address origin) public { + // Perform the prank + address oldOrigin = tx.origin; + Victim victim = new Victim(); + cheats.startPrank(sender, origin); + victim.assertCallerAndOrigin( + sender, "msg.sender was not set during prank", origin, "tx.origin was not set during prank" + ); + + // Overwrite the prank + cheats.startPrank(sender); + victim.assertCallerAndOrigin( + sender, "msg.sender was not set during prank", oldOrigin, "tx.origin invariant failed" + ); + + cheats.stopPrank(); + // Ensure we cleaned up correctly after stopping the prank + victim.assertCallerAndOrigin( + address(this), "msg.sender was not cleaned up", oldOrigin, "tx.origin invariant failed" + ); + } + function testStartPrank1AfterStartPrank0(address sender, address origin) public { // Perform the prank address oldOrigin = tx.origin; @@ -186,9 +208,7 @@ contract PrankTest is DSTest { // Overwrite the prank cheats.startPrank(sender, origin); - victim.assertCallerAndOrigin( - sender, "msg.sender was not set during prank", origin, "tx.origin was not set" - ); + victim.assertCallerAndOrigin(sender, "msg.sender was not set during prank", origin, "tx.origin was not set"); // Ensure prank is still up as startPrank covers multiple calls victim.assertCallerAndOrigin( @@ -210,7 +230,7 @@ contract PrankTest is DSTest { // try to overwrite the prank. This should fail. cheats.startPrank(address(this), origin); } - + function testFailOverwriteUnusedPrankAfterSuccessfulPrank(address sender, address origin) public { // Set the prank, but not use it address oldOrigin = tx.origin;