From 6ad0f7d88d3432732b7449ea9bedb375c67cda86 Mon Sep 17 00:00:00 2001 From: Ryan Pate Date: Fri, 22 Dec 2023 14:01:35 -0800 Subject: [PATCH] fix(forge): Optionally use create2 factory in tests --- crates/cheatcodes/src/config.rs | 4 ++ crates/cheatcodes/src/inspector.rs | 76 +++++++++++++++++++++++------- crates/common/src/evm.rs | 12 +++++ crates/config/README.md | 1 + crates/config/src/lib.rs | 6 +++ crates/evm/core/src/opts.rs | 3 ++ crates/forge/tests/cli/config.rs | 1 + crates/forge/tests/it/repros.rs | 7 +++ testdata/repros/Issue5529.t.sol | 37 +++++++++++++++ 9 files changed, 129 insertions(+), 18 deletions(-) create mode 100644 testdata/repros/Issue5529.t.sol diff --git a/crates/cheatcodes/src/config.rs b/crates/cheatcodes/src/config.rs index 85b2dcaab61f..d1b1d716f5b7 100644 --- a/crates/cheatcodes/src/config.rs +++ b/crates/cheatcodes/src/config.rs @@ -20,6 +20,8 @@ use std::{ pub struct CheatsConfig { /// Whether the FFI cheatcode is enabled. pub ffi: bool, + /// Use the create 2 factory in all cases including tests and non-broadcasting scripts. + pub always_use_create_2_factory: bool, /// RPC storage caching settings determines what chains and endpoints to cache pub rpc_storage_caching: StorageCachingConfig, /// All known endpoints and their aliases @@ -50,6 +52,7 @@ impl CheatsConfig { Self { ffi: evm_opts.ffi, + always_use_create_2_factory: evm_opts.always_use_create_2_factory, rpc_storage_caching: config.rpc_storage_caching.clone(), rpc_endpoints, paths: config.project_paths(), @@ -164,6 +167,7 @@ impl Default for CheatsConfig { fn default() -> Self { Self { ffi: false, + always_use_create_2_factory: false, rpc_storage_caching: Default::default(), rpc_endpoints: Default::default(), paths: ProjectPathsConfig::builder().build_with_root("./"), diff --git a/crates/cheatcodes/src/inspector.rs b/crates/cheatcodes/src/inspector.rs index 854397480cba..ce2cd91af8e5 100644 --- a/crates/cheatcodes/src/inspector.rs +++ b/crates/cheatcodes/src/inspector.rs @@ -1183,18 +1183,12 @@ impl Inspector for Cheatcodes { data.env.tx.caller = broadcast.new_origin; if data.journaled_state.depth() == broadcast.depth { - let (bytecode, to, nonce) = match process_create( + let (bytecode, to, nonce) = process_broadcast_create( broadcast.new_origin, call.init_code.clone(), data, call, - ) { - Ok(val) => val, - Err(err) => { - return (InstructionResult::Revert, None, gas, Error::encode(err)) - } - }; - + ); let is_fixed_gas_limit = check_if_fixed_gas_limit(data, call.gas_limit); self.broadcastable_transactions.push_back(BroadcastableTransaction { @@ -1222,6 +1216,14 @@ impl Inspector for Cheatcodes { } } + // Apply the Create2 deployer + if self.broadcast.is_some() || self.config.always_use_create_2_factory { + match apply_create2_deployer(data, call, &self.prank, &self.broadcast) { + Ok(_val) => {} + Err(err) => return (InstructionResult::Revert, None, gas, Error::encode(err)), + }; + } + // allow cheatcodes from the address of the new contract // Compute the address *after* any possible broadcast updates, so it's based on the updated // call inputs @@ -1395,18 +1397,31 @@ fn mstore_revert_string(interpreter: &mut Interpreter<'_>, bytes: &[u8]) { interpreter.return_len = interpreter.shared_memory.len() - starting_offset } -fn process_create( - broadcast_sender: Address, - bytecode: Bytes, +/// Applies the default CREATE2 deployer for contract creation. +/// +/// This function is invoked during the contract creation process and updates the caller of the +/// contract creation transaction to be the `DEFAULT_CREATE2_DEPLOYER` if the `CreateScheme` is +/// `Create2` and the current execution depth matches the depth at which the `prank` or `broadcast` +/// was started, or the default depth of 1 if no prank or broadcast is currently active. +/// +/// Returns a `DatabaseError::MissingCreate2Deployer` if the `DEFAULT_CREATE2_DEPLOYER` account is +/// not found or if it does not have any associated bytecode. +fn apply_create2_deployer( data: &mut EVMData<'_, DB>, call: &mut CreateInputs, -) -> Result<(Bytes, Option
, u64), DB::Error> { - match call.scheme { - CreateScheme::Create => { - call.caller = broadcast_sender; - Ok((bytecode, None, data.journaled_state.account(broadcast_sender).info.nonce)) + prank: &Option, + broadcast: &Option, +) -> Result<(), DB::Error> { + if let CreateScheme::Create2 { salt: _ } = call.scheme { + let mut base_depth = 1; + if let Some(prank) = &prank { + base_depth = prank.depth; + } else if let Some(broadcast) = &broadcast { + base_depth = broadcast.depth; } - CreateScheme::Create2 { salt } => { + // If the create scheme is Create2 and the depth equals the broadcast/prank/default + // depth, then use the default create2 factory as the deployer + if data.journaled_state.depth() == base_depth { // Sanity checks for our CREATE2 deployer let info = &data.journaled_state.load_account(DEFAULT_CREATE2_DEPLOYER, data.db)?.0.info; @@ -1419,7 +1434,32 @@ fn process_create( } call.caller = DEFAULT_CREATE2_DEPLOYER; + } + } + Ok(()) +} +/// Processes the creation of a new contract when broadcasting, preparing the necessary data for the +/// transaction to deploy the contract. +/// +/// Returns the transaction calldata and the target address. +/// +/// If the CreateScheme is Create, then this function returns the input bytecode without +/// modification and no address since it will be filled in later. If the CreateScheme is Create2, +/// then this function returns the calldata for the call to the create2 deployer which must be the +/// salt and init code concatenated. +fn process_broadcast_create( + broadcast_sender: Address, + bytecode: Bytes, + data: &mut EVMData<'_, DB>, + call: &mut CreateInputs, +) -> (Bytes, Option
, u64) { + match call.scheme { + CreateScheme::Create => { + call.caller = broadcast_sender; + (bytecode, None, data.journaled_state.account(broadcast_sender).info.nonce) + } + CreateScheme::Create2 { salt } => { // We have to increment the nonce of the user address, since this create2 will be done // by the create2_deployer let account = data.journaled_state.state().get_mut(&broadcast_sender).unwrap(); @@ -1429,7 +1469,7 @@ fn process_create( // Proxy deployer requires the data to be `salt ++ init_code` let calldata = [&salt.to_be_bytes::<32>()[..], &bytecode[..]].concat(); - Ok((calldata.into(), Some(DEFAULT_CREATE2_DEPLOYER), prev)) + (calldata.into(), Some(DEFAULT_CREATE2_DEPLOYER), prev) } } } diff --git a/crates/common/src/evm.rs b/crates/common/src/evm.rs index 2080951abb4a..a924fe5251a1 100644 --- a/crates/common/src/evm.rs +++ b/crates/common/src/evm.rs @@ -95,6 +95,11 @@ pub struct EvmArgs { #[serde(skip)] pub ffi: bool, + /// Use the create 2 factory in all cases including tests and non-broadcasting scripts. + #[clap(long)] + #[serde(skip)] + pub always_use_create_2_factory: bool, + /// Verbosity of the EVM. /// /// Pass multiple times to increase the verbosity (e.g. -v, -vv, -vvv). @@ -161,6 +166,13 @@ impl Provider for EvmArgs { dict.insert("ffi".to_string(), self.ffi.into()); } + if self.always_use_create_2_factory { + dict.insert( + "always_use_create_2_factory".to_string(), + self.always_use_create_2_factory.into(), + ); + } + if self.no_storage_caching { dict.insert("no_storage_caching".to_string(), self.no_storage_caching.into()); } diff --git a/crates/config/README.md b/crates/config/README.md index 5308cdfab341..ac422e13361a 100644 --- a/crates/config/README.md +++ b/crates/config/README.md @@ -114,6 +114,7 @@ no_match_contract = "Bar" match_path = "*/Foo*" no_match_path = "*/Bar*" ffi = false +always_use_create_2_factory = false # These are the default callers, generated using `address(uint160(uint256(keccak256("foundry default caller"))))` sender = '0x1804c8AB1F12E6bbf3894d4083f33e07309d1f38' tx_origin = '0x1804c8AB1F12E6bbf3894d4083f33e07309d1f38' diff --git a/crates/config/src/lib.rs b/crates/config/src/lib.rs index 618f55a1e863..62a273615744 100644 --- a/crates/config/src/lib.rs +++ b/crates/config/src/lib.rs @@ -239,6 +239,8 @@ pub struct Config { pub invariant: InvariantConfig, /// Whether to allow ffi cheatcodes in test pub ffi: bool, + /// Use the create 2 factory in all cases including tests and non-broadcasting scripts. + pub always_use_create_2_factory: bool, /// The address which will be executing all tests pub sender: Address, /// The tx.origin value during EVM execution @@ -1789,6 +1791,7 @@ impl Default for Config { path_pattern_inverse: None, fuzz: Default::default(), invariant: Default::default(), + always_use_create_2_factory: false, ffi: false, sender: Config::DEFAULT_SENDER, tx_origin: Config::DEFAULT_SENDER, @@ -3356,6 +3359,7 @@ mod tests { revert_strings = "strip" allow_paths = ["allow", "paths"] build_info_path = "build-info" + always_use_create_2_factory = true [rpc_endpoints] optimism = "https://example.com/" @@ -3407,6 +3411,7 @@ mod tests { ), ]), build_info_path: Some("build-info".into()), + always_use_create_2_factory: true, ..Config::default() } ); @@ -3458,6 +3463,7 @@ mod tests { evm_version = 'london' extra_output = [] extra_output_files = [] + always_use_create_2_factory = false ffi = false force = false gas_limit = 9223372036854775807 diff --git a/crates/evm/core/src/opts.rs b/crates/evm/core/src/opts.rs index 14b3e2cfcf33..a43eb7eefb61 100644 --- a/crates/evm/core/src/opts.rs +++ b/crates/evm/core/src/opts.rs @@ -53,6 +53,9 @@ pub struct EvmOpts { /// Enables the FFI cheatcode. pub ffi: bool, + /// Use the create 2 factory in all cases including tests and non-broadcasting scripts. + pub always_use_create_2_factory: bool, + /// Verbosity mode of EVM output as number of occurrences. pub verbosity: u8, diff --git a/crates/forge/tests/cli/config.rs b/crates/forge/tests/cli/config.rs index 1b69369db52a..c6c04bf5d756 100644 --- a/crates/forge/tests/cli/config.rs +++ b/crates/forge/tests/cli/config.rs @@ -64,6 +64,7 @@ forgetest!(can_extract_config_values, |prj, cmd| { }, invariant: InvariantConfig { runs: 256, ..Default::default() }, ffi: true, + always_use_create_2_factory: false, sender: "00a329c0648769A73afAc7F9381D08FB43dBEA72".parse().unwrap(), tx_origin: "00a329c0648769A73afAc7F9F81E08FB43dBEA72".parse().unwrap(), initial_balance: U256::from(0xffffffffffffffffffffffffu128), diff --git a/crates/forge/tests/it/repros.rs b/crates/forge/tests/it/repros.rs index b53002ae2d17..2da469e4f69d 100644 --- a/crates/forge/tests/it/repros.rs +++ b/crates/forge/tests/it/repros.rs @@ -290,3 +290,10 @@ test_repro!(6554; |config| { // https://github.com/foundry-rs/foundry/issues/6759 test_repro!(6759); + +// https://github.com/foundry-rs/foundry/issues/5529 +test_repro!(5529; |config| { + let mut cheats_config = config.runner.cheats_config.as_ref().clone(); + cheats_config.always_use_create_2_factory = true; + config.runner.cheats_config = std::sync::Arc::new(cheats_config); +}); diff --git a/testdata/repros/Issue5529.t.sol b/testdata/repros/Issue5529.t.sol new file mode 100644 index 000000000000..801a1167ec0b --- /dev/null +++ b/testdata/repros/Issue5529.t.sol @@ -0,0 +1,37 @@ +// SPDX-License-Identifier: UNLICENSED +pragma solidity ^0.8.13; + +import "ds-test/test.sol"; +import "../cheats/Vm.sol"; + +contract Counter { + uint256 public number; + + function setNumber(uint256 newNumber) public { + number = newNumber; + } + + function increment() public { + number++; + } +} + +contract CounterTest is DSTest { + Vm constant vm = Vm(HEVM_ADDRESS); + + Counter public counter; + address public constant default_create2_factory = 0x4e59b44847b379578588920cA78FbF26c0B4956C; + + function testCreate2FactoryUsedInTests() public { + address a = vm.computeCreate2Address(0, keccak256(type(Counter).creationCode), address(default_create2_factory)); + address b = address(new Counter{salt: 0}()); + require(a == b, "create2 address mismatch"); + } + + function testCreate2FactoryUsedWhenPranking() public { + vm.startPrank(address(1234)); + address a = vm.computeCreate2Address(0, keccak256(type(Counter).creationCode), address(default_create2_factory)); + address b = address(new Counter{salt: 0}()); + require(a == b, "create2 address mismatch"); + } +}