feat(forge): readCallers cheatcode

evm/src/executor/abi/mod.rs

@@ -76,6 +76,7 @@ abigen!(
 
         prank(address)
         prank(address,address)
+        readPrank()(bool,address,address)
         startPrank(address)
         startPrank(address,address)
         stopPrank()

evm/src/executor/inspector/cheatcodes/env.rs

@@ -10,7 +10,7 @@ use crate::{
 use ethers::{
     abi::{self, AbiEncode, RawLog, Token, Tokenizable, Tokenize},
     signers::{LocalWallet, Signer},
-    types::{Address, Bytes, U256},
+    types::{Address, Bytes, H160, U256},
 };
 use foundry_config::Config;
 use revm::{primitives::Bytecode, Database, EVMData};
@@ -108,6 +108,24 @@ fn prank(
     Ok(Bytes::new())
 }
 
+fn read_prank(prank: &Option<Prank>) -> Bytes {
+    let is_prank_active = prank.is_some();
+
+    let data = if let Some(prank) = prank {
+        [
+            Token::Bool(is_prank_active),
+            Token::Address(prank.new_caller),
+            Token::Address(prank.new_origin.unwrap_or_default()),
+        ]
+    } else {
+        let zero_address = H160::default();
+
+        [Token::Bool(is_prank_active), Token::Address(zero_address), Token::Address(zero_address)]
+    };
+
+    abi::encode(&data).into()
+}
+
 #[derive(Clone, Debug, Default)]
 pub struct RecordAccess {
     pub reads: BTreeMap<Address, Vec<U256>>,
@@ -308,6 +326,7 @@ pub fn apply<DB: DatabaseExt>(
             state.prank = None;
             Bytes::new()
         }
+        HEVMCalls::ReadPrank(_) => read_prank(&state.prank),
         HEVMCalls::Record(_) => {
             start_record(state);
             Bytes::new()

testdata/cheats/Cheats.sol

@@ -152,6 +152,9 @@ interface Cheats {
     // Resets subsequent calls' msg.sender to be `address(this)`
     function stopPrank() external;
 
+    // Reads the current prank from state
+    function readPrank() external returns (bool, address, address);
+
     // Sets an address' balance, (who, newBalance)
     function deal(address, uint256) external;
 

testdata/cheats/ReadPrank.t.sol

@@ -0,0 +1,169 @@
+// SPDX-License-Identifier: Unlicense
+pragma solidity >=0.8.0;
+
+import "ds-test/test.sol";
+import "./Cheats.sol";
+import "../logs/console.sol";
+
+contract Target {
+    function consumePrank() external {}
+}
+
+contract ReadPrankTest is DSTest {
+    Cheats constant cheats = Cheats(HEVM_ADDRESS);
+
+    function testReadPrankWithNoActivePrank() public {
+        // Act
+        (bool isActive, address newSender, address newOrigin) = cheats
+            .readPrank();
+
+        // Assert
+        assertEq(isActive, false);
+        assertEq(newSender, address(0));
+        assertEq(newOrigin, address(0));
+    }
+
+    function testReadPrankWithActivePrankForMsgSender(address sender) public {
+        // Arrange
+        cheats.prank(sender);
+
+        // Act
+        (bool isActive, address newSender, address newOrigin) = cheats
+            .readPrank();
+
+        // Assert
+        assertTrue(isActive);
+        assertEq(newSender, sender);
+        assertEq(newOrigin, address(0));
+    }
+
+    function testReadPrankWithActivePrankForMsgSenderAndTxOrigin(
+        address sender,
+        address origin
+    ) public {
+        // Arrange
+        cheats.prank(sender, origin);
+
+        // Act
+        (bool isActive, address newSender, address newOrigin) = cheats
+            .readPrank();
+
+        // Assert
+        assertTrue(isActive);
+        assertEq(newSender, sender);
+        assertEq(newOrigin, origin);
+    }
+
+    function testReadPrankAfterConsumingMsgSenderPrank(address sender) public {
+        // Arrange
+        Target target = new Target();
+        cheats.prank(sender);
+
+        // Act
+        target.consumePrank();
+        (bool isActive, address newSender, address newOrigin) = cheats
+            .readPrank();
+
+        // Assert
+        assertEq(isActive, false);
+        assertEq(newSender, address(0));
+        assertEq(newOrigin, address(0));
+    }
+
+    function testReadPrankAfterConsumingMsgSenderAndTxOriginPrank(
+        address sender,
+        address origin
+    ) public {
+        // Arrange
+        Target target = new Target();
+        cheats.prank(sender, origin);
+
+        // Act
+        target.consumePrank();
+        (bool isActive, address newSender, address newOrigin) = cheats
+            .readPrank();
+
+        // Assert
+        assertEq(isActive, false);
+        assertEq(newSender, address(0));
+        assertEq(newOrigin, address(0));
+    }
+
+    function testReadPrankWithActiveRecurrentMsgSenderPrank(address sender)
+        public
+    {
+        // Arrange
+        Target target = new Target();
+        cheats.startPrank(sender);
+
+        for (uint i = 0; i < 5; i++) {
+            // Act
+            target.consumePrank();
+            (bool isActive, address newSender, address newOrigin) = cheats
+                .readPrank();
+
+            // Assert
+            assertTrue(isActive);
+            assertEq(newSender, sender);
+            assertEq(newOrigin, address(0));
+        }
+    }
+
+    function testReadPrankWithActiveRecurrentMsgSenderAndTxOriginPrank(
+        address sender,
+        address origin
+    ) public {
+        // Arrange
+        Target target = new Target();
+        cheats.startPrank(sender, origin);
+
+        for (uint i = 0; i < 5; i++) {
+            // Act
+            target.consumePrank();
+            (bool isActive, address newSender, address newOrigin) = cheats
+                .readPrank();
+
+            // Assert
+            assertTrue(isActive);
+            assertEq(newSender, sender);
+            assertEq(newOrigin, origin);
+        }
+    }
+
+    function testReadPrankAfterStoppingRecurrentMsgSenderPrank(address sender)
+        public
+    {
+        // Arrange
+        cheats.startPrank(sender);
+
+        // Act
+        cheats.stopPrank();
+
+        (bool isActive, address newSender, address newOrigin) = cheats
+            .readPrank();
+
+        // Assert
+        assertEq(isActive, false);
+        assertEq(newSender, address(0));
+        assertEq(newOrigin, address(0));
+    }
+
+    function testReadPrankAfterStoppingRecurrentMsgSenderAndTxOriginPrank(
+        address sender,
+        address origin
+    ) public {
+        // Arrange
+        cheats.startPrank(sender, origin);
+
+        // Act
+        cheats.stopPrank();
+
+        (bool isActive, address newSender, address newOrigin) = cheats
+            .readPrank();
+
+        // Assert
+        assertEq(isActive, false);
+        assertEq(newSender, address(0));
+        assertEq(newOrigin, address(0));
+    }
+}