diff --git a/crates/cheatcodes/assets/cheatcodes.json b/crates/cheatcodes/assets/cheatcodes.json index 7198f615288a..c4f6cc8d738b 100644 --- a/crates/cheatcodes/assets/cheatcodes.json +++ b/crates/cheatcodes/assets/cheatcodes.json @@ -6933,6 +6933,26 @@ "status": "stable", "safety": "safe" }, + { + "func": { + "id": "stopExpectSafeMemory", + "description": "Stops all safe memory expectation in the current subcontext.", + "declaration": "function stopExpectSafeMemory() external;", + "visibility": "external", + "mutability": "", + "signature": "stopExpectSafeMemory()", + "selector": "0x0956441b", + "selectorBytes": [ + 9, + 86, + 68, + 27 + ] + }, + "group": "testing", + "status": "stable", + "safety": "unsafe" + }, { "func": { "id": "stopMappingRecording", diff --git a/crates/cheatcodes/spec/src/vm.rs b/crates/cheatcodes/spec/src/vm.rs index 6ab18d10935b..370c8b27a4cf 100644 --- a/crates/cheatcodes/spec/src/vm.rs +++ b/crates/cheatcodes/spec/src/vm.rs @@ -701,6 +701,10 @@ interface Vm { #[cheatcode(group = Testing, safety = Unsafe)] function expectSafeMemory(uint64 min, uint64 max) external; + /// Stops all safe memory expectation in the current subcontext. + #[cheatcode(group = Testing, safety = Unsafe)] + function stopExpectSafeMemory() external; + /// Only allows memory writes to offsets [0x00, 0x60) ∪ [min, max) in the next created subcontext. /// If any other memory is written to, the test will fail. Can be called multiple times to add more ranges /// to the set. diff --git a/crates/cheatcodes/src/test/expect.rs b/crates/cheatcodes/src/test/expect.rs index 002a23acd28e..e4cc34986f5b 100644 --- a/crates/cheatcodes/src/test/expect.rs +++ b/crates/cheatcodes/src/test/expect.rs @@ -282,6 +282,14 @@ impl Cheatcode for expectSafeMemoryCall { } } +impl Cheatcode for stopExpectSafeMemoryCall { + fn apply_full(&self, ccx: &mut CheatsCtxt) -> Result { + let Self {} = self; + ccx.state.allowed_mem_writes.remove(&ccx.data.journaled_state.depth()); + Ok(Default::default()) + } +} + impl Cheatcode for expectSafeMemoryCallCall { fn apply_full(&self, ccx: &mut CheatsCtxt) -> Result { let Self { min, max } = *self; diff --git a/testdata/cheats/MemSafety.t.sol b/testdata/cheats/MemSafety.t.sol index 0da135d0c2a8..48205233b1e3 100644 --- a/testdata/cheats/MemSafety.t.sol +++ b/testdata/cheats/MemSafety.t.sol @@ -670,6 +670,47 @@ contract MemSafetyTest is DSTest { } } + //////////////////////////////////////////////////////////////// + // `stopExpectSafeMemory` cheatcode // + //////////////////////////////////////////////////////////////// + + /// @dev Tests that the `stopExpectSafeMemory` cheatcode works as expected. + function testStopExpectSafeMemory() public { + uint64 initPtr; + assembly { + initPtr := mload(0x40) + } + + vm.expectSafeMemory(initPtr, initPtr + 0x20); + assembly { + // write to allowed range + mstore(initPtr, 0x01) + } + + vm.stopExpectSafeMemory(); + + assembly { + // write ouside allowed range, this should be fine + mstore(add(initPtr, 0x20), 0x01) + } + } + + /// @dev Tests that the `stopExpectSafeMemory` cheatcode does not cause violations not being noticed. + function testFailStopExpectSafeMemory() public { + uint64 initPtr; + assembly { + initPtr := mload(0x40) + } + + vm.expectSafeMemory(initPtr, initPtr + 0x20); + assembly { + // write outside of allowed range, this should revert + mstore(add(initPtr, 0x20), 0x01) + } + + vm.stopExpectSafeMemory(); + } + //////////////////////////////////////////////////////////////// // HELPERS // //////////////////////////////////////////////////////////////// diff --git a/testdata/cheats/Vm.sol b/testdata/cheats/Vm.sol index bac97febb2f9..4bc85d457ce0 100644 --- a/testdata/cheats/Vm.sol +++ b/testdata/cheats/Vm.sol @@ -344,6 +344,7 @@ interface Vm { function startStateDiffRecording() external; function stopAndReturnStateDiff() external returns (AccountAccess[] memory accountAccesses); function stopBroadcast() external; + function stopExpectSafeMemory() external; function stopMappingRecording() external; function stopPrank() external; function store(address target, bytes32 slot, bytes32 value) external;