feat(era:cheatcodes): expectRevert v0

crates/era-cheatcodes/src/cheatcodes.rs

@@ -61,6 +61,7 @@ const INTERNAL_CONTRACT_ADDRESSES: [H160; 20] = [
 #[derive(Debug, Default, Clone)]
 pub struct CheatcodeTracer {
     one_time_actions: Vec<FinishCycleOneTimeActions>,
+    next_execution_actions: Vec<NextExecutionOneTimeActions>,
     permanent_actions: FinishCyclePermanentActions,
     return_data: Option<Vec<U256>>,
     return_ptr: Option<FatPointer>,
@@ -72,6 +73,13 @@ pub struct CheatcodeTracer {
 enum FinishCycleOneTimeActions {
     StorageWrite { key: StorageKey, read_value: H256, write_value: H256 },
     StoreFactoryDep { hash: U256, bytecode: Vec<U256> },
+    ForceRevert { error: Vec<u8> },
+    ForceReturn { data: Vec<u8> },
+}
+
+#[derive(Debug, Clone)]
+enum NextExecutionOneTimeActions {
+    ExpectRevert { reason: Option<Vec<u8>>, depth: usize },
 }
 
 #[derive(Debug, Default, Clone)]
@@ -95,6 +103,37 @@ impl<S: DatabaseExt + Send, H: HistoryMode> DynTracer<EraDb<S>, SimpleMemory<H>>
         memory: &SimpleMemory<H>,
         storage: StoragePtr<EraDb<S>>,
     ) {
+        //Only execute "next execution" actions when a cheatcode isn't being invoked
+        if state.vm_local_state.callstack.current.code_address != CHEATCODE_ADDRESS {
+            // in `handle_action`, when true is returned the current action will
+            // be kept in the queue
+            let handle_action = |action: &NextExecutionOneTimeActions| match action {
+                NextExecutionOneTimeActions::ExpectRevert { reason, depth }
+                    if state.vm_local_state.callstack.depth() > *depth =>
+                {
+                    match data.opcode.variant.opcode {
+                        Opcode::Ret(op) => {
+                            self.one_time_actions.push(
+                                Self::handle_except_revert(reason.as_ref(), op, &state, memory)
+                                    .map(|_| FinishCycleOneTimeActions::ForceReturn {
+                                        //dummy data
+                                        data: vec![0u8; 8192],
+                                    })
+                                    .unwrap_or_else(|error| {
+                                        FinishCycleOneTimeActions::ForceRevert { error }
+                                    }),
+                            );
+                            false
+                        }
+                        _ => true,
+                    }
+                }
+                _ => true,
+            };
+
+            self.next_execution_actions.retain(handle_action);
+        }
+
         if self.return_data.is_some() {
             if let Opcode::Ret(_call) = data.opcode.variant.opcode {
                 if self.near_calls == 0 {
@@ -175,6 +214,19 @@ impl<S: DatabaseExt + Send, H: HistoryMode> VmTracer<EraDb<S>, H> for CheatcodeT
                 FinishCycleOneTimeActions::StoreFactoryDep { hash, bytecode } => state
                     .decommittment_processor
                     .populate(vec![(hash, bytecode)], Timestamp(state.local_state.timestamp)),
+                FinishCycleOneTimeActions::ForceReturn { data: _ } => {
+                    //TODO: override return data with the given one and force return (instead of
+                    // revert)
+                }
+                FinishCycleOneTimeActions::ForceRevert { error } => {
+                    return TracerExecutionStatus::Stop(
+                        multivm::interface::tracer::TracerExecutionStopReason::Abort(
+                            multivm::interface::Halt::Unknown(VmRevertReason::from(
+                                error.as_slice(),
+                            )),
+                        ),
+                    )
+                }
             }
         }
 
@@ -209,6 +261,7 @@ impl CheatcodeTracer {
     pub fn new() -> Self {
         CheatcodeTracer {
             one_time_actions: vec![],
+            next_execution_actions: vec![],
             permanent_actions: FinishCyclePermanentActions { start_prank: None },
             near_calls: 0,
             return_data: None,
@@ -219,7 +272,7 @@ impl CheatcodeTracer {
 
     pub fn dispatch_cheatcode<S: DatabaseExt + Send, H: HistoryMode>(
         &mut self,
-        _state: VmLocalStateData<'_>,
+        state: VmLocalStateData<'_>,
         _data: AfterExecutionData,
         _memory: &SimpleMemory<H>,
         storage: StoragePtr<EraDb<S>>,
@@ -253,6 +306,14 @@ impl CheatcodeTracer {
                 self.store_factory_dep(hash, code);
                 self.write_storage(code_key, u256_to_h256(hash), &mut storage.borrow_mut());
             }
+            expectRevert_0(expectRevert_0Call {}) => {
+                self.add_except_revert(None, state.vm_local_sate.callstack.depth())
+            }
+            expectRevert_1(expectRevert_1Call { revertData }) |
+            expectRevert_2(expectRevert_2Call { revertData }) => self.add_except_revert(
+                Some(revertData.to_vec()),
+                state.vm_local_sate.callstack.depth(),
+            ),
             getNonce_0(getNonce_0Call { account }) => {
                 tracing::info!("👷 Getting nonce for {account:?}");
                 let mut storage = storage.borrow_mut();
@@ -515,4 +576,72 @@ impl CheatcodeTracer {
 
         self.return_data = Some(data);
     }
+
+    fn add_except_revert(&mut self, reason: Option<Vec<u8>>, depth: usize) {
+        self.next_execution_actions
+            .push(NextExecutionOneTimeActions::ExpectRevert { reason, depth });
+    }
+
+    fn handle_except_revert<H: HistoryMode>(
+        reason: Option<&Vec<u8>>,
+        op: zkevm_opcode_defs::RetOpcode,
+        state: &VmLocalStateData<'_>,
+        memory: &SimpleMemory<H>,
+    ) -> Result<(), Vec<u8>> {
+        match (op, reason) {
+            (zkevm_opcode_defs::RetOpcode::Revert, Some(expected_reason)) => {
+                let retdata = {
+                    let ptr = state.vm_local_state.registers
+                        [CALL_IMPLICIT_CALLDATA_FAT_PTR_REGISTER as usize];
+                    assert!(ptr.is_pointer);
+                    let fat_data_pointer = FatPointer::from_u256(ptr.value);
+                    memory.read_unaligned_bytes(
+                        fat_data_pointer.memory_page as usize,
+                        fat_data_pointer.start as usize,
+                        fat_data_pointer.length as usize,
+                    )
+                };
+
+                if !expected_reason.is_empty() && retdata.is_empty() {
+                    return Err("call reverted as expected, but without data".to_string().into())
+                }
+
+                let mut actual_revert: Vec<u8> = retdata.into();
+
+                // Try decoding as known errors
+                // alloy_sol_types::Revert = "Error(string)" => [0x08, 0xc3, 0x79, 0xa0]
+                // CheatCodeError = "CheatcodeError(string)" => [0xee, 0xaa, 0x9e, 0x6f]
+                if matches!(
+                    actual_revert.get(..4),
+                    Some(&[0x08, 0xc3, 0x79, 0xa0] | &[0xee, 0xaa, 0x9e, 0x6f])
+                ) {
+                    if let Ok(decoded) = Vec::<u8>::decode(&actual_revert[4..]) {
+                        actual_revert = decoded;
+                    }
+                }
+
+                if &actual_revert == expected_reason {
+                    Ok(())
+                } else {
+                    let stringify = |data: &[u8]| {
+                        String::decode(data)
+                            .ok()
+                            .or_else(|| std::str::from_utf8(data).ok().map(ToOwned::to_owned))
+                            .unwrap_or_else(|| data.to_vec().encode_hex())
+                    };
+                    Err(format!(
+                        "Error != expected error: {} != {}",
+                        stringify(&actual_revert),
+                        stringify(expected_reason),
+                    )
+                    .into())
+                }
+            }
+            (zkevm_opcode_defs::RetOpcode::Revert, None) => Ok(()),
+            (zkevm_opcode_defs::RetOpcode::Ok, _) => {
+                Err("expected revert but call succeeded".to_string().into())
+            }
+            (zkevm_opcode_defs::RetOpcode::Panic, _) => todo!("ignore/return error ?"),
+        }
+    }
 }

crates/era-cheatcodes/tests/src/cheatcodes/ExpectRevert.t.sol

@@ -0,0 +1,195 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.18;
+
+import {Test, console2 as console} from "../../lib/forge-std/src/Test.sol";
+import {Constants} from "./Constants.sol";
+
+interface Cheatcodes {
+    function expectRevert() external;
+    function expectRevert(bytes4 revertData) external;
+    function expectRevert(bytes calldata revertData) external;
+}
+
+contract Reverter {
+    error CustomError();
+
+    function revertWithMessage(string memory message) public pure {
+        revert(message);
+    }
+
+    function doNotRevert() public pure {}
+
+    function panic() public pure returns (uint256) {
+        return uint256(100) - uint256(101);
+    }
+
+    function revertWithCustomError() public pure {
+        revert CustomError();
+    }
+
+    function nestedRevert(Reverter inner, string memory message) public pure {
+        inner.revertWithMessage(message);
+    }
+
+    function callThenRevert(Dummy dummy, string memory message) public pure {
+        dummy.callMe();
+        revert(message);
+    }
+
+    function revertWithoutReason() public pure {
+        revert();
+    }
+}
+
+contract ConstructorReverter {
+    constructor(string memory message) {
+        revert(message);
+    }
+}
+
+/// Used to ensure that the dummy data from `cheatcodes.expectRevert`
+/// is large enough to decode big structs.
+///
+/// The struct is based on issue #2454
+struct LargeDummyStruct {
+    address a;
+    uint256 b;
+    bool c;
+    address d;
+    address e;
+    string f;
+    address[8] g;
+    address h;
+    uint256 i;
+}
+
+contract Dummy {
+    function callMe() public pure returns (string memory) {
+        return "thanks for calling";
+    }
+
+    function largeReturnType() public pure returns (LargeDummyStruct memory) {
+        revert("reverted with large return type");
+    }
+}
+
+contract ExpectRevertTest is Test {
+    Cheatcodes constant cheatcodes = Cheatcodes(Constants.CHEATCODE_ADDRESS);
+
+    function shouldRevert() internal {
+        revert();
+    }
+
+    function testExpectRevertString() public {
+        Reverter reverter = new Reverter();
+        cheatcodes.expectRevert("revert");
+        reverter.revertWithMessage("revert");
+    }
+
+    function testFailExpectRevertWrongString() public {
+        Reverter reverter = new Reverter();
+        cheatcodes.expectRevert("my not so cool error");
+        reverter.revertWithMessage("my cool error");
+    }
+
+    function testFailRevertNotOnImmediateNextCall() public {
+        Reverter reverter = new Reverter();
+        // expectRevert should only work for the next call. However,
+        // we do not immediately revert, so,
+        // we fail.
+        cheatcodes.expectRevert("revert");
+        reverter.doNotRevert();
+        reverter.revertWithMessage("revert");
+    }
+
+    function testExpectRevertConstructor() public {
+        cheatcodes.expectRevert("constructor revert");
+        new ConstructorReverter("constructor revert");
+    }
+
+    function testExpectRevertBuiltin() public {
+        Reverter reverter = new Reverter();
+        cheatcodes.expectRevert(abi.encodeWithSignature("Panic(uint256)", 0x11));
+        reverter.panic();
+    }
+
+    function testExpectRevertCustomError() public {
+        Reverter reverter = new Reverter();
+        cheatcodes.expectRevert(abi.encodePacked(Reverter.CustomError.selector));
+        reverter.revertWithCustomError();
+    }
+
+    function testExpectRevertNested() public {
+        Reverter reverter = new Reverter();
+        Reverter inner = new Reverter();
+        cheatcodes.expectRevert("nested revert");
+        reverter.nestedRevert(inner, "nested revert");
+    }
+
+    function testExpectRevertCallsThenReverts() public {
+        Reverter reverter = new Reverter();
+        Dummy dummy = new Dummy();
+        cheatcodes.expectRevert("called a function and then reverted");
+        reverter.callThenRevert(dummy, "called a function and then reverted");
+    }
+
+    function testDummyReturnDataForBigType() public {
+        Dummy dummy = new Dummy();
+        cheatcodes.expectRevert("reverted with large return type");
+        dummy.largeReturnType();
+    }
+
+    function testFailExpectRevertErrorDoesNotMatch() public {
+        Reverter reverter = new Reverter();
+        cheatcodes.expectRevert("should revert with this message");
+        reverter.revertWithMessage("but reverts with this message");
+    }
+
+    function testFailExpectRevertDidNotRevert() public {
+        Reverter reverter = new Reverter();
+        cheatcodes.expectRevert("does not revert, but we think it should");
+        reverter.doNotRevert();
+    }
+
+    function testExpectRevertNoReason() public {
+        Reverter reverter = new Reverter();
+        cheatcodes.expectRevert(bytes(""));
+        reverter.revertWithoutReason();
+    }
+
+    function testExpectRevertAnyRevert() public {
+        cheatcodes.expectRevert();
+        new ConstructorReverter("hello this is a revert message");
+
+        Reverter reverter = new Reverter();
+        cheatcodes.expectRevert();
+        reverter.revertWithMessage("this is also a revert message");
+
+        cheatcodes.expectRevert();
+        reverter.panic();
+
+        cheatcodes.expectRevert();
+        reverter.revertWithCustomError();
+
+        Reverter reverter2 = new Reverter();
+        cheatcodes.expectRevert();
+        reverter.nestedRevert(reverter2, "this too is a revert message");
+
+        Dummy dummy = new Dummy();
+        cheatcodes.expectRevert();
+        reverter.callThenRevert(dummy, "revert message 4 i ran out of synonims for also");
+
+        cheatcodes.expectRevert();
+        reverter.revertWithoutReason();
+    }
+
+    function testFailExpectRevertAnyRevertDidNotRevert() public {
+        Reverter reverter = new Reverter();
+        cheatcodes.expectRevert();
+        reverter.doNotRevert();
+    }
+
+    function testFailExpectRevertDangling() public {
+        cheatcodes.expectRevert("dangling");
+    }
+}