From 08bf7bf8a0eff64250dfb2294722ed2461d8b1d2 Mon Sep 17 00:00:00 2001 From: Artur Cygan Date: Fri, 6 Jan 2023 21:56:57 +0100 Subject: [PATCH 1/2] Implement prank(address) cheatcode --- src/EVM.hs | 16 +++++++++++++--- test/contracts/pass/cheatCodes.sol | 15 +++++++++++++++ 2 files changed, 28 insertions(+), 3 deletions(-) diff --git a/src/EVM.hs b/src/EVM.hs index 5dbd011ec..a58e157c5 100644 --- 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 index 975d19c64..068849bd6 100644 --- 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)); + } } From 560ee43ab3708b695e1d2f956cd95606ae808c9c Mon Sep 17 00:00:00 2001 From: Artur Cygan Date: Mon, 9 Jan 2023 13:21:24 +0100 Subject: [PATCH 2/2] Update changelog and docs with prank cheatcode --- CHANGELOG.md | 6 ++++++ doc/src/controlling-the-unit-testing-environment.md | 3 +++ 2 files changed, 9 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 90d67805f..1375eca19 100644 --- 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 index b04e24105..7939ef5fe 100644 --- 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: