diff --git a/Cargo.lock b/Cargo.lock index 69bee29216009..68d1581afc8a8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4213,6 +4213,7 @@ dependencies = [ "jiff", "num-format", "path-slash", + "regex", "reqwest", "semver 1.0.26", "serde", diff --git a/crates/cast/src/cmd/call.rs b/crates/cast/src/cmd/call.rs index c329dbe3cc717..65dc59c465833 100644 --- a/crates/cast/src/cmd/call.rs +++ b/crates/cast/src/cmd/call.rs @@ -34,6 +34,8 @@ use regex::Regex; use revm::context::TransactionType; use std::{str::FromStr, sync::LazyLock}; +use super::run::fetch_contracts_bytecode_from_trace; + // matches override pattern
:: // e.g. 0x123:0x1:0x1234 static OVERRIDE_PATTERN: LazyLock = @@ -301,10 +303,12 @@ impl CallArgs { ), }; + let contracts_bytecode = fetch_contracts_bytecode_from_trace(&provider, &trace).await?; handle_traces( trace, &config, chain, + &contracts_bytecode, labels, with_local_artifacts, debug, diff --git a/crates/cast/src/cmd/run.rs b/crates/cast/src/cmd/run.rs index 04f52957ef668..1af816ad5b40b 100644 --- a/crates/cast/src/cmd/run.rs +++ b/crates/cast/src/cmd/run.rs @@ -1,6 +1,10 @@ use alloy_consensus::Transaction; use alloy_network::{AnyNetwork, TransactionResponse}; -use alloy_provider::Provider; +use alloy_primitives::{ + map::{HashMap, HashSet}, + Address, Bytes, +}; +use alloy_provider::{Provider, RootProvider}; use alloy_rpc_types::BlockTransactions; use clap::Parser; use eyre::{Result, WrapErr}; @@ -21,7 +25,7 @@ use foundry_config::{ use foundry_evm::{ executors::{EvmError, TracingExecutor}, opts::EvmOpts, - traces::{InternalTraceMode, TraceMode}, + traces::{InternalTraceMode, TraceMode, Traces}, utils::configure_tx_env, Env, }; @@ -272,10 +276,12 @@ impl RunArgs { } }; + let contracts_bytecode = fetch_contracts_bytecode_from_trace(&provider, &result).await?; handle_traces( result, &config, chain, + &contracts_bytecode, self.label, self.with_local_artifacts, self.debug, @@ -287,6 +293,47 @@ impl RunArgs { } } +pub async fn fetch_contracts_bytecode_from_trace( + provider: &RootProvider, + result: &TraceResult, +) -> Result> { + let mut contracts_bytecode = HashMap::default(); + if let Some(ref traces) = result.traces { + let addresses = gather_trace_addresses(traces); + let results = futures::future::join_all(addresses.into_iter().map(async |a| { + ( + a, + provider.get_code_at(a).await.unwrap_or_else(|e| { + sh_warn!("Failed to fetch code for {a:?}: {e:?}").ok(); + Bytes::new() + }), + ) + })) + .await; + for (address, code) in results { + if !code.is_empty() { + contracts_bytecode.insert(address, code); + } + } + } + Ok(contracts_bytecode) +} + +fn gather_trace_addresses(traces: &Traces) -> HashSet
{ + let mut addresses = HashSet::default(); + for (_, trace) in traces { + for node in trace.arena.nodes() { + if !node.trace.address.is_zero() { + addresses.insert(node.trace.address); + } + if !node.trace.caller.is_zero() { + addresses.insert(node.trace.caller); + } + } + } + addresses +} + impl figment::Provider for RunArgs { fn metadata(&self) -> Metadata { Metadata::named("RunArgs") diff --git a/crates/cli/src/utils/cmd.rs b/crates/cli/src/utils/cmd.rs index 21f411d446235..e3e6a166fdeac 100644 --- a/crates/cli/src/utils/cmd.rs +++ b/crates/cli/src/utils/cmd.rs @@ -1,5 +1,5 @@ use alloy_json_abi::JsonAbi; -use alloy_primitives::Address; +use alloy_primitives::{map::HashMap, Address, Bytes}; use eyre::{Result, WrapErr}; use foundry_common::{ compile::ProjectCompiler, fs, selectors::SelectorKind, shell, ContractsByArtifact, @@ -330,10 +330,12 @@ impl TryFrom> for TraceResult { } /// labels the traces, conditionally prints them or opens the debugger +#[expect(clippy::too_many_arguments)] pub async fn handle_traces( mut result: TraceResult, config: &Config, chain: Option, + contracts_bytecode: &HashMap, labels: Vec, with_local_artifacts: bool, debug: bool, @@ -372,7 +374,7 @@ pub async fn handle_traces( let mut identifier = TraceIdentifiers::new().with_etherscan(config, chain)?; if let Some(contracts) = &known_contracts { builder = builder.with_known_contracts(contracts); - identifier = identifier.with_local(contracts); + identifier = identifier.with_local_and_bytecodes(contracts, contracts_bytecode); } let mut decoder = builder.build(); diff --git a/crates/common/Cargo.toml b/crates/common/Cargo.toml index b440a5f7cdb84..76f42de2f0f6b 100644 --- a/crates/common/Cargo.toml +++ b/crates/common/Cargo.toml @@ -57,6 +57,7 @@ itertools.workspace = true jiff.workspace = true num-format.workspace = true path-slash.workspace = true +regex.workspace = true reqwest.workspace = true semver.workspace = true serde = { workspace = true, features = ["derive"] } diff --git a/crates/common/src/contracts.rs b/crates/common/src/contracts.rs index 789491d074dbf..1ebb284e4133a 100644 --- a/crates/common/src/contracts.rs +++ b/crates/common/src/contracts.rs @@ -1,6 +1,6 @@ //! Commonly used contract types and functions. -use crate::compile::PathOrContractInfo; +use crate::{compile::PathOrContractInfo, strip_bytecode_placeholders}; use alloy_dyn_abi::JsonAbiExt; use alloy_json_abi::{Event, Function, JsonAbi}; use alloy_primitives::{hex, Address, Bytes, Selector, B256}; @@ -87,6 +87,16 @@ impl ContractData { pub fn deployed_bytecode(&self) -> Option<&Bytes> { self.deployed_bytecode.as_ref()?.bytes().filter(|b| !b.is_empty()) } + + /// Returns the bytecode without placeholders, if present. + pub fn bytecode_without_placeholders(&self) -> Option { + strip_bytecode_placeholders(self.bytecode.as_ref()?.object.as_ref()?) + } + + /// Returns the deployed bytecode without placeholders, if present. + pub fn deployed_bytecode_without_placeholders(&self) -> Option { + strip_bytecode_placeholders(self.deployed_bytecode.as_ref()?.object.as_ref()?) + } } type ArtifactWithContractRef<'a> = (&'a ArtifactId, &'a ContractData); diff --git a/crates/common/src/utils.rs b/crates/common/src/utils.rs index a655d1e3d713f..3a9294c1b40cc 100644 --- a/crates/common/src/utils.rs +++ b/crates/common/src/utils.rs @@ -1,6 +1,13 @@ //! Uncategorised utilities. -use alloy_primitives::{keccak256, B256, U256}; +use alloy_primitives::{hex, keccak256, Bytes, B256, U256}; +use foundry_compilers::artifacts::BytecodeObject; +use regex::Regex; +use std::sync::LazyLock; + +static BYTECODE_PLACEHOLDER_RE: LazyLock = + LazyLock::new(|| Regex::new(r"__\$.{34}\$__").expect("invalid regex")); + /// Block on a future using the current tokio runtime on the current thread. pub fn block_on(future: F) -> F::Output { block_on_handle(&tokio::runtime::Handle::current(), future) @@ -54,3 +61,19 @@ pub fn ignore_metadata_hash(bytecode: &[u8]) -> &[u8] { bytecode } } + +/// Strips all __$xxx$__ placeholders from the bytecode if it's an unlinked bytecode. +/// by replacing them with 20 zero bytes. +/// This is useful for matching bytecodes to a contract source, and for the source map, +/// in which the actual address of the placeholder isn't important. +pub fn strip_bytecode_placeholders(bytecode: &BytecodeObject) -> Option { + match &bytecode { + BytecodeObject::Bytecode(bytes) => Some(bytes.clone()), + BytecodeObject::Unlinked(s) => { + // Replace all __$xxx$__ placeholders with 32 zero bytes + let s = (*BYTECODE_PLACEHOLDER_RE).replace_all(s, "00".repeat(40)); + let bytes = hex::decode(s.as_bytes()); + Some(bytes.ok()?.into()) + } + } +} diff --git a/crates/evm/traces/src/debug/sources.rs b/crates/evm/traces/src/debug/sources.rs index cfd7056e5a8b6..12aaafb29ad5a 100644 --- a/crates/evm/traces/src/debug/sources.rs +++ b/crates/evm/traces/src/debug/sources.rs @@ -1,5 +1,5 @@ use eyre::{Context, Result}; -use foundry_common::compact_to_contract; +use foundry_common::{compact_to_contract, strip_bytecode_placeholders}; use foundry_compilers::{ artifacts::{ sourcemap::{SourceElement, SourceMap}, @@ -94,9 +94,9 @@ impl ArtifactData { }) }; - // Only parse bytecode if it's not empty. - let pc_ic_map = if let Some(bytes) = b.bytes() { - (!bytes.is_empty()).then(|| PcIcMap::new(bytes)) + // Only parse bytecode if it's not empty, stripping placeholders if necessary. + let pc_ic_map = if let Some(bytes) = strip_bytecode_placeholders(&b.object) { + (!bytes.is_empty()).then(|| PcIcMap::new(bytes.as_ref())) } else { None }; diff --git a/crates/evm/traces/src/identifier/local.rs b/crates/evm/traces/src/identifier/local.rs index f0bb8cd9d7d76..94b24c341a493 100644 --- a/crates/evm/traces/src/identifier/local.rs +++ b/crates/evm/traces/src/identifier/local.rs @@ -1,6 +1,7 @@ use super::{IdentifiedAddress, TraceIdentifier}; use alloy_dyn_abi::JsonAbiExt; use alloy_json_abi::JsonAbi; +use alloy_primitives::{map::HashMap, Address, Bytes}; use foundry_common::contracts::{bytecode_diff_score, ContractsByArtifact}; use foundry_compilers::ArtifactId; use revm_inspectors::tracing::types::CallTraceNode; @@ -12,6 +13,8 @@ pub struct LocalTraceIdentifier<'a> { known_contracts: &'a ContractsByArtifact, /// Vector of pairs of artifact ID and the runtime code length of the given artifact. ordered_ids: Vec<(&'a ArtifactId, usize)>, + /// The contracts bytecode. + contracts_bytecode: Option<&'a HashMap>, } impl<'a> LocalTraceIdentifier<'a> { @@ -24,7 +27,12 @@ impl<'a> LocalTraceIdentifier<'a> { .map(|(id, bytecode)| (id, bytecode.len())) .collect::>(); ordered_ids.sort_by_key(|(_, len)| *len); - Self { known_contracts, ordered_ids } + Self { known_contracts, ordered_ids, contracts_bytecode: None } + } + + pub fn with_bytecodes(mut self, contracts_bytecode: &'a HashMap) -> Self { + self.contracts_bytecode = Some(contracts_bytecode); + self } /// Returns the known contracts. @@ -48,9 +56,9 @@ impl<'a> LocalTraceIdentifier<'a> { let contract = self.known_contracts.get(id)?; // Select bytecodes to compare based on `is_creation` flag. let (contract_bytecode, current_bytecode) = if is_creation { - (contract.bytecode(), creation_code) + (contract.bytecode_without_placeholders(), creation_code) } else { - (contract.deployed_bytecode(), runtime_code) + (contract.deployed_bytecode_without_placeholders(), runtime_code) }; if let Some(bytecode) = contract_bytecode { @@ -67,7 +75,7 @@ impl<'a> LocalTraceIdentifier<'a> { } } - let score = bytecode_diff_score(bytecode, current_bytecode); + let score = bytecode_diff_score(&bytecode, current_bytecode); if score == 0.0 { trace!(target: "evm::traces::local", "found exact match"); return Some((id, &contract.abi)); @@ -161,7 +169,18 @@ impl TraceIdentifier for LocalTraceIdentifier<'_> { let _span = trace_span!(target: "evm::traces::local", "identify", %address).entered(); - let (id, abi) = self.identify_code(runtime_code?, creation_code?)?; + // In order to identify the addresses, we need at least the runtime code. It can be + // obtained from the trace itself (if it's a CREATE* call), or from the fetched + // bytecodes. + let (runtime_code, creation_code) = match (runtime_code, creation_code) { + (Some(runtime_code), Some(creation_code)) => (runtime_code, creation_code), + (Some(runtime_code), _) => (runtime_code, &[] as &[u8]), + _ => { + let code = self.contracts_bytecode?.get(&address)?; + (code.as_ref(), &[] as &[u8]) + } + }; + let (id, abi) = self.identify_code(runtime_code, creation_code)?; trace!(target: "evm::traces::local", id=%id.identifier(), "identified"); Some(IdentifiedAddress { diff --git a/crates/evm/traces/src/identifier/mod.rs b/crates/evm/traces/src/identifier/mod.rs index 0654b5087764f..00045863847a8 100644 --- a/crates/evm/traces/src/identifier/mod.rs +++ b/crates/evm/traces/src/identifier/mod.rs @@ -1,5 +1,5 @@ use alloy_json_abi::JsonAbi; -use alloy_primitives::Address; +use alloy_primitives::{map::HashMap, Address, Bytes}; use foundry_common::ContractsByArtifact; use foundry_compilers::ArtifactId; use foundry_config::{Chain, Config}; @@ -79,6 +79,17 @@ impl<'a> TraceIdentifiers<'a> { self } + /// Sets the local identifier. + pub fn with_local_and_bytecodes( + mut self, + known_contracts: &'a ContractsByArtifact, + contracts_bytecode: &'a HashMap, + ) -> Self { + self.local = + Some(LocalTraceIdentifier::new(known_contracts).with_bytecodes(contracts_bytecode)); + self + } + /// Sets the etherscan identifier. pub fn with_etherscan(mut self, config: &Config, chain: Option) -> eyre::Result { self.etherscan = EtherscanIdentifier::new(config, chain)?;