From 8e11fdb57e7652e8efe0c860833c49fb807e2078 Mon Sep 17 00:00:00 2001 From: DaniPopes <57450786+DaniPopes@users.noreply.github.com> Date: Thu, 7 Mar 2024 21:43:53 +0100 Subject: [PATCH] chore(traces): add a trace identifier stack/builder --- crates/chisel/src/dispatcher.rs | 18 +++--- crates/cli/src/utils/cmd.rs | 12 +++- crates/evm/traces/src/decoder/mod.rs | 19 ++++-- crates/evm/traces/src/identifier/etherscan.rs | 62 ++++++++----------- crates/evm/traces/src/identifier/mod.rs | 58 ++++++++++++++++- crates/forge/bin/cmd/script/mod.rs | 27 ++++---- crates/forge/bin/cmd/test/mod.rs | 16 +++-- 7 files changed, 133 insertions(+), 79 deletions(-) diff --git a/crates/chisel/src/dispatcher.rs b/crates/chisel/src/dispatcher.rs index 46086f1dc56d..b6a4411ac907 100644 --- a/crates/chisel/src/dispatcher.rs +++ b/crates/chisel/src/dispatcher.rs @@ -17,7 +17,7 @@ use foundry_config::{Config, RpcEndpoint}; use foundry_evm::{ decode::decode_console_logs, traces::{ - identifier::{EtherscanIdentifier, SignaturesIdentifier}, + identifier::{SignaturesIdentifier, TraceIdentifiers}, render_trace_arena, CallTraceDecoder, CallTraceDecoderBuilder, TraceKind, }, }; @@ -893,11 +893,6 @@ impl ChiselDispatcher { result: &mut ChiselResult, // known_contracts: &ContractsByArtifact, ) -> eyre::Result { - let mut etherscan_identifier = EtherscanIdentifier::new( - &session_config.foundry_config, - session_config.evm_opts.get_remote_chain_id(), - )?; - let mut decoder = CallTraceDecoderBuilder::new() .with_labels(result.labeled_addresses.clone()) .with_signature_identifier(SignaturesIdentifier::new( @@ -906,9 +901,14 @@ impl ChiselDispatcher { )?) .build(); - for (_, trace) in &mut result.traces { - // decoder.identify(trace, &mut local_identifier); - decoder.identify(trace, &mut etherscan_identifier); + let mut identifier = TraceIdentifiers::new().with_etherscan( + &session_config.foundry_config, + session_config.evm_opts.get_remote_chain_id(), + )?; + if !identifier.is_empty() { + for (_, trace) in &mut result.traces { + decoder.identify(trace, &mut identifier); + } } Ok(decoder) } diff --git a/crates/cli/src/utils/cmd.rs b/crates/cli/src/utils/cmd.rs index 00a76002c831..1ea54d3f6413 100644 --- a/crates/cli/src/utils/cmd.rs +++ b/crates/cli/src/utils/cmd.rs @@ -397,12 +397,18 @@ pub async fn handle_traces( .build(); let mut etherscan_identifier = EtherscanIdentifier::new(config, chain)?; - for (_, trace) in &mut result.traces { - decoder.identify(trace, &mut etherscan_identifier); + if let Some(etherscan_identifier) = &mut etherscan_identifier { + for (_, trace) in &mut result.traces { + decoder.identify(trace, etherscan_identifier); + } } if debug { - let sources = etherscan_identifier.get_compiled_contracts().await?; + let sources = if let Some(etherscan_identifier) = etherscan_identifier { + etherscan_identifier.get_compiled_contracts().await? + } else { + Default::default() + }; let mut debugger = Debugger::builder() .debug_arena(&result.debug) .decoder(&decoder) diff --git a/crates/evm/traces/src/decoder/mod.rs b/crates/evm/traces/src/decoder/mod.rs index 6b5aaf35deaf..fbe797a0030b 100644 --- a/crates/evm/traces/src/decoder/mod.rs +++ b/crates/evm/traces/src/decoder/mod.rs @@ -7,7 +7,9 @@ use crate::{ use alloy_dyn_abi::{DecodedEvent, DynSolValue, EventExt, FunctionExt, JsonAbiExt}; use alloy_json_abi::{Error, Event, Function, JsonAbi}; use alloy_primitives::{Address, LogData, Selector, B256}; -use foundry_common::{abi::get_indexed_event, fmt::format_token, SELECTOR_LEN}; +use foundry_common::{ + abi::get_indexed_event, fmt::format_token, ContractsByArtifact, SELECTOR_LEN, +}; use foundry_evm_core::{ abi::{Console, HardhatConsole, Vm, HARDHAT_CONSOLE_SELECTOR_PATCHES}, constants::{ @@ -50,17 +52,22 @@ impl CallTraceDecoderBuilder { self } - /// Add known contracts to the decoder from a `LocalTraceIdentifier`. + /// Add known contracts to the decoder. #[inline] - pub fn with_local_identifier_abis(mut self, identifier: &LocalTraceIdentifier<'_>) -> Self { - let contracts = identifier.contracts(); - trace!(target: "evm::traces", len=contracts.len(), "collecting local identifier ABIs"); + pub fn with_known_contracts(mut self, contracts: &ContractsByArtifact) -> Self { + trace!(target: "evm::traces", len=contracts.len(), "collecting known contract ABIs"); for (abi, _) in contracts.values() { self.decoder.collect_abi(abi, None); } self } + /// Add known contracts to the decoder from a `LocalTraceIdentifier`. + #[inline] + pub fn with_local_identifier_abis(self, identifier: &LocalTraceIdentifier<'_>) -> Self { + self.with_known_contracts(identifier.contracts()) + } + /// Sets the verbosity level of the decoder. #[inline] pub fn with_verbosity(mut self, level: u8) -> Self { @@ -225,7 +232,7 @@ impl CallTraceDecoder { fn addresses<'a>( &'a self, arena: &'a CallTraceArena, - ) -> impl Iterator)> + 'a { + ) -> impl Iterator)> + Clone + 'a { arena .nodes() .iter() diff --git a/crates/evm/traces/src/identifier/etherscan.rs b/crates/evm/traces/src/identifier/etherscan.rs index 43552f12b250..50c273d07b0f 100644 --- a/crates/evm/traces/src/identifier/etherscan.rs +++ b/crates/evm/traces/src/identifier/etherscan.rs @@ -25,10 +25,9 @@ use std::{ use tokio::time::{Duration, Interval}; /// A trace identifier that tries to identify addresses using Etherscan. -#[derive(Default)] pub struct EtherscanIdentifier { /// The Etherscan client - client: Option>, + client: Arc, /// Tracks whether the API key provides was marked as invalid /// /// After the first [EtherscanError::InvalidApiKey] this will get set to true, so we can @@ -40,22 +39,21 @@ pub struct EtherscanIdentifier { impl EtherscanIdentifier { /// Creates a new Etherscan identifier with the given client - pub fn new(config: &Config, chain: Option) -> eyre::Result { + pub fn new(config: &Config, chain: Option) -> eyre::Result> { + // In offline mode, don't use Etherscan. if config.offline { - // offline mode, don't use etherscan - return Ok(Default::default()) - } - if let Some(config) = config.get_etherscan_config_with_chain(chain)? { - trace!(target: "etherscanidentifier", chain=?config.chain, url=?config.api_url, "using etherscan identifier"); - Ok(Self { - client: Some(Arc::new(config.into_client()?)), - invalid_api_key: Arc::new(Default::default()), - contracts: BTreeMap::new(), - sources: BTreeMap::new(), - }) - } else { - Ok(Default::default()) + return Ok(None); } + let Some(config) = config.get_etherscan_config_with_chain(chain)? else { + return Ok(None); + }; + trace!(target: "traces::etherscan", chain=?config.chain, url=?config.api_url, "using etherscan identifier"); + Ok(Some(Self { + client: Arc::new(config.into_client()?), + invalid_api_key: Arc::new(AtomicBool::new(false)), + contracts: BTreeMap::new(), + sources: BTreeMap::new(), + })) } /// Goes over the list of contracts we have pulled from the traces, clones their source from @@ -101,18 +99,13 @@ impl TraceIdentifier for EtherscanIdentifier { { trace!(target: "evm::traces", "identify {:?} addresses", addresses.size_hint().1); - let Some(client) = self.client.clone() else { - // no client was configured - return Vec::new() - }; - if self.invalid_api_key.load(Ordering::Relaxed) { // api key was marked as invalid return Vec::new() } let mut fetcher = EtherscanFetcher::new( - client, + self.client.clone(), Duration::from_secs(1), 5, Arc::clone(&self.invalid_api_key), @@ -191,16 +184,13 @@ impl EtherscanFetcher { fn queue_next_reqs(&mut self) { while self.in_progress.len() < self.concurrency { - if let Some(addr) = self.queue.pop() { - let client = Arc::clone(&self.client); - trace!(target: "etherscanidentifier", "fetching info for {:?}", addr); - self.in_progress.push(Box::pin(async move { - let res = client.contract_source_code(addr).await; - (addr, res) - })); - } else { - break - } + let Some(addr) = self.queue.pop() else { break }; + let client = Arc::clone(&self.client); + self.in_progress.push(Box::pin(async move { + trace!(target: "traces::etherscan", ?addr, "fetching info"); + let res = client.contract_source_code(addr).await; + (addr, res) + })); } } } @@ -234,24 +224,24 @@ impl Stream for EtherscanFetcher { } } Err(EtherscanError::RateLimitExceeded) => { - warn!(target: "etherscanidentifier", "rate limit exceeded on attempt"); + warn!(target: "traces::etherscan", "rate limit exceeded on attempt"); pin.backoff = Some(tokio::time::interval(pin.timeout)); pin.queue.push(addr); } Err(EtherscanError::InvalidApiKey) => { - warn!(target: "etherscanidentifier", "invalid api key"); + warn!(target: "traces::etherscan", "invalid api key"); // mark key as invalid pin.invalid_api_key.store(true, Ordering::Relaxed); return Poll::Ready(None) } Err(EtherscanError::BlockedByCloudflare) => { - warn!(target: "etherscanidentifier", "blocked by cloudflare"); + warn!(target: "traces::etherscan", "blocked by cloudflare"); // mark key as invalid pin.invalid_api_key.store(true, Ordering::Relaxed); return Poll::Ready(None) } Err(err) => { - warn!(target: "etherscanidentifier", "could not get etherscan info: {:?}", err); + warn!(target: "traces::etherscan", "could not get etherscan info: {:?}", err); } } } diff --git a/crates/evm/traces/src/identifier/mod.rs b/crates/evm/traces/src/identifier/mod.rs index 6d86b072abf6..a16b108d8537 100644 --- a/crates/evm/traces/src/identifier/mod.rs +++ b/crates/evm/traces/src/identifier/mod.rs @@ -1,6 +1,8 @@ use alloy_json_abi::JsonAbi; use alloy_primitives::Address; +use foundry_common::ContractsByArtifact; use foundry_compilers::ArtifactId; +use foundry_config::{Chain, Config}; use std::borrow::Cow; mod local; @@ -33,5 +35,59 @@ pub trait TraceIdentifier { /// Attempts to identify an address in one or more call traces. fn identify_addresses<'a, A>(&mut self, addresses: A) -> Vec> where - A: Iterator)>; + A: Iterator)> + Clone; +} + +/// A collection of trace identifiers. +pub struct TraceIdentifiers<'a> { + /// The local trace identifier. + pub local: Option>, + /// The optional Etherscan trace identifier. + pub etherscan: Option, +} + +impl Default for TraceIdentifiers<'_> { + fn default() -> Self { + Self::new() + } +} + +impl TraceIdentifier for TraceIdentifiers<'_> { + fn identify_addresses<'a, A>(&mut self, addresses: A) -> Vec> + where + A: Iterator)> + Clone, + { + let mut identities = Vec::new(); + if let Some(local) = &mut self.local { + identities.extend(local.identify_addresses(addresses.clone())); + } + if let Some(etherscan) = &mut self.etherscan { + identities.extend(etherscan.identify_addresses(addresses)); + } + identities + } +} + +impl<'a> TraceIdentifiers<'a> { + /// Creates a new, empty instance. + pub const fn new() -> Self { + Self { local: None, etherscan: None } + } + + /// Sets the local identifier. + pub fn with_local(mut self, known_contracts: &'a ContractsByArtifact) -> Self { + self.local = Some(LocalTraceIdentifier::new(known_contracts)); + self + } + + /// Sets the etherscan identifier. + pub fn with_etherscan(mut self, config: &Config, chain: Option) -> eyre::Result { + self.etherscan = EtherscanIdentifier::new(config, chain)?; + Ok(self) + } + + /// Returns `true` if there are no set identifiers. + pub fn is_empty(&self) -> bool { + self.local.is_none() && self.etherscan.is_none() + } } diff --git a/crates/forge/bin/cmd/script/mod.rs b/crates/forge/bin/cmd/script/mod.rs index 79c8f937ff12..98c3e223db76 100644 --- a/crates/forge/bin/cmd/script/mod.rs +++ b/crates/forge/bin/cmd/script/mod.rs @@ -13,8 +13,8 @@ use forge::{ decode::decode_console_logs, opts::EvmOpts, traces::{ - identifier::{EtherscanIdentifier, LocalTraceIdentifier, SignaturesIdentifier}, - render_trace_arena, CallTraceDecoder, CallTraceDecoderBuilder, TraceKind, Traces, + identifier::SignaturesIdentifier, render_trace_arena, CallTraceDecoder, + CallTraceDecoderBuilder, TraceKind, Traces, }, }; use forge_verify::RetryArgs; @@ -42,6 +42,7 @@ use foundry_evm::{ constants::DEFAULT_CREATE2_DEPLOYER, decode::RevertDecoder, inspectors::cheatcodes::{BroadcastableTransaction, BroadcastableTransactions}, + traces::identifier::TraceIdentifiers, }; use foundry_wallets::MultiWalletOpts; use futures::future; @@ -198,33 +199,29 @@ impl ScriptArgs { result: &mut ScriptResult, known_contracts: &ContractsByArtifact, ) -> Result { - let verbosity = script_config.evm_opts.verbosity; - let mut etherscan_identifier = EtherscanIdentifier::new( - &script_config.config, - script_config.evm_opts.get_remote_chain_id(), - )?; - - let mut local_identifier = LocalTraceIdentifier::new(known_contracts); let mut decoder = CallTraceDecoderBuilder::new() .with_labels(result.labeled_addresses.clone()) - .with_verbosity(verbosity) - .with_local_identifier_abis(&local_identifier) + .with_verbosity(script_config.evm_opts.verbosity) + .with_known_contracts(known_contracts) .with_signature_identifier(SignaturesIdentifier::new( Config::foundry_cache_dir(), script_config.config.offline, )?) .build(); + let mut identifier = TraceIdentifiers::new() + .with_local(known_contracts) + .with_etherscan(&script_config.config, script_config.evm_opts.get_remote_chain_id())?; // Decoding traces using etherscan is costly as we run into rate limits, // causing scripts to run for a very long time unnecessarily. // Therefore, we only try and use etherscan if the user has provided an API key. let should_use_etherscan_traces = script_config.config.etherscan_api_key.is_some(); + if !should_use_etherscan_traces { + identifier.etherscan = None; + } for (_, trace) in &mut result.traces { - decoder.identify(trace, &mut local_identifier); - if should_use_etherscan_traces { - decoder.identify(trace, &mut etherscan_identifier); - } + decoder.identify(trace, &mut identifier); } Ok(decoder) } diff --git a/crates/forge/bin/cmd/test/mod.rs b/crates/forge/bin/cmd/test/mod.rs index 2a0e7bbdbbfa..3e7d148484c4 100644 --- a/crates/forge/bin/cmd/test/mod.rs +++ b/crates/forge/bin/cmd/test/mod.rs @@ -7,10 +7,7 @@ use forge::{ gas_report::GasReport, inspectors::CheatsConfig, result::{SuiteResult, TestOutcome, TestStatus}, - traces::{ - identifier::{EtherscanIdentifier, LocalTraceIdentifier, SignaturesIdentifier}, - CallTraceDecoderBuilder, TraceKind, - }, + traces::{identifier::SignaturesIdentifier, CallTraceDecoderBuilder, TraceKind}, MultiContractRunner, MultiContractRunnerBuilder, TestOptions, TestOptionsBuilder, }; use foundry_cli::{ @@ -31,6 +28,7 @@ use foundry_config::{ get_available_profiles, Config, }; use foundry_debugger::Debugger; +use foundry_evm::traces::identifier::TraceIdentifiers; use regex::Regex; use std::{sync::mpsc::channel, time::Instant}; use watchexec::config::{InitConfig, RuntimeConfig}; @@ -296,9 +294,10 @@ impl TestArgs { // Set up trace identifiers. let known_contracts = runner.known_contracts.clone(); - let mut local_identifier = LocalTraceIdentifier::new(&known_contracts); let remote_chain_id = runner.evm_opts.get_remote_chain_id(); - let mut etherscan_identifier = EtherscanIdentifier::new(&config, remote_chain_id)?; + let mut identifier = TraceIdentifiers::new() + .with_local(&known_contracts) + .with_etherscan(&config, remote_chain_id)?; // Run tests. let (tx, rx) = channel::<(String, SuiteResult)>(); @@ -313,7 +312,7 @@ impl TestArgs { // Build the trace decoder. let mut builder = CallTraceDecoderBuilder::new() - .with_local_identifier_abis(&local_identifier) + .with_known_contracts(&known_contracts) .with_verbosity(verbosity); // Signatures are of no value for gas reports. if !self.gas_report { @@ -379,8 +378,7 @@ impl TestArgs { let mut decoded_traces = Vec::with_capacity(result.traces.len()); for (kind, arena) in &result.traces { if identify_addresses { - decoder.identify(arena, &mut local_identifier); - decoder.identify(arena, &mut etherscan_identifier); + decoder.identify(arena, &mut identifier); } // verbosity: