Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

refactor: inject call to CREATE2 factory through custom revm handler #7653

Merged
merged 10 commits into from
Apr 17, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion crates/anvil/src/eth/backend/mem/inspector.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,13 +6,14 @@ use foundry_evm::{
call_inspectors,
decode::decode_console_logs,
inspectors::{LogCollector, TracingInspector},
revm,
revm::{
self,
interpreter::{CallInputs, CallOutcome, CreateInputs, CreateOutcome, Interpreter},
primitives::U256,
EvmContext,
},
traces::TracingInspectorConfig,
InspectorExt,
};

/// The [`revm::Inspector`] used when transacting in the evm
Expand Down Expand Up @@ -136,6 +137,8 @@ impl<DB: Database> revm::Inspector<DB> for Inspector {
}
}

impl<DB: Database> InspectorExt<DB> for Inspector {}

/// Prints all the logs
#[inline]
pub fn print_logs(logs: &[Log]) {
Expand Down
3 changes: 2 additions & 1 deletion crates/anvil/src/eth/backend/mem/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,7 @@ use foundry_evm::{
},
},
utils::new_evm_with_inspector_ref,
InspectorExt,
};
use futures::channel::mpsc::{unbounded, UnboundedSender};
use parking_lot::{Mutex, RwLock};
Expand Down Expand Up @@ -819,7 +820,7 @@ impl Backend {
) -> revm::Evm<'_, I, WrapDatabaseRef<DB>>
where
DB: revm::DatabaseRef,
I: revm::Inspector<WrapDatabaseRef<DB>>,
I: InspectorExt<WrapDatabaseRef<DB>>,
{
let mut evm = new_evm_with_inspector_ref(db, env, inspector);
if let Some(ref factory) = self.precompile_factory {
Expand Down
171 changes: 33 additions & 138 deletions crates/cheatcodes/src/inspector.rs
Original file line number Diff line number Diff line change
Expand Up @@ -21,8 +21,9 @@ use alloy_sol_types::{SolInterface, SolValue};
use foundry_common::{evm::Breakpoints, provider::alloy::RpcUrl, SELECTOR_LEN};
use foundry_evm_core::{
abi::Vm::stopExpectSafeMemoryCall,
backend::{DatabaseError, DatabaseExt, RevertDiagnostic},
constants::{CHEATCODE_ADDRESS, DEFAULT_CREATE2_DEPLOYER, HARDHAT_CONSOLE_ADDRESS},
backend::{DatabaseExt, RevertDiagnostic},
constants::{CHEATCODE_ADDRESS, HARDHAT_CONSOLE_ADDRESS},
InspectorExt,
};
use itertools::Itertools;
use revm::{
Expand Down Expand Up @@ -1327,22 +1328,19 @@ impl<DB: DatabaseExt> Inspector<DB> for Cheatcodes {
ecx.env.tx.caller = broadcast.new_origin;

if ecx.journaled_state.depth() == broadcast.depth {
let (bytecode, to, nonce) = process_broadcast_create(
broadcast.new_origin,
call.init_code.clone(),
ecx,
call,
);
call.caller = broadcast.new_origin;
let is_fixed_gas_limit = check_if_fixed_gas_limit(ecx, call.gas_limit);

let account = &ecx.journaled_state.state()[&broadcast.new_origin];

self.broadcastable_transactions.push_back(BroadcastableTransaction {
rpc: ecx.db.active_fork_url(),
transaction: TransactionRequest {
from: Some(broadcast.new_origin),
to,
to: None,
value: Some(call.value),
input: TransactionInput::new(bytecode),
nonce: Some(nonce),
input: TransactionInput::new(call.init_code.clone()),
nonce: Some(account.info.nonce),
gas: if is_fixed_gas_limit {
Some(call.gas_limit as u128)
} else {
Expand All @@ -1351,6 +1349,7 @@ impl<DB: DatabaseExt> Inspector<DB> for Cheatcodes {
..Default::default()
},
});

let kind = match call.scheme {
CreateScheme::Create => "create",
CreateScheme::Create2 { .. } => "create2",
Expand All @@ -1360,29 +1359,6 @@ impl<DB: DatabaseExt> Inspector<DB> for Cheatcodes {
}
}

// Apply the Create2 deployer
if self.broadcast.is_some() || self.config.always_use_create_2_factory {
match apply_create2_deployer(
ecx,
call,
self.prank.as_ref(),
self.broadcast.as_ref(),
self.recorded_account_diffs_stack.as_mut(),
) {
Ok(_) => {}
Err(err) => {
return Some(CreateOutcome {
result: InterpreterResult {
result: InstructionResult::Revert,
output: Error::encode(err),
gas,
},
address: None,
})
}
};
}

// 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
Expand Down Expand Up @@ -1526,6 +1502,29 @@ impl<DB: DatabaseExt> Inspector<DB> for Cheatcodes {
}
}

impl<DB: DatabaseExt> InspectorExt<DB> for Cheatcodes {
fn should_use_create2_factory(
&mut self,
ecx: &mut EvmContext<DB>,
inputs: &mut CreateInputs,
) -> bool {
if let CreateScheme::Create2 { .. } = inputs.scheme {
let target_depth = if let Some(prank) = &self.prank {
prank.depth
} else if let Some(broadcast) = &self.broadcast {
broadcast.depth
} else {
1
};

ecx.journaled_state.depth() == target_depth &&
(self.broadcast.is_some() || self.config.always_use_create_2_factory)
} else {
false
}
}
}

/// Helper that expands memory, stores a revert string pertaining to a disallowed memory write,
/// and sets the return range to the revert string's location in memory.
///
Expand Down Expand Up @@ -1554,110 +1553,6 @@ fn disallowed_mem_write(
};
}

/// 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<DB: DatabaseExt>(
ecx: &mut InnerEvmContext<DB>,
call: &mut CreateInputs,
prank: Option<&Prank>,
broadcast: Option<&Broadcast>,
diffs_stack: Option<&mut Vec<Vec<AccountAccess>>>,
) -> 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;
}

// 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 ecx.journaled_state.depth() == base_depth {
// Record the call to the create2 factory in the state diff
if let Some(recorded_account_diffs_stack) = diffs_stack {
let calldata = [&salt.to_be_bytes::<32>()[..], &call.init_code[..]].concat();
recorded_account_diffs_stack.push(vec![AccountAccess {
chainInfo: crate::Vm::ChainInfo {
forkId: ecx.db.active_fork_id().unwrap_or_default(),
chainId: U256::from(ecx.env.cfg.chain_id),
},
accessor: call.caller,
account: DEFAULT_CREATE2_DEPLOYER,
kind: crate::Vm::AccountAccessKind::Call,
initialized: true,
oldBalance: U256::ZERO, // updated on create_end
newBalance: U256::ZERO, // updated on create_end
value: call.value,
data: calldata.into(),
reverted: false,
deployedCode: Bytes::new(), // updated on create_end
storageAccesses: vec![], // updated on create_end
depth: ecx.journaled_state.depth(),
}])
}

// Sanity checks for our CREATE2 deployer
// TODO: use ecx.load_account
let info =
&ecx.journaled_state.load_account(DEFAULT_CREATE2_DEPLOYER, &mut ecx.db)?.0.info;
match &info.code {
Some(code) if code.is_empty() => return Err(DatabaseError::MissingCreate2Deployer),
None if ecx.db.code_by_hash(info.code_hash)?.is_empty() => {
return Err(DatabaseError::MissingCreate2Deployer)
}
_ => {}
}

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<DB: DatabaseExt>(
broadcast_sender: Address,
bytecode: Bytes,
ecx: &mut InnerEvmContext<DB>,
call: &mut CreateInputs,
) -> (Bytes, Option<Address>, u64) {
call.caller = broadcast_sender;
match call.scheme {
CreateScheme::Create => {
(bytecode, None, ecx.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 = ecx.journaled_state.state().get_mut(&broadcast_sender).unwrap();
let prev = account.info.nonce;
// Touch account to ensure that incremented nonce is committed
account.mark_touch();
account.info.nonce += 1;
debug!(target: "cheatcodes", address=%broadcast_sender, nonce=prev+1, prev, "incremented nonce in create2");
// Proxy deployer requires the data to be `salt ++ init_code`
let calldata = [&salt.to_be_bytes::<32>()[..], &bytecode[..]].concat();
(calldata.into(), Some(DEFAULT_CREATE2_DEPLOYER), prev)
}
}
}

// Determines if the gas limit on a given call was manually set in the script and should therefore
// not be overwritten by later estimations
fn check_if_fixed_gas_limit<DB: DatabaseExt>(
Expand Down
7 changes: 4 additions & 3 deletions crates/evm/core/src/backend/cow.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ use crate::{
RevertSnapshotAction,
},
fork::{CreateFork, ForkId},
InspectorExt,
};
use alloy_genesis::GenesisAccount;
use alloy_primitives::{Address, B256, U256};
Expand All @@ -16,7 +17,7 @@ use revm::{
Account, AccountInfo, Bytecode, Env, EnvWithHandlerCfg, HashMap as Map, ResultAndState,
SpecId,
},
Database, DatabaseCommit, Inspector, JournaledState,
Database, DatabaseCommit, JournaledState,
};
use std::{borrow::Cow, collections::BTreeMap};

Expand Down Expand Up @@ -58,7 +59,7 @@ impl<'a> CowBackend<'a> {
///
/// Note: in case there are any cheatcodes executed that modify the environment, this will
/// update the given `env` with the new values.
pub fn inspect<'b, I: Inspector<&'b mut Self>>(
pub fn inspect<'b, I: InspectorExt<&'b mut Self>>(
&'b mut self,
env: &mut EnvWithHandlerCfg,
inspector: I,
Expand Down Expand Up @@ -176,7 +177,7 @@ impl<'a> DatabaseExt for CowBackend<'a> {
self.backend_mut(env).roll_fork_to_transaction(id, transaction, env, journaled_state)
}

fn transact<I: Inspector<Backend>>(
fn transact<I: InspectorExt<Backend>>(
&mut self,
id: Option<LocalForkId>,
transaction: B256,
Expand Down
11 changes: 6 additions & 5 deletions crates/evm/core/src/backend/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ use crate::{
fork::{CreateFork, ForkId, MultiFork, SharedBackend},
snapshot::Snapshots,
utils::configure_tx_env,
InspectorExt,
};
use alloy_genesis::GenesisAccount;
use alloy_primitives::{b256, keccak256, Address, B256, U256};
Expand All @@ -19,7 +20,7 @@ use revm::{
Account, AccountInfo, Bytecode, CreateScheme, Env, EnvWithHandlerCfg, HashMap as Map, Log,
ResultAndState, SpecId, State, StorageSlot, TransactTo, KECCAK_EMPTY,
},
Database, DatabaseCommit, Inspector, JournaledState,
Database, DatabaseCommit, JournaledState,
};
use std::{
collections::{BTreeMap, HashMap, HashSet},
Expand Down Expand Up @@ -188,7 +189,7 @@ pub trait DatabaseExt: Database<Error = DatabaseError> {
) -> eyre::Result<()>;

/// Fetches the given transaction for the fork and executes it, committing the state in the DB
fn transact<I: Inspector<Backend>>(
fn transact<I: InspectorExt<Backend>>(
&mut self,
id: Option<LocalForkId>,
transaction: B256,
Expand Down Expand Up @@ -780,7 +781,7 @@ impl Backend {
///
/// Note: in case there are any cheatcodes executed that modify the environment, this will
/// update the given `env` with the new values.
pub fn inspect<'a, I: Inspector<&'a mut Self>>(
pub fn inspect<'a, I: InspectorExt<&'a mut Self>>(
&'a mut self,
env: &mut EnvWithHandlerCfg,
inspector: I,
Expand Down Expand Up @@ -1228,7 +1229,7 @@ impl DatabaseExt for Backend {
Ok(())
}

fn transact<I: Inspector<Backend>>(
fn transact<I: InspectorExt<Backend>>(
&mut self,
maybe_id: Option<LocalForkId>,
transaction: B256,
Expand Down Expand Up @@ -1868,7 +1869,7 @@ fn update_env_block(env: &mut Env, fork_block: u64, block: &Block) {

/// Executes the given transaction and commits state changes to the database _and_ the journaled
/// state, with an optional inspector
fn commit_transaction<I: Inspector<Backend>>(
fn commit_transaction<I: InspectorExt<Backend>>(
tx: WithOtherFields<Transaction>,
mut env: EnvWithHandlerCfg,
journaled_state: &mut JournaledState,
Expand Down
24 changes: 24 additions & 0 deletions crates/evm/core/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,10 @@

#![warn(unused_crate_dependencies)]

use auto_impl::auto_impl;
use revm::{inspectors::NoOpInspector, interpreter::CreateInputs, Database, EvmContext, Inspector};
use revm_inspectors::access_list::AccessListInspector;

#[macro_use]
extern crate tracing;

Expand All @@ -19,3 +23,23 @@ pub mod opcodes;
pub mod opts;
pub mod snapshot;
pub mod utils;

/// An extension trait that allows us to add additional hooks to Inspector for later use in
/// handlers.
#[auto_impl(&mut, Box)]
pub trait InspectorExt<DB: Database>: Inspector<DB> {
/// Determines whether the `DEFAULT_CREATE2_DEPLOYER` should be used for a CREATE2 frame.
///
/// If this function returns true, we'll replace CREATE2 frame with a CALL frame to CREATE2
/// factory.
fn should_use_create2_factory(
&mut self,
_context: &mut EvmContext<DB>,
_inputs: &mut CreateInputs,
) -> bool {
false
}
}

impl<DB: Database> InspectorExt<DB> for NoOpInspector {}
impl<DB: Database> InspectorExt<DB> for AccessListInspector {}
Loading
Loading