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..703d7db126e9 100644 --- a/crates/cheatcodes/src/evm.rs +++ b/crates/cheatcodes/src/evm.rs @@ -13,6 +13,7 @@ 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}, InnerEvmContext, @@ -89,7 +90,25 @@ 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())?; + + if val.is_cold && val.data.is_zero() { + 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. + 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 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. + let rand_value = ccx.state.rng().gen(); + val.data = + ccx.state.arbitrary_storage.copy(ccx.ecx, target, slot.into(), rand_value); + } + } + 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..e09e1e8c88b9 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; @@ -254,6 +254,89 @@ 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, + ) { + self.values.get_mut(&address).expect("missing arbitrary address entry").insert(slot, data); + if let Ok(mut account) = ecx.load_account(address) { + 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; @@ -320,6 +403,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 +454,9 @@ pub struct Cheatcodes { /// Ignored traces. pub ignored_traces: IgnoredTraces, + + /// Addresses with arbitrary storage. + pub arbitrary_storage: ArbitraryStorage, } // This is not derived because calling this in `fn new` with `..Default::default()` creates a second @@ -396,6 +485,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 +500,7 @@ impl Cheatcodes { breakpoints: Default::default(), rng: Default::default(), ignored_traces: Default::default(), + arbitrary_storage: Default::default(), } } @@ -1045,7 +1136,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 +1144,14 @@ impl Inspector for Cheatcodes { if self.gas_metering.touched { self.meter_gas_check(interpreter); } + + // `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); + } } fn log(&mut self, interpreter: &mut Interpreter, _ecx: &mut EvmContext, log: &Log) { @@ -1465,6 +1564,43 @@ impl Cheatcodes { } } + /// 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 arbitrary_storage_end( + &mut self, + interpreter: &mut Interpreter, + ecx: &mut EvmContext, + ) { + 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, + ); + } else { + self.arbitrary_storage.save( + &mut ecx.inner, + target_address, + key, + arbitrary_value, + ); + } + } + } + } + /// 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..0896f2b31631 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,31 @@ 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.mark_arbitrary(target); + + Ok(Default::default()) + } +} + +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); + } + } + + Ok(Default::default()) + } +} diff --git a/crates/evm/evm/src/inspectors/stack.rs b/crates/evm/evm/src/inspectors/stack.rs index e44d499eafa2..3df8dc8f01da 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) { + // 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; + } + } + 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..a60602cbc1cc 100644 --- a/crates/forge/tests/it/cheats.rs +++ b/crates/forge/tests/it/cheats.rs @@ -7,14 +7,16 @@ use crate::{ TEST_DATA_MULTI_VERSION, }, }; +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) { @@ -22,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(); @@ -32,7 +34,7 @@ async fn test_cheats_local(test_data: &ForgeTestData) { 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}*")); @@ -43,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 @@ -53,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/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..86910279e98e --- /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 AContractArbitraryStorageWithSeedTest 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 SymbolicStorageWithSeedTest 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..89584749745e --- /dev/null +++ b/testdata/default/cheats/CopyStorage.t.sol @@ -0,0 +1,158 @@ +// 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; + int256[] public c; + + function setA(uint256 _a) public { + a = _a; + } + + function setB(address _b) public { + b = _b; + } +} + +contract CounterWithSeedTest is DSTest { + Counter public counter; + Counter public counter1; + Vm vm = Vm(HEVM_ADDRESS); + + 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)); + + 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)); + } + + 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 { + uint256 public x; +} + +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 { + vm.store(contractAddress, bytes32(slot), bytes32(value)); + } + + function setUp() public { + csc_1 = new CopyStorageContract(); + csc_2 = new CopyStorageContract(); + csc_3 = 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); + } + + 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); + } +} 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); + } +}