From 10748dc2d5235a5cef4e8077a3b7bf475e1ef17b Mon Sep 17 00:00:00 2001 From: grandizzy Date: Fri, 6 Sep 2024 19:50:30 +0300 Subject: [PATCH 1/5] feat(cheatcodes): additional cheatcodes to aid in symbolic testing --- crates/cheatcodes/assets/cheatcodes.json | 60 +++++++++ crates/cheatcodes/spec/src/vm.rs | 17 +++ crates/cheatcodes/src/evm.rs | 11 +- crates/cheatcodes/src/evm/mock.rs | 9 ++ crates/cheatcodes/src/inspector.rs | 37 +++++- crates/cheatcodes/src/utils.rs | 24 +++- crates/evm/evm/src/inspectors/stack.rs | 12 ++ crates/forge/tests/it/cheats.rs | 2 + testdata/cheats/Vm.sol | 3 + .../default/cheats/ArbitraryStorage.t.sol | 125 ++++++++++++++++++ testdata/default/cheats/CopyStorage.t.sol | 78 +++++++++++ testdata/default/cheats/MockFunction.t.sol | 74 +++++++++++ 12 files changed, 447 insertions(+), 5 deletions(-) create mode 100644 testdata/default/cheats/ArbitraryStorage.t.sol create mode 100644 testdata/default/cheats/CopyStorage.t.sol create mode 100644 testdata/default/cheats/MockFunction.t.sol diff --git a/crates/cheatcodes/assets/cheatcodes.json b/crates/cheatcodes/assets/cheatcodes.json index 4517f075e7fb..da501a11ed68 100644 --- a/crates/cheatcodes/assets/cheatcodes.json +++ b/crates/cheatcodes/assets/cheatcodes.json @@ -3331,6 +3331,26 @@ "status": "stable", "safety": "safe" }, + { + "func": { + "id": "copyStorage", + "description": "Utility cheatcode to copy storage of `from` contract to another `to` contract.", + "declaration": "function copyStorage(address from, address to) external;", + "visibility": "external", + "mutability": "", + "signature": "copyStorage(address,address)", + "selector": "0x203dac0d", + "selectorBytes": [ + 32, + 61, + 172, + 13 + ] + }, + "group": "utilities", + "status": "stable", + "safety": "safe" + }, { "func": { "id": "createDir", @@ -5591,6 +5611,26 @@ "status": "stable", "safety": "unsafe" }, + { + "func": { + "id": "mockFunction", + "description": "Whenever a call is made to `callee` with calldata `data`, this cheatcode instead calls\n`target` with the same calldata. This functionality is similar to a delegate call made to\n`target` contract from `callee`.\nCan be used to substitute a call to a function with another implementation that captures\nthe primary logic of the original function but is easier to reason about.\nIf calldata is not a strict match then partial match by selector is attempted.", + "declaration": "function mockFunction(address callee, address target, bytes calldata data) external;", + "visibility": "external", + "mutability": "", + "signature": "mockFunction(address,address,bytes)", + "selector": "0xadf84d21", + "selectorBytes": [ + 173, + 248, + 77, + 33 + ] + }, + "group": "evm", + "status": "stable", + "safety": "unsafe" + }, { "func": { "id": "parseAddress", @@ -7791,6 +7831,26 @@ "status": "stable", "safety": "safe" }, + { + "func": { + "id": "setArbitraryStorage", + "description": "Utility cheatcode to set arbitrary storage for given target address.", + "declaration": "function setArbitraryStorage(address target) external;", + "visibility": "external", + "mutability": "", + "signature": "setArbitraryStorage(address)", + "selector": "0xe1631837", + "selectorBytes": [ + 225, + 99, + 24, + 55 + ] + }, + "group": "utilities", + "status": "stable", + "safety": "safe" + }, { "func": { "id": "setBlockhash", diff --git a/crates/cheatcodes/spec/src/vm.rs b/crates/cheatcodes/spec/src/vm.rs index 980bab066a3a..d0a921485c7f 100644 --- a/crates/cheatcodes/spec/src/vm.rs +++ b/crates/cheatcodes/spec/src/vm.rs @@ -473,6 +473,15 @@ interface Vm { function mockCallRevert(address callee, uint256 msgValue, bytes calldata data, bytes calldata revertData) external; + /// Whenever a call is made to `callee` with calldata `data`, this cheatcode instead calls + /// `target` with the same calldata. This functionality is similar to a delegate call made to + /// `target` contract from `callee`. + /// Can be used to substitute a call to a function with another implementation that captures + /// the primary logic of the original function but is easier to reason about. + /// If calldata is not a strict match then partial match by selector is attempted. + #[cheatcode(group = Evm, safety = Unsafe)] + function mockFunction(address callee, address target, bytes calldata data) external; + // --- Impersonation (pranks) --- /// Sets the *next* call's `msg.sender` to be the input address. @@ -2303,6 +2312,14 @@ interface Vm { /// Unpauses collection of call traces. #[cheatcode(group = Utilities)] function resumeTracing() external view; + + /// Utility cheatcode to copy storage of `from` contract to another `to` contract. + #[cheatcode(group = Utilities)] + function copyStorage(address from, address to) external; + + /// Utility cheatcode to set arbitrary storage for given target address. + #[cheatcode(group = Utilities)] + function setArbitraryStorage(address target) external; } } diff --git a/crates/cheatcodes/src/evm.rs b/crates/cheatcodes/src/evm.rs index a3d517387d3d..37652ca93fa7 100644 --- a/crates/cheatcodes/src/evm.rs +++ b/crates/cheatcodes/src/evm.rs @@ -13,8 +13,9 @@ use foundry_evm_core::{ backend::{DatabaseExt, RevertSnapshotAction}, constants::{CALLER, CHEATCODE_ADDRESS, HARDHAT_CONSOLE_ADDRESS, TEST_CONTRACT_ADDRESS}, }; +use rand::Rng; use revm::{ - primitives::{Account, Bytecode, SpecId, KECCAK_EMPTY}, + primitives::{Account, Bytecode, EvmStorageSlot, SpecId, KECCAK_EMPTY}, InnerEvmContext, }; use std::{ @@ -89,7 +90,13 @@ impl Cheatcode for loadCall { let Self { target, slot } = *self; ensure_not_precompile!(&target, ccx); ccx.ecx.load_account(target)?; - let val = ccx.ecx.sload(target, slot.into())?; + let mut val = ccx.ecx.sload(target, slot.into())?; + // Generate random value if target should have arbitrary storage and storage slot untouched. + if ccx.state.arbitrary_storage.contains(&target) && val.is_cold && val.data.is_zero() { + val.data = ccx.state.rng().gen(); + let mut account = ccx.ecx.load_account(target)?; + account.storage.insert(slot.into(), EvmStorageSlot::new(val.data)); + } Ok(val.abi_encode()) } } diff --git a/crates/cheatcodes/src/evm/mock.rs b/crates/cheatcodes/src/evm/mock.rs index 1a6ffb46a51f..cd7c459b6b17 100644 --- a/crates/cheatcodes/src/evm/mock.rs +++ b/crates/cheatcodes/src/evm/mock.rs @@ -89,6 +89,15 @@ impl Cheatcode for mockCallRevert_1Call { } } +impl Cheatcode for mockFunctionCall { + fn apply(&self, state: &mut Cheatcodes) -> Result { + let Self { callee, target, data } = self; + state.mocked_functions.entry(*callee).or_default().insert(data.clone(), *target); + + Ok(Default::default()) + } +} + #[allow(clippy::ptr_arg)] // Not public API, doesn't matter fn mock_call( state: &mut Cheatcodes, diff --git a/crates/cheatcodes/src/inspector.rs b/crates/cheatcodes/src/inspector.rs index f5238d810c8b..bc51a4e4a35f 100644 --- a/crates/cheatcodes/src/inspector.rs +++ b/crates/cheatcodes/src/inspector.rs @@ -41,7 +41,7 @@ use revm::{ EOFCreateInputs, EOFCreateKind, Gas, InstructionResult, Interpreter, InterpreterAction, InterpreterResult, }, - primitives::{BlockEnv, CreateScheme, EVMError, SpecId, EOF_MAGIC_BYTES}, + primitives::{BlockEnv, CreateScheme, EVMError, EvmStorageSlot, SpecId, EOF_MAGIC_BYTES}, EvmContext, InnerEvmContext, Inspector, }; use rustc_hash::FxHashMap; @@ -320,6 +320,9 @@ pub struct Cheatcodes { // **Note**: inner must a BTreeMap because of special `Ord` impl for `MockCallDataContext` pub mocked_calls: HashMap>, + /// Mocked functions. Maps target address to be mocked to pair of (calldata, mock address). + pub mocked_functions: HashMap>, + /// Expected calls pub expected_calls: ExpectedCallTracker, /// Expected emits @@ -368,6 +371,10 @@ pub struct Cheatcodes { /// Ignored traces. pub ignored_traces: IgnoredTraces, + + /// Addresses that should have arbitrary storage generated (SLOADs return random value if + /// storage slot wasn't accessed). + pub arbitrary_storage: Vec
, } // This is not derived because calling this in `fn new` with `..Default::default()` creates a second @@ -396,6 +403,7 @@ impl Cheatcodes { recorded_account_diffs_stack: Default::default(), recorded_logs: Default::default(), mocked_calls: Default::default(), + mocked_functions: Default::default(), expected_calls: Default::default(), expected_emits: Default::default(), allowed_mem_writes: Default::default(), @@ -410,6 +418,7 @@ impl Cheatcodes { breakpoints: Default::default(), rng: Default::default(), ignored_traces: Default::default(), + arbitrary_storage: Default::default(), } } @@ -1045,7 +1054,7 @@ impl Inspector for Cheatcodes { } #[inline] - fn step_end(&mut self, interpreter: &mut Interpreter, _ecx: &mut EvmContext) { + fn step_end(&mut self, interpreter: &mut Interpreter, ecx: &mut EvmContext) { if self.gas_metering.paused { self.meter_gas_end(interpreter); } @@ -1053,6 +1062,10 @@ impl Inspector for Cheatcodes { if self.gas_metering.touched { self.meter_gas_check(interpreter); } + + if self.arbitrary_storage.contains(&interpreter.contract().target_address) { + self.ensure_arbitrary_storage(interpreter, ecx); + } } fn log(&mut self, interpreter: &mut Interpreter, _ecx: &mut EvmContext, log: &Log) { @@ -1465,6 +1478,26 @@ impl Cheatcodes { } } + /// Generates arbitrary values for storage slots. + #[cold] + fn ensure_arbitrary_storage( + &mut self, + interpreter: &mut Interpreter, + ecx: &mut EvmContext, + ) { + if interpreter.current_opcode() == op::SLOAD { + let key = try_or_return!(interpreter.stack().peek(0)); + let target_address = interpreter.contract().target_address; + if let Ok(value) = ecx.sload(target_address, key) { + if value.is_cold && value.data.is_zero() { + if let Ok(mut target_account) = ecx.load_account(target_address) { + target_account.storage.insert(key, EvmStorageSlot::new(self.rng().gen())); + } + } + } + } + } + /// Records storage slots reads and writes. #[cold] fn record_accesses(&mut self, interpreter: &mut Interpreter) { diff --git a/crates/cheatcodes/src/utils.rs b/crates/cheatcodes/src/utils.rs index 642cf83abb99..b08762945b75 100644 --- a/crates/cheatcodes/src/utils.rs +++ b/crates/cheatcodes/src/utils.rs @@ -1,6 +1,6 @@ //! Implementations of [`Utilities`](spec::Group::Utilities) cheatcodes. -use crate::{Cheatcode, Cheatcodes, Result, Vm::*}; +use crate::{Cheatcode, Cheatcodes, CheatsCtxt, Result, Vm::*}; use alloy_primitives::{Address, U256}; use alloy_sol_types::SolValue; use foundry_common::ens::namehash; @@ -149,3 +149,25 @@ impl Cheatcode for resumeTracingCall { Ok(Default::default()) } } +impl Cheatcode for setArbitraryStorageCall { + fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { + let Self { target } = self; + ccx.state.arbitrary_storage.push(*target); + + Ok(Default::default()) + } +} + +impl Cheatcode for copyStorageCall { + fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { + let Self { from, to } = self; + if let Ok(from_account) = ccx.load_account(*from) { + let from_storage = from_account.storage.clone(); + if let Ok(mut to_account) = ccx.load_account(*to) { + to_account.storage = from_storage; + } + } + + Ok(Default::default()) + } +} diff --git a/crates/evm/evm/src/inspectors/stack.rs b/crates/evm/evm/src/inspectors/stack.rs index e44d499eafa2..d39fbb8756b5 100644 --- a/crates/evm/evm/src/inspectors/stack.rs +++ b/crates/evm/evm/src/inspectors/stack.rs @@ -726,6 +726,18 @@ impl<'a, DB: DatabaseExt> Inspector for InspectorStackRefMut<'a> { ecx.journaled_state.depth += self.in_inner_context as usize; if let Some(cheatcodes) = self.cheatcodes.as_deref_mut() { + // Handle mocked functions, replace bytecode address with mock if matched. + if let Some(mocks) = cheatcodes.mocked_functions.get(&call.target_address) { + if let Some(target) = mocks.get(&call.input) { + call.bytecode_address = *target; + } else { + // Check if we have a catch-all mock set for selector. + if let Some(target) = mocks.get(&call.input.slice(..4)) { + call.bytecode_address = *target; + } + } + } + if let Some(output) = cheatcodes.call_with_executor(ecx, call, self.inner) { if output.result.result != InstructionResult::Continue { ecx.journaled_state.depth -= self.in_inner_context as usize; diff --git a/crates/forge/tests/it/cheats.rs b/crates/forge/tests/it/cheats.rs index 2bbbee90289c..19c78c973d89 100644 --- a/crates/forge/tests/it/cheats.rs +++ b/crates/forge/tests/it/cheats.rs @@ -7,6 +7,7 @@ use crate::{ TEST_DATA_MULTI_VERSION, }, }; +use alloy_primitives::U256; use foundry_config::{fs_permissions::PathPermission, FsPermissions}; use foundry_test_utils::Filter; @@ -26,6 +27,7 @@ async fn test_cheats_local(test_data: &ForgeTestData) { } let mut config = test_data.config.clone(); + config.fuzz.seed = Some(U256::from(100)); config.fs_permissions = FsPermissions::new(vec![PathPermission::read_write("./")]); let runner = test_data.runner_with_config(config); diff --git a/testdata/cheats/Vm.sol b/testdata/cheats/Vm.sol index 5b6750237add..edc7dad3249e 100644 --- a/testdata/cheats/Vm.sol +++ b/testdata/cheats/Vm.sol @@ -162,6 +162,7 @@ interface Vm { function computeCreateAddress(address deployer, uint256 nonce) external pure returns (address); function cool(address target) external; function copyFile(string calldata from, string calldata to) external returns (uint64 copied); + function copyStorage(address from, address to) external; function createDir(string calldata path, bool recursive) external; function createFork(string calldata urlOrAlias) external returns (uint256 forkId); function createFork(string calldata urlOrAlias, uint256 blockNumber) external returns (uint256 forkId); @@ -275,6 +276,7 @@ interface Vm { function mockCallRevert(address callee, uint256 msgValue, bytes calldata data, bytes calldata revertData) external; function mockCall(address callee, bytes calldata data, bytes calldata returnData) external; function mockCall(address callee, uint256 msgValue, bytes calldata data, bytes calldata returnData) external; + function mockFunction(address callee, address target, bytes calldata data) external; function parseAddress(string calldata stringifiedValue) external pure returns (address parsedValue); function parseBool(string calldata stringifiedValue) external pure returns (bool parsedValue); function parseBytes(string calldata stringifiedValue) external pure returns (bytes memory parsedValue); @@ -385,6 +387,7 @@ interface Vm { function serializeUintToHex(string calldata objectKey, string calldata valueKey, uint256 value) external returns (string memory json); function serializeUint(string calldata objectKey, string calldata valueKey, uint256 value) external returns (string memory json); function serializeUint(string calldata objectKey, string calldata valueKey, uint256[] calldata values) external returns (string memory json); + function setArbitraryStorage(address target) external; function setBlockhash(uint256 blockNumber, bytes32 blockHash) external; function setEnv(string calldata name, string calldata value) external; function setNonce(address account, uint64 newNonce) external; diff --git a/testdata/default/cheats/ArbitraryStorage.t.sol b/testdata/default/cheats/ArbitraryStorage.t.sol new file mode 100644 index 000000000000..325bdbde6f98 --- /dev/null +++ b/testdata/default/cheats/ArbitraryStorage.t.sol @@ -0,0 +1,125 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 +pragma solidity 0.8.18; + +import "ds-test/test.sol"; +import "cheats/Vm.sol"; + +contract Counter { + uint256 public a; + address public b; + int8 public c; + address[] public owners; + + function setA(uint256 _a) public { + a = _a; + } + + function setB(address _b) public { + b = _b; + } + + function getOwner(uint256 pos) public view returns (address) { + return owners[pos]; + } + + function setOwner(uint256 pos, address owner) public { + owners[pos] = owner; + } +} + +contract CounterArbitraryStorageWithSeedTest is DSTest { + Vm vm = Vm(HEVM_ADDRESS); + + function test_fresh_storage() public { + uint256 index = 55; + Counter counter = new Counter(); + vm.setArbitraryStorage(address(counter)); + // Next call would fail with array out of bounds without arbitrary storage. + address owner = counter.getOwner(index); + // Subsequent calls should retrieve same value + assertEq(counter.getOwner(index), owner); + // Change slot and make sure new value retrieved + counter.setOwner(index, address(111)); + assertEq(counter.getOwner(index), address(111)); + } + + function test_arbitrary_storage_warm() public { + Counter counter = new Counter(); + vm.setArbitraryStorage(address(counter)); + assertGt(counter.a(), 0); + counter.setA(0); + // This should remain 0 if explicitly set. + assertEq(counter.a(), 0); + counter.setA(11); + assertEq(counter.a(), 11); + } + + function test_arbitrary_storage_multiple_read_writes() public { + Counter counter = new Counter(); + vm.setArbitraryStorage(address(counter)); + uint256 slot1 = vm.randomUint(0, 100); + uint256 slot2 = vm.randomUint(0, 100); + require(slot1 != slot2, "random positions should be different"); + address alice = counter.owners(slot1); + address bob = counter.owners(slot2); + require(alice != bob, "random storage values should be different"); + counter.setOwner(slot1, bob); + counter.setOwner(slot2, alice); + assertEq(alice, counter.owners(slot2)); + assertEq(bob, counter.owners(slot1)); + } +} + +contract AContract { + uint256[] public a; + address[] public b; + int8[] public c; + bytes32[] public d; +} + +contract AContractArbitraryStorageTest is DSTest { + Vm vm = Vm(HEVM_ADDRESS); + + function test_arbitrary_storage_with_seed() public { + AContract target = new AContract(); + vm.setArbitraryStorage(address(target)); + assertEq(target.a(11), 85286582241781868037363115933978803127245343755841464083427462398552335014708); + assertEq(target.b(22), 0x939180Daa938F9e18Ff0E76c112D25107D358B02); + assertEq(target.c(33), -104); + assertEq(target.d(44), 0x6c178fa9c434f142df61a5355cc2b8d07be691b98dabf5b1a924f2bce97a19c7); + } +} + +contract SymbolicStore { + uint256 public testNumber = 1337; // slot 0 + + constructor() {} +} + +contract SymbolicStorageTest is DSTest { + Vm vm = Vm(HEVM_ADDRESS); + + function test_SymbolicStorage() public { + uint256 slot = vm.randomUint(0, 100); + address addr = 0xEA674fdDe714fd979de3EdF0F56AA9716B898ec8; + vm.setArbitraryStorage(addr); + bytes32 value = vm.load(addr, bytes32(slot)); + assertEq(uint256(value), 85286582241781868037363115933978803127245343755841464083427462398552335014708); + // Load slot again and make sure we get same value. + bytes32 value1 = vm.load(addr, bytes32(slot)); + assertEq(uint256(value), uint256(value1)); + } + + function test_SymbolicStorage1() public { + uint256 slot = vm.randomUint(0, 100); + SymbolicStore myStore = new SymbolicStore(); + vm.setArbitraryStorage(address(myStore)); + bytes32 value = vm.load(address(myStore), bytes32(uint256(slot))); + assertEq(uint256(value), 85286582241781868037363115933978803127245343755841464083427462398552335014708); + } + + function testEmptyInitialStorage(uint256 slot) public { + bytes32 storage_value = vm.load(address(vm), bytes32(slot)); + assertEq(uint256(storage_value), 0); + } +} diff --git a/testdata/default/cheats/CopyStorage.t.sol b/testdata/default/cheats/CopyStorage.t.sol new file mode 100644 index 000000000000..c497b342f323 --- /dev/null +++ b/testdata/default/cheats/CopyStorage.t.sol @@ -0,0 +1,78 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 +pragma solidity 0.8.18; + +import "ds-test/test.sol"; +import "cheats/Vm.sol"; + +contract Counter { + uint256 public a; + address public b; + + function setA(uint256 _a) public { + a = _a; + } + + function setB(address _b) public { + b = _b; + } +} + +contract CounterTest is DSTest { + Counter public counter; + Counter public counter1; + Vm vm = Vm(HEVM_ADDRESS); + + function setUp() public { + counter = new Counter(); + counter.setA(1000); + counter.setB(address(27)); + counter1 = new Counter(); + counter1.setA(11); + counter1.setB(address(50)); + } + + function test_copy_storage() public { + assertEq(counter.a(), 1000); + assertEq(counter.b(), address(27)); + assertEq(counter1.a(), 11); + assertEq(counter1.b(), address(50)); + vm.copyStorage(address(counter), address(counter1)); + assertEq(counter.a(), 1000); + assertEq(counter.b(), address(27)); + assertEq(counter1.a(), 1000); + assertEq(counter1.b(), address(27)); + } +} + +contract CopyStorageContract { + uint256 public x; +} + +contract CopyStorageTest is DSTest { + CopyStorageContract csc_1; + CopyStorageContract csc_2; + Vm vm = Vm(HEVM_ADDRESS); + + function _storeUInt256(address contractAddress, uint256 slot, uint256 value) internal { + vm.store(contractAddress, bytes32(slot), bytes32(value)); + } + + function setUp() public { + csc_1 = new CopyStorageContract(); + csc_2 = new CopyStorageContract(); + } + + function test_copy_storage() public { + // Make the storage of first contract symbolic + vm.setArbitraryStorage(address(csc_1)); + // and explicitly put a constrained symbolic value into the slot for `x` + uint256 x_1 = vm.randomUint(); + _storeUInt256(address(csc_1), 0, x_1); + // `x` of second contract is uninitialized + assert(csc_2.x() == 0); + // Copy storage from first to second contract + vm.copyStorage(address(csc_1), address(csc_2)); + // `x` of second contract is now the `x` of the first + assert(csc_2.x() == x_1); + } +} diff --git a/testdata/default/cheats/MockFunction.t.sol b/testdata/default/cheats/MockFunction.t.sol new file mode 100644 index 000000000000..9cf1004ca279 --- /dev/null +++ b/testdata/default/cheats/MockFunction.t.sol @@ -0,0 +1,74 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 +pragma solidity 0.8.18; + +import "ds-test/test.sol"; +import "cheats/Vm.sol"; + +contract MockFunctionContract { + uint256 public a; + + function mocked_function() public { + a = 321; + } + + function mocked_args_function(uint256 x) public { + a = 321 + x; + } +} + +contract ModelMockFunctionContract { + uint256 public a; + + function mocked_function() public { + a = 123; + } + + function mocked_args_function(uint256 x) public { + a = 123 + x; + } +} + +contract MockFunctionTest is DSTest { + MockFunctionContract my_contract; + ModelMockFunctionContract model_contract; + Vm vm = Vm(HEVM_ADDRESS); + + function setUp() public { + my_contract = new MockFunctionContract(); + model_contract = new ModelMockFunctionContract(); + } + + function test_mock_function() public { + vm.mockFunction( + address(my_contract), + address(model_contract), + abi.encodeWithSelector(MockFunctionContract.mocked_function.selector) + ); + my_contract.mocked_function(); + assertEq(my_contract.a(), 123); + } + + function test_mock_function_concrete_args() public { + vm.mockFunction( + address(my_contract), + address(model_contract), + abi.encodeWithSelector(MockFunctionContract.mocked_args_function.selector, 456) + ); + my_contract.mocked_args_function(456); + assertEq(my_contract.a(), 123 + 456); + my_contract.mocked_args_function(567); + assertEq(my_contract.a(), 321 + 567); + } + + function test_mock_function_all_args() public { + vm.mockFunction( + address(my_contract), + address(model_contract), + abi.encodeWithSelector(MockFunctionContract.mocked_args_function.selector) + ); + my_contract.mocked_args_function(678); + assertEq(my_contract.a(), 123 + 678); + my_contract.mocked_args_function(789); + assertEq(my_contract.a(), 123 + 789); + } +} From 16f068f6e9f557571464b9ad5d39d9f24ff40bb5 Mon Sep 17 00:00:00 2001 From: grandizzy Date: Tue, 10 Sep 2024 08:46:24 +0300 Subject: [PATCH 2/5] Support copies from arbitrary storage, docs --- crates/cheatcodes/src/evm.rs | 23 +++- crates/cheatcodes/src/inspector.rs | 140 ++++++++++++++++++++-- crates/cheatcodes/src/utils.rs | 7 +- testdata/default/cheats/CopyStorage.t.sol | 86 ++++++++++++- 4 files changed, 239 insertions(+), 17 deletions(-) diff --git a/crates/cheatcodes/src/evm.rs b/crates/cheatcodes/src/evm.rs index 37652ca93fa7..dc4fe2ba5ac0 100644 --- a/crates/cheatcodes/src/evm.rs +++ b/crates/cheatcodes/src/evm.rs @@ -15,7 +15,7 @@ use foundry_evm_core::{ }; use rand::Rng; use revm::{ - primitives::{Account, Bytecode, EvmStorageSlot, SpecId, KECCAK_EMPTY}, + primitives::{Account, Bytecode, SpecId, KECCAK_EMPTY}, InnerEvmContext, }; use std::{ @@ -91,12 +91,23 @@ impl Cheatcode for loadCall { ensure_not_precompile!(&target, ccx); ccx.ecx.load_account(target)?; let mut val = ccx.ecx.sload(target, slot.into())?; - // Generate random value if target should have arbitrary storage and storage slot untouched. - if ccx.state.arbitrary_storage.contains(&target) && val.is_cold && val.data.is_zero() { - val.data = ccx.state.rng().gen(); - let mut account = ccx.ecx.load_account(target)?; - account.storage.insert(slot.into(), EvmStorageSlot::new(val.data)); + + if val.is_cold && val.data.is_zero() { + let rand_value = ccx.state.rng().gen(); + let arbitrary_storage = &mut ccx.state.arbitrary_storage; + if arbitrary_storage.is_arbitrary(&target) { + // If storage slot is untouched and load from a target with arbitrary storage, + // then set random value for current slot. + arbitrary_storage.save(ccx.ecx, target, slot.into(), rand_value); + val.data = rand_value; + } else if arbitrary_storage.is_copy(&target) { + // If storage slot is untouched and load from a target that copies storage from + // a source address with arbitrary storage, then copy existing arbitrary value. + // If no arbitrary value generated yet, then the random one is saved and set. + val.data = arbitrary_storage.copy(ccx.ecx, target, slot.into(), rand_value); + } } + Ok(val.abi_encode()) } } diff --git a/crates/cheatcodes/src/inspector.rs b/crates/cheatcodes/src/inspector.rs index bc51a4e4a35f..213f8e5f4c40 100644 --- a/crates/cheatcodes/src/inspector.rs +++ b/crates/cheatcodes/src/inspector.rs @@ -254,6 +254,92 @@ impl GasMetering { } } +/// Holds data about arbitrary storage. +#[derive(Clone, Debug, Default)] +pub struct ArbitraryStorage { + /// Mapping of arbitrary storage addresses to generated values (slot, arbitrary value). + /// (SLOADs return random value if storage slot wasn't accessed). + /// Changed values are recorded and used to copy storage to different addresses. + pub values: HashMap>, + /// Mapping of address with storage copied to arbitrary storage address source. + pub copies: HashMap, +} + +impl ArbitraryStorage { + /// Whether the given address has arbitrary storage. + pub fn is_arbitrary(&self, address: &Address) -> bool { + self.values.contains_key(address) + } + + /// Whether the given address is a copy of an address with arbitrary storage. + pub fn is_copy(&self, address: &Address) -> bool { + self.copies.contains_key(address) + } + + /// Marks an address with arbitrary storage. + pub fn mark_arbitrary(&mut self, address: &Address) { + self.values.insert(*address, HashMap::default()); + } + + /// Maps an address that copies storage with the arbitrary storage address. + pub fn mark_copy(&mut self, from: &Address, to: &Address) { + if self.is_arbitrary(from) { + self.copies.insert(*to, *from); + } + } + + /// Saves arbitrary storage value for a given address: + /// - store value in changed values cache. + /// - update account's storage with given value. + pub fn save( + &mut self, + ecx: &mut InnerEvmContext, + address: Address, + slot: U256, + data: U256, + ) { + if let Ok(mut account) = ecx.load_account(address) { + self.values + .get_mut(&address) + .expect("missing arbitrary address entry") + .insert(slot, data); + account.storage.insert(slot, EvmStorageSlot::new(data)); + } + } + + /// Copies arbitrary storage value from source address to the given target address: + /// - if a value is present in arbitrary values cache, then update target storage and return + /// existing value. + /// - if no value was yet generated for given slot, then save new value in cache and update both + /// source and target storages. + pub fn copy( + &mut self, + ecx: &mut InnerEvmContext, + target: Address, + slot: U256, + new_value: U256, + ) -> U256 { + let source = self.copies.get(&target).expect("missing arbitrary copy target entry"); + let storage_cache = self.values.get_mut(source).expect("missing arbitrary source storage"); + let value = match storage_cache.get(&slot) { + Some(value) => *value, + None => { + storage_cache.insert(slot, new_value); + // Update source storage with new value. + if let Ok(mut source_account) = ecx.load_account(*source) { + source_account.storage.insert(slot, EvmStorageSlot::new(new_value)); + } + new_value + } + }; + // Update target storage with new value. + if let Ok(mut target_account) = ecx.load_account(target) { + target_account.storage.insert(slot, EvmStorageSlot::new(value)); + } + value + } +} + /// List of transactions that can be broadcasted. pub type BroadcastableTransactions = VecDeque; @@ -372,9 +458,8 @@ pub struct Cheatcodes { /// Ignored traces. pub ignored_traces: IgnoredTraces, - /// Addresses that should have arbitrary storage generated (SLOADs return random value if - /// storage slot wasn't accessed). - pub arbitrary_storage: Vec
, + /// Addresses with arbitrary storage. + pub arbitrary_storage: ArbitraryStorage, } // This is not derived because calling this in `fn new` with `..Default::default()` creates a second @@ -1063,9 +1148,13 @@ impl Inspector for Cheatcodes { self.meter_gas_check(interpreter); } - if self.arbitrary_storage.contains(&interpreter.contract().target_address) { + if self.arbitrary_storage.is_arbitrary(&interpreter.contract().target_address) { self.ensure_arbitrary_storage(interpreter, ecx); } + + if self.arbitrary_storage.is_copy(&interpreter.contract().target_address) { + self.copy_arbitrary_storage(interpreter, ecx); + } } fn log(&mut self, interpreter: &mut Interpreter, _ecx: &mut EvmContext, log: &Log) { @@ -1479,6 +1568,10 @@ impl Cheatcodes { } /// Generates arbitrary values for storage slots. + /// Invoked in inspector `step_end`, when the current opcode is not executed. + /// If current opcode to execute is `SLOAD` and storage slot is cold, then an arbitrary value + /// is generated and saved in target address storage (therefore when `SLOAD` opcode is executed, + /// the arbitrary value will be returned). #[cold] fn ensure_arbitrary_storage( &mut self, @@ -1490,9 +1583,42 @@ impl Cheatcodes { let target_address = interpreter.contract().target_address; if let Ok(value) = ecx.sload(target_address, key) { if value.is_cold && value.data.is_zero() { - if let Ok(mut target_account) = ecx.load_account(target_address) { - target_account.storage.insert(key, EvmStorageSlot::new(self.rng().gen())); - } + let arbitrary_value = self.rng().gen(); + self.arbitrary_storage.save( + &mut ecx.inner, + target_address, + key, + arbitrary_value, + ); + } + } + } + } + + /// Copies arbitrary values for storage slots. + /// Invoked in inspector `step_end`, when the current opcode is not executed. + /// If current opcode to execute is `SLOAD` and storage slot is cold, it copies the existing + /// arbitrary storage value (or the new generated one if no value in cache) from mapped source + /// address to the target address (therefore when `SLOAD` opcode is executed, the arbitrary + /// value will be returned). + #[cold] + fn copy_arbitrary_storage( + &mut self, + interpreter: &mut Interpreter, + ecx: &mut EvmContext, + ) { + if interpreter.current_opcode() == op::SLOAD { + let key = try_or_return!(interpreter.stack().peek(0)); + let target_address = interpreter.contract().target_address; + if let Ok(value) = ecx.sload(target_address, key) { + if value.is_cold && value.data.is_zero() { + let arbitrary_value = self.rng().gen(); + self.arbitrary_storage.copy( + &mut ecx.inner, + target_address, + key, + arbitrary_value, + ); } } } diff --git a/crates/cheatcodes/src/utils.rs b/crates/cheatcodes/src/utils.rs index b08762945b75..d4a1681f4dfc 100644 --- a/crates/cheatcodes/src/utils.rs +++ b/crates/cheatcodes/src/utils.rs @@ -152,7 +152,7 @@ impl Cheatcode for resumeTracingCall { impl Cheatcode for setArbitraryStorageCall { fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { let Self { target } = self; - ccx.state.arbitrary_storage.push(*target); + ccx.state.arbitrary_storage.mark_arbitrary(target); Ok(Default::default()) } @@ -161,10 +161,15 @@ impl Cheatcode for setArbitraryStorageCall { impl Cheatcode for copyStorageCall { fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { let Self { from, to } = self; + ensure!( + !ccx.state.arbitrary_storage.is_arbitrary(to), + "target address cannot have arbitrary storage" + ); if let Ok(from_account) = ccx.load_account(*from) { let from_storage = from_account.storage.clone(); if let Ok(mut to_account) = ccx.load_account(*to) { to_account.storage = from_storage; + ccx.state.arbitrary_storage.mark_copy(from, to); } } diff --git a/testdata/default/cheats/CopyStorage.t.sol b/testdata/default/cheats/CopyStorage.t.sol index c497b342f323..7d10ab2c3756 100644 --- a/testdata/default/cheats/CopyStorage.t.sol +++ b/testdata/default/cheats/CopyStorage.t.sol @@ -7,6 +7,7 @@ import "cheats/Vm.sol"; contract Counter { uint256 public a; address public b; + int256[] public c; function setA(uint256 _a) public { a = _a; @@ -22,16 +23,14 @@ contract CounterTest is DSTest { Counter public counter1; Vm vm = Vm(HEVM_ADDRESS); - function setUp() public { + function test_copy_storage() public { counter = new Counter(); counter.setA(1000); counter.setB(address(27)); counter1 = new Counter(); counter1.setA(11); counter1.setB(address(50)); - } - function test_copy_storage() public { assertEq(counter.a(), 1000); assertEq(counter.b(), address(27)); assertEq(counter1.a(), 11); @@ -42,6 +41,26 @@ contract CounterTest is DSTest { assertEq(counter1.a(), 1000); assertEq(counter1.b(), address(27)); } + + function test_copy_storage_from_arbitrary() public { + counter = new Counter(); + counter1 = new Counter(); + vm.setArbitraryStorage(address(counter)); + vm.copyStorage(address(counter), address(counter1)); + + // Make sure untouched storage has same values. + assertEq(counter.a(), counter1.a()); + assertEq(counter.b(), counter1.b()); + assertEq(counter.c(33), counter1.c(33)); + + // Change storage in source storage contract and make sure copy is not changed. + counter.setA(1000); + counter1.setB(address(50)); + assertEq(counter.a(), 1000); + assertEq(counter1.a(), 40426841063417815470953489044557166618267862781491517122018165313568904172524); + assertEq(counter.b(), 0x485E9Cc0ef187E54A3AB45b50c3DcE43f2C223B1); + assertEq(counter1.b(), address(50)); + } } contract CopyStorageContract { @@ -51,6 +70,7 @@ contract CopyStorageContract { contract CopyStorageTest is DSTest { CopyStorageContract csc_1; CopyStorageContract csc_2; + CopyStorageContract csc_3; Vm vm = Vm(HEVM_ADDRESS); function _storeUInt256(address contractAddress, uint256 slot, uint256 value) internal { @@ -60,6 +80,7 @@ contract CopyStorageTest is DSTest { function setUp() public { csc_1 = new CopyStorageContract(); csc_2 = new CopyStorageContract(); + csc_3 = new CopyStorageContract(); } function test_copy_storage() public { @@ -75,4 +96,63 @@ contract CopyStorageTest is DSTest { // `x` of second contract is now the `x` of the first assert(csc_2.x() == x_1); } + + function test_copy_storage_same_values_on_load() public { + // Make the storage of first contract symbolic + vm.setArbitraryStorage(address(csc_1)); + vm.copyStorage(address(csc_1), address(csc_2)); + uint256 slot1 = vm.randomUint(0, 100); + uint256 slot2 = vm.randomUint(0, 100); + bytes32 value1 = vm.load(address(csc_1), bytes32(slot1)); + bytes32 value2 = vm.load(address(csc_1), bytes32(slot2)); + + bytes32 value3 = vm.load(address(csc_2), bytes32(slot1)); + bytes32 value4 = vm.load(address(csc_2), bytes32(slot2)); + + // Check storage values are the same for both source and target contracts. + assertEq(value1, value3); + assertEq(value2, value4); + } + + function test_copy_storage_consistent_values() public { + // Make the storage of first contract symbolic. + vm.setArbitraryStorage(address(csc_1)); + // Copy arbitrary storage to 2 contracts. + vm.copyStorage(address(csc_1), address(csc_2)); + vm.copyStorage(address(csc_1), address(csc_3)); + uint256 slot1 = vm.randomUint(0, 100); + uint256 slot2 = vm.randomUint(0, 100); + + // Load slot 1 from 1st copied contract and slot2 from symbolic contract. + bytes32 value3 = vm.load(address(csc_2), bytes32(slot1)); + bytes32 value2 = vm.load(address(csc_1), bytes32(slot2)); + + bytes32 value1 = vm.load(address(csc_1), bytes32(slot1)); + bytes32 value4 = vm.load(address(csc_2), bytes32(slot2)); + + // Make sure same values for both copied and symbolic contract. + assertEq(value3, value1); + assertEq(value2, value4); + + uint256 x_1 = vm.randomUint(); + // Change slot1 of 1st copied contract. + _storeUInt256(address(csc_2), slot1, x_1); + value3 = vm.load(address(csc_2), bytes32(slot1)); + bytes32 value5 = vm.load(address(csc_3), bytes32(slot1)); + // Make sure value for 1st contract copied is different than symbolic contract value. + assert(value3 != value1); + // Make sure same values for 2nd contract copied and symbolic contract. + assertEq(value5, value1); + + uint256 x_2 = vm.randomUint(); + // Change slot2 of symbolic contract. + _storeUInt256(address(csc_1), slot2, x_2); + value2 = vm.load(address(csc_1), bytes32(slot2)); + bytes32 value6 = vm.load(address(csc_3), bytes32(slot2)); + // Make sure value for symbolic contract value is different than 1st contract copied. + assert(value2 != value4); + // Make sure value for symbolic contract value is different than 2nd contract copied. + assert(value2 != value6); + assertEq(value4, value6); + } } From 7e69723609d13ceb3ce744c0f7bbdd3e156f1296 Mon Sep 17 00:00:00 2001 From: grandizzy Date: Wed, 11 Sep 2024 08:03:50 +0300 Subject: [PATCH 3/5] Changes after review: - separate cheatcodes tests with specific seed - better way to match mocked function - arbitrary_storage_end instead multiple calls - generate arbitrary value only when needed --- crates/cheatcodes/src/evm.rs | 13 ++-- crates/cheatcodes/src/inspector.rs | 71 +++++++------------ crates/evm/evm/src/inspectors/stack.rs | 12 ++-- crates/forge/tests/it/cheats.rs | 24 +++++-- .../default/cheats/ArbitraryStorage.t.sol | 4 +- testdata/default/cheats/CopyStorage.t.sol | 2 +- 6 files changed, 60 insertions(+), 66 deletions(-) diff --git a/crates/cheatcodes/src/evm.rs b/crates/cheatcodes/src/evm.rs index dc4fe2ba5ac0..703d7db126e9 100644 --- a/crates/cheatcodes/src/evm.rs +++ b/crates/cheatcodes/src/evm.rs @@ -93,18 +93,19 @@ impl Cheatcode for loadCall { let mut val = ccx.ecx.sload(target, slot.into())?; if val.is_cold && val.data.is_zero() { - let rand_value = ccx.state.rng().gen(); - let arbitrary_storage = &mut ccx.state.arbitrary_storage; - if arbitrary_storage.is_arbitrary(&target) { + if ccx.state.arbitrary_storage.is_arbitrary(&target) { // If storage slot is untouched and load from a target with arbitrary storage, // then set random value for current slot. - arbitrary_storage.save(ccx.ecx, target, slot.into(), rand_value); + let rand_value = ccx.state.rng().gen(); + ccx.state.arbitrary_storage.save(ccx.ecx, target, slot.into(), rand_value); val.data = rand_value; - } else if arbitrary_storage.is_copy(&target) { + } else if ccx.state.arbitrary_storage.is_copy(&target) { // If storage slot is untouched and load from a target that copies storage from // a source address with arbitrary storage, then copy existing arbitrary value. // If no arbitrary value generated yet, then the random one is saved and set. - val.data = arbitrary_storage.copy(ccx.ecx, target, slot.into(), rand_value); + let rand_value = ccx.state.rng().gen(); + val.data = + ccx.state.arbitrary_storage.copy(ccx.ecx, target, slot.into(), rand_value); } } diff --git a/crates/cheatcodes/src/inspector.rs b/crates/cheatcodes/src/inspector.rs index 213f8e5f4c40..e09e1e8c88b9 100644 --- a/crates/cheatcodes/src/inspector.rs +++ b/crates/cheatcodes/src/inspector.rs @@ -298,11 +298,8 @@ impl ArbitraryStorage { slot: U256, data: U256, ) { + self.values.get_mut(&address).expect("missing arbitrary address entry").insert(slot, data); if let Ok(mut account) = ecx.load_account(address) { - self.values - .get_mut(&address) - .expect("missing arbitrary address entry") - .insert(slot, data); account.storage.insert(slot, EvmStorageSlot::new(data)); } } @@ -1148,12 +1145,12 @@ impl Inspector for Cheatcodes { self.meter_gas_check(interpreter); } - if self.arbitrary_storage.is_arbitrary(&interpreter.contract().target_address) { - self.ensure_arbitrary_storage(interpreter, ecx); - } - - if self.arbitrary_storage.is_copy(&interpreter.contract().target_address) { - self.copy_arbitrary_storage(interpreter, ecx); + // `setArbitraryStorage` and `copyStorage`: add arbitrary values to storage. + if (self.arbitrary_storage.is_arbitrary(&interpreter.contract().target_address) || + self.arbitrary_storage.is_copy(&interpreter.contract().target_address)) && + interpreter.current_opcode() == op::SLOAD + { + self.arbitrary_storage_end(interpreter, ecx); } } @@ -1567,53 +1564,33 @@ impl Cheatcodes { } } - /// Generates arbitrary values for storage slots. - /// Invoked in inspector `step_end`, when the current opcode is not executed. - /// If current opcode to execute is `SLOAD` and storage slot is cold, then an arbitrary value - /// is generated and saved in target address storage (therefore when `SLOAD` opcode is executed, - /// the arbitrary value will be returned). + /// Generates or copies arbitrary values for storage slots. + /// Invoked in inspector `step_end` (when the current opcode is not executed), if current opcode + /// to execute is `SLOAD` and storage slot is cold. + /// Ensures that in next step (when `SLOAD` opcode is executed) an arbitrary value is returned: + /// - copies the existing arbitrary storage value (or the new generated one if no value in + /// cache) from mapped source address to the target address. + /// - generates arbitrary value and saves it in target address storage. #[cold] - fn ensure_arbitrary_storage( + fn arbitrary_storage_end( &mut self, interpreter: &mut Interpreter, ecx: &mut EvmContext, ) { - if interpreter.current_opcode() == op::SLOAD { - let key = try_or_return!(interpreter.stack().peek(0)); - let target_address = interpreter.contract().target_address; - if let Ok(value) = ecx.sload(target_address, key) { - if value.is_cold && value.data.is_zero() { - let arbitrary_value = self.rng().gen(); - self.arbitrary_storage.save( + let key = try_or_return!(interpreter.stack().peek(0)); + let target_address = interpreter.contract().target_address; + if let Ok(value) = ecx.sload(target_address, key) { + if value.is_cold && value.data.is_zero() { + let arbitrary_value = self.rng().gen(); + if self.arbitrary_storage.is_copy(&target_address) { + self.arbitrary_storage.copy( &mut ecx.inner, target_address, key, arbitrary_value, ); - } - } - } - } - - /// Copies arbitrary values for storage slots. - /// Invoked in inspector `step_end`, when the current opcode is not executed. - /// If current opcode to execute is `SLOAD` and storage slot is cold, it copies the existing - /// arbitrary storage value (or the new generated one if no value in cache) from mapped source - /// address to the target address (therefore when `SLOAD` opcode is executed, the arbitrary - /// value will be returned). - #[cold] - fn copy_arbitrary_storage( - &mut self, - interpreter: &mut Interpreter, - ecx: &mut EvmContext, - ) { - if interpreter.current_opcode() == op::SLOAD { - let key = try_or_return!(interpreter.stack().peek(0)); - let target_address = interpreter.contract().target_address; - if let Ok(value) = ecx.sload(target_address, key) { - if value.is_cold && value.data.is_zero() { - let arbitrary_value = self.rng().gen(); - self.arbitrary_storage.copy( + } else { + self.arbitrary_storage.save( &mut ecx.inner, target_address, key, diff --git a/crates/evm/evm/src/inspectors/stack.rs b/crates/evm/evm/src/inspectors/stack.rs index d39fbb8756b5..3df8dc8f01da 100644 --- a/crates/evm/evm/src/inspectors/stack.rs +++ b/crates/evm/evm/src/inspectors/stack.rs @@ -728,13 +728,13 @@ impl<'a, DB: DatabaseExt> Inspector for InspectorStackRefMut<'a> { if let Some(cheatcodes) = self.cheatcodes.as_deref_mut() { // Handle mocked functions, replace bytecode address with mock if matched. if let Some(mocks) = cheatcodes.mocked_functions.get(&call.target_address) { - if let Some(target) = mocks.get(&call.input) { + // Check if any mock function set for call data or if catch-all mock function set + // for selector. + if let Some(target) = mocks + .get(&call.input) + .or_else(|| call.input.get(..4).and_then(|selector| mocks.get(selector))) + { call.bytecode_address = *target; - } else { - // Check if we have a catch-all mock set for selector. - if let Some(target) = mocks.get(&call.input.slice(..4)) { - call.bytecode_address = *target; - } } } diff --git a/crates/forge/tests/it/cheats.rs b/crates/forge/tests/it/cheats.rs index 19c78c973d89..4bbeb9fe3d89 100644 --- a/crates/forge/tests/it/cheats.rs +++ b/crates/forge/tests/it/cheats.rs @@ -11,11 +11,12 @@ use alloy_primitives::U256; use foundry_config::{fs_permissions::PathPermission, FsPermissions}; use foundry_test_utils::Filter; -/// Executes all cheat code tests but not fork cheat codes or tests that require isolation mode +/// Executes all cheat code tests but not fork cheat codes or tests that require isolation mode or +/// specific seed. async fn test_cheats_local(test_data: &ForgeTestData) { let mut filter = Filter::new(".*", ".*", &format!(".*cheats{RE_PATH_SEPARATOR}*")) .exclude_paths("Fork") - .exclude_contracts("Isolated"); + .exclude_contracts("(Isolated|WithSeed)"); // Exclude FFI tests on Windows because no `echo`, and file tests that expect certain file paths if cfg!(windows) { @@ -27,14 +28,13 @@ async fn test_cheats_local(test_data: &ForgeTestData) { } let mut config = test_data.config.clone(); - config.fuzz.seed = Some(U256::from(100)); config.fs_permissions = FsPermissions::new(vec![PathPermission::read_write("./")]); let runner = test_data.runner_with_config(config); TestConfig::with_filter(runner, filter).run().await; } -/// Executes subset of all cheat code tests in isolation mode +/// Executes subset of all cheat code tests in isolation mode. async fn test_cheats_local_isolated(test_data: &ForgeTestData) { let filter = Filter::new(".*", ".*(Isolated)", &format!(".*cheats{RE_PATH_SEPARATOR}*")); @@ -45,6 +45,17 @@ async fn test_cheats_local_isolated(test_data: &ForgeTestData) { TestConfig::with_filter(runner, filter).run().await; } +/// Executes subset of all cheat code tests using a specific seed. +async fn test_cheats_local_with_seed(test_data: &ForgeTestData) { + let filter = Filter::new(".*", ".*(WithSeed)", &format!(".*cheats{RE_PATH_SEPARATOR}*")); + + let mut config = test_data.config.clone(); + config.fuzz.seed = Some(U256::from(100)); + let runner = test_data.runner_with_config(config); + + TestConfig::with_filter(runner, filter).run().await; +} + #[tokio::test(flavor = "multi_thread")] async fn test_cheats_local_default() { test_cheats_local(&TEST_DATA_DEFAULT).await @@ -55,6 +66,11 @@ async fn test_cheats_local_default_isolated() { test_cheats_local_isolated(&TEST_DATA_DEFAULT).await } +#[tokio::test(flavor = "multi_thread")] +async fn test_cheats_local_default_with_seed() { + test_cheats_local_with_seed(&TEST_DATA_DEFAULT).await +} + #[tokio::test(flavor = "multi_thread")] async fn test_cheats_local_multi_version() { test_cheats_local(&TEST_DATA_MULTI_VERSION).await diff --git a/testdata/default/cheats/ArbitraryStorage.t.sol b/testdata/default/cheats/ArbitraryStorage.t.sol index 325bdbde6f98..86910279e98e 100644 --- a/testdata/default/cheats/ArbitraryStorage.t.sol +++ b/testdata/default/cheats/ArbitraryStorage.t.sol @@ -77,7 +77,7 @@ contract AContract { bytes32[] public d; } -contract AContractArbitraryStorageTest is DSTest { +contract AContractArbitraryStorageWithSeedTest is DSTest { Vm vm = Vm(HEVM_ADDRESS); function test_arbitrary_storage_with_seed() public { @@ -96,7 +96,7 @@ contract SymbolicStore { constructor() {} } -contract SymbolicStorageTest is DSTest { +contract SymbolicStorageWithSeedTest is DSTest { Vm vm = Vm(HEVM_ADDRESS); function test_SymbolicStorage() public { diff --git a/testdata/default/cheats/CopyStorage.t.sol b/testdata/default/cheats/CopyStorage.t.sol index 7d10ab2c3756..89584749745e 100644 --- a/testdata/default/cheats/CopyStorage.t.sol +++ b/testdata/default/cheats/CopyStorage.t.sol @@ -18,7 +18,7 @@ contract Counter { } } -contract CounterTest is DSTest { +contract CounterWithSeedTest is DSTest { Counter public counter; Counter public counter1; Vm vm = Vm(HEVM_ADDRESS); From b4f36db1be9eec4bf3439e9d0e91f1b923fbfcc0 Mon Sep 17 00:00:00 2001 From: grandizzy <38490174+grandizzy@users.noreply.github.com> Date: Wed, 11 Sep 2024 10:25:45 +0300 Subject: [PATCH 4/5] Update crates/cheatcodes/src/utils.rs Co-authored-by: DaniPopes <57450786+DaniPopes@users.noreply.github.com> --- crates/cheatcodes/src/utils.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/crates/cheatcodes/src/utils.rs b/crates/cheatcodes/src/utils.rs index d4a1681f4dfc..0896f2b31631 100644 --- a/crates/cheatcodes/src/utils.rs +++ b/crates/cheatcodes/src/utils.rs @@ -149,6 +149,7 @@ impl Cheatcode for resumeTracingCall { Ok(Default::default()) } } + impl Cheatcode for setArbitraryStorageCall { fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { let Self { target } = self; From 02086f94db6b6db262ded77b53d5c480f9f4fc8a Mon Sep 17 00:00:00 2001 From: grandizzy Date: Wed, 11 Sep 2024 13:13:38 +0300 Subject: [PATCH 5/5] Fix tests with isolate-by-default --- crates/forge/tests/it/cheats.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/forge/tests/it/cheats.rs b/crates/forge/tests/it/cheats.rs index 4bbeb9fe3d89..a60602cbc1cc 100644 --- a/crates/forge/tests/it/cheats.rs +++ b/crates/forge/tests/it/cheats.rs @@ -24,7 +24,7 @@ async fn test_cheats_local(test_data: &ForgeTestData) { } if cfg!(feature = "isolate-by-default") { - filter = filter.exclude_contracts("LastCallGasDefaultTest"); + filter = filter.exclude_contracts("(LastCallGasDefaultTest|MockFunctionTest|WithSeed)"); } let mut config = test_data.config.clone();