From a77c19f71253013d508327cf74f3a3005d45d5ae Mon Sep 17 00:00:00 2001 From: alec Date: Tue, 14 May 2024 11:41:21 -0400 Subject: [PATCH 1/2] added startPrank() and stopPrank() cheatcodes + tests --- doc/src/ds-test-tutorial.md | 2 ++ src/EVM.hs | 28 +++++++++++++++-- src/EVM/Types.hs | 1 + test/contracts/pass/cheatCodes.sol | 43 ++++++++++++++++++++++++++ test/contracts/pass/cheatCodesFork.sol | 2 ++ 5 files changed, 73 insertions(+), 3 deletions(-) diff --git a/doc/src/ds-test-tutorial.md b/doc/src/ds-test-tutorial.md index 3aa9f9e2e..ca92b0e5e 100644 --- a/doc/src/ds-test-tutorial.md +++ b/doc/src/ds-test-tutorial.md @@ -191,6 +191,8 @@ implementing the following methods: | Function | Description | | --- | --- | |`function prank(address sender) public`| Sets `msg.sender` to the specified `sender` for the next call.| +|`function startPrank(address sender) public`| Sets `msg.sender` to the specified `sender` until `stopPrank()` is called.| +|`function stopPrank() public`| Resets `msg.sender` to the default sender.| |`function deal(address usr, uint amt) public`| Sets the eth balance of `usr` to `amt`. Note that if `usr` is a symbolic address, then it must be the address of a contract that has already been deployed. This restriction is in place to ensure soundness of our symbolic address encoding with respect to potential aliasing of symbolic addresses.| |`function store(address c, bytes32 loc, bytes32 val) public`| Sets the slot `loc` of contract `c` to `val`.| |`function warp(uint x) public`| Sets the block timestamp to `x`.| diff --git a/src/EVM.hs b/src/EVM.hs index 3754c7466..a3a1ec78d 100644 --- a/src/EVM.hs +++ b/src/EVM.hs @@ -144,6 +144,7 @@ makeVm o = do , iterations = mempty , config = RuntimeConfig { allowFFI = o.allowFFI + , resetCaller = True , overrideCaller = Nothing , baseState = o.baseState } @@ -770,7 +771,9 @@ exec1 = do assign #callvalue xValue assign #caller from' assign #contract callee - assign (#config % #overrideCaller) Nothing + do + resetCaller <- use (#config % #resetCaller) + when (resetCaller) $ assign (#config % #overrideCaller) Nothing touchAccount from' touchAccount callee transfer from' callee xValue @@ -788,7 +791,9 @@ exec1 = do zoom #state $ do assign #callvalue xValue assign #caller $ fromMaybe self vm.config.overrideCaller - assign (#config % #overrideCaller) Nothing + do + resetCaller <- use (#config % #resetCaller) + when (resetCaller) $ assign (#config % #overrideCaller) Nothing touchAccount self _ -> underrun @@ -878,7 +883,9 @@ exec1 = do assign #caller $ fromMaybe self (vm.config.overrideCaller) assign #contract callee assign #static True - assign (#config % #overrideCaller) Nothing + do + resetCaller <- use (#config % #resetCaller) + when (resetCaller) $ assign (#config % #overrideCaller) Nothing touchAccount self touchAccount callee _ -> @@ -1657,6 +1664,21 @@ cheatActions = Map.fromList Just a -> assign (#config % #overrideCaller) (Just a) Nothing -> vmError (BadCheatCode sig) _ -> vmError (BadCheatCode sig) + + , action "startPrank(address)" $ + \sig _ _ input -> case decodeStaticArgs 0 1 input of + [addr] -> case wordToAddr addr of + Just a -> do + assign (#config % #overrideCaller) (Just a) + assign (#config % #resetCaller) False + Nothing -> vmError (BadCheatCode sig) + _ -> vmError (BadCheatCode sig) + + , action "stopPrank()" $ + \_ _ _ _ -> do + assign (#config % #overrideCaller) Nothing + assign (#config % #resetCaller) True + , action "createFork(string)" $ \sig outOffset _ input -> case decodeBuf [AbiStringType] input of diff --git a/src/EVM/Types.hs b/src/EVM/Types.hs index cae8d4811..59fe1eab7 100644 --- a/src/EVM/Types.hs +++ b/src/EVM/Types.hs @@ -652,6 +652,7 @@ data BaseState data RuntimeConfig = RuntimeConfig { allowFFI :: Bool , overrideCaller :: Maybe (Expr EAddr) + , resetCaller :: Bool , baseState :: BaseState } deriving (Show) diff --git a/test/contracts/pass/cheatCodes.sol b/test/contracts/pass/cheatCodes.sol index c77cd998b..769834de5 100644 --- a/test/contracts/pass/cheatCodes.sol +++ b/test/contracts/pass/cheatCodes.sol @@ -13,6 +13,8 @@ interface Hevm { function addr(uint256) external returns (address); function ffi(string[] calldata) external returns (bytes memory); function prank(address) external; + function startPrank(address) external; + function stopPrank() external; function label(address addr, string calldata label) external; } @@ -129,6 +131,28 @@ contract CheatCodes is DSTest { assertEq(prankster.prankme(), address(this)); } + + function prove_startPrank(address caller) public { + Prankster prankster = new Prankster(); + assertEq(prankster.prankme(), address(this)); + + hevm.startPrank(address(0xdeadbeef)); + assertEq(prankster.prankme(), address(0xdeadbeef)); + assertEq(prankster.prankme(), address(0xdeadbeef)); + hevm.stopPrank(); + + hevm.startPrank(caller); + assertEq(prankster.prankme(), caller); + assertEq(prankster.prankme(), caller); + hevm.stopPrank(); + + assertEq(prankster.prankme(), address(this)); + + hevm.prank(caller); + assertEq(prankster.prankme(), caller); + assertEq(prankster.prankme(), address(this)); + } + // this is not supported yet due to restrictions around symbolic address aliasing... function proveFail_deal_unknown_address(address e, uint val) public { hevm.deal(e, val); @@ -212,6 +236,25 @@ contract CheatCodes is DSTest { assertEq(b.balance, 1); } + function prove_startPrank_val() public { + address payable a = payable(address(new Payable())); + address payable b = payable(address(new Payable())); + + // send this.balance to a + a.call{value: address(this).balance}(""); + uint aBal = a.balance; + + // send 1 wei from a to b + hevm.startPrank(a); + address(b).call{value: 1}(""); + address(b).call{value: 1}(""); + hevm.stopPrank(); + + // check balances + assertEq(a.balance, aBal - 2); + assertEq(b.balance, 2); + } + function prove_label_works() public { hevm.label(address(this), "label"); assert(true); diff --git a/test/contracts/pass/cheatCodesFork.sol b/test/contracts/pass/cheatCodesFork.sol index b0d4c16f6..2f716070b 100644 --- a/test/contracts/pass/cheatCodesFork.sol +++ b/test/contracts/pass/cheatCodesFork.sol @@ -11,6 +11,8 @@ interface Hevm { function addr(uint256) external returns (address); function ffi(string[] calldata) external returns (bytes memory); function prank(address) external; + function startPrank(address) external; + function stopPrank() external; function deal(address,uint256) external; function createFork(string calldata urlOrAlias) external returns (uint256); function selectFork(uint256 forkId) external; From e07ddafb174a7cb5bf2ba38f7409b85a61ef2b23 Mon Sep 17 00:00:00 2001 From: alec Date: Tue, 14 May 2024 11:45:48 -0400 Subject: [PATCH 2/2] updated changelog --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 545092700..5ce826d55 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,6 +16,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - More PEq, PLEq, and PLT rules - New `label` cheatcode. - Updated Bitwuzla to newer version +- New cheatcodes `startPrank()` & `stopPrank()` ## Fixed - `concat` is a 2-ary, not an n-ary function in SMT2LIB, declare-const does not exist in QF_AUFBV, replacing