Added startPrank() and stopPrank() cheatcodes to hevm.

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

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`.|

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

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)

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);

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;