Merge pull request #167 from arcz/prank-cheatcode

ethereum · Jan 9, 2023 · 2f49891 · 2f49891
2 parents f82cbd4 + 560ee43
commit 2f49891
Show file tree

Hide file tree

Showing 4 changed files with 37 additions and 3 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -5,6 +5,12 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## Unreleased
+
+### Added
+
+- New cheatcode `prank(address)` that sets `msg.sender` to the specified address for the next call.
+
 ## [0.50.2] - 2023-01-06
 
 ### Fixed

diff --git a/doc/src/controlling-the-unit-testing-environment.md b/doc/src/controlling-the-unit-testing-environment.md
@@ -29,6 +29,9 @@ These can be accessed by calling into a contract (typically called `Hevm`) at ad
   Executes the arguments as a command in the system shell and returns stdout. Expects abi encoded values to be returned from the shell or an error will be thrown. Note that this
   cheatcode means test authors can execute arbitrary code on user machines as part of a call to `dapp test`, for this reason all calls to `ffi` will fail unless the `--ffi` flag is passed.
 
+- `function prank(address sender) public`
+  Sets `msg.sender` to the specified `sender` for the next call.
+
 ## Environment Variables
 
 These environment variables can be used to control block parameters:

diff --git a/src/EVM.hs b/src/EVM.hs
@@ -112,6 +112,7 @@ data VM = VM
   , _constraints    :: [Prop]
   , _keccakEqs      :: [Prop]
   , _allowFFI       :: Bool
+  , _overrideCaller :: Maybe (Expr EWord)
   }
   deriving (Show)
 
@@ -536,6 +537,7 @@ makeVm o =
   , _keccakEqs = mempty
   , _iterations = mempty
   , _allowFFI = vmoptAllowFFI o
+  , _overrideCaller = Nothing
   }
 
 -- | Initialize empty contract with given code
@@ -1178,8 +1180,9 @@ exec1 = do
                   delegateCall this (num xGas) xTo xTo xValue xInOffset xInSize xOutOffset xOutSize xs $ \callee -> do
                     zoom state $ do
                       assign callvalue (Lit xValue)
-                      assign caller (litAddr self)
+                      assign caller $ fromMaybe (litAddr self) (vm ^. overrideCaller)
                       assign contract callee
+                    assign overrideCaller Nothing
                     transfer self callee xValue
                     touchAccount self
                     touchAccount callee
@@ -1202,7 +1205,8 @@ exec1 = do
                   delegateCall this (num xGas) xTo (litAddr self) xValue xInOffset xInSize xOutOffset xOutSize xs $ \_ -> do
                     zoom state $ do
                       assign callvalue (Lit xValue)
-                      assign caller (litAddr self)
+                      assign caller $ fromMaybe (litAddr self) (vm ^. overrideCaller)
+                    assign overrideCaller Nothing
                     touchAccount self
             _ ->
               underrun
@@ -1293,9 +1297,10 @@ exec1 = do
                 delegateCall this (num xGas) xTo xTo 0 xInOffset xInSize xOutOffset xOutSize xs $ \callee -> do
                   zoom state $ do
                     assign callvalue (Lit 0)
-                    assign caller (litAddr self)
+                    assign caller $ fromMaybe (litAddr self) (vm ^. overrideCaller)
                     assign contract callee
                     assign static True
+                  assign overrideCaller Nothing
                   touchAccount self
                   touchAccount callee
             _ ->
@@ -2050,6 +2055,11 @@ cheatActions =
                       addr = Lit . W256 . word256 . BS.drop 12 . BS.take 32 . keccakBytes $ pub
                     assign (state . returndata . word256At (Lit 0)) addr
                     assign (state . memory . word256At outOffset) addr
+          _ -> vmError (BadCheatCode sig),
+
+      action "prank(address)" $
+        \sig _ _ input -> case decodeStaticArgs 0 1 input of
+          [addr]  -> assign overrideCaller (Just addr)
           _ -> vmError (BadCheatCode sig)
 
     ]

diff --git a/test/contracts/pass/cheatCodes.sol b/test/contracts/pass/cheatCodes.sol
@@ -10,12 +10,19 @@ interface Hevm {
     function sign(uint256,bytes32) external returns (uint8,bytes32,bytes32);
     function addr(uint256) external returns (address);
     function ffi(string[] calldata) external returns (bytes memory);
+    function prank(address) external;
 }
 
 contract HasStorage {
     uint slot0 = 10;
 }
 
+contract Prankster {
+    function prankme() public returns (address) {
+        return msg.sender;
+    }
+}
+
 contract CheatCodes is DSTest {
     address store = address(new HasStorage());
     Hevm hevm = Hevm(HEVM_ADDRESS);
@@ -86,4 +93,12 @@ contract CheatCodes is DSTest {
         (string memory output) = abi.decode(hevm.ffi(inputs), (string));
         assertEq(output, "acab");
     }
+
+    function test_prank() public {
+        Prankster prankster = new Prankster();
+        assertEq(prankster.prankme(), address(this));
+        hevm.prank(address(0xdeadbeef));
+        assertEq(prankster.prankme(), address(0xdeadbeef));
+        assertEq(prankster.prankme(), address(this));
+    }
 }