Skip to content

Commit

Permalink
refactor: inject call to CREATE2 factory through custom revm handler (#…
Browse files Browse the repository at this point in the history
…7653)

* wip

* wip

* add docs

* clippy

* update doc

* simplify logic

* review fixes

* doc

* review fixes

* fix
  • Loading branch information
klkvr authored Apr 17, 2024
1 parent 19871fc commit 63072be
Show file tree
Hide file tree
Showing 9 changed files with 231 additions and 158 deletions.
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

0 comments on commit 63072be

Please sign in to comment.