diff --git a/Cargo.lock b/Cargo.lock index 8d94667f..e84a7c7d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -77,6 +77,15 @@ dependencies = [ "term", ] +[[package]] +name = "async-convert" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6d416feee97712e43152cd42874de162b8f9b77295b1c85e5d92725cc8310bae" +dependencies = [ + "async-trait", +] + [[package]] name = "async-openai" version = "0.10.3" @@ -1713,6 +1722,7 @@ dependencies = [ name = "heimdall-core" version = "0.6.5" dependencies = [ + "async-convert", "backtrace", "clap", "clap-verbosity-flag", @@ -1721,6 +1731,7 @@ dependencies = [ "derive_builder", "ethers", "fancy-regex", + "futures", "heimdall-cache", "heimdall-common", "heimdall-config", diff --git a/common/src/resources/transpose.rs b/common/src/resources/transpose.rs index 4688cc1d..7afc5987 100644 --- a/common/src/resources/transpose.rs +++ b/common/src/resources/transpose.rs @@ -21,6 +21,7 @@ struct TransposeResponse { } /// executes a transpose SQL query and returns the response +/// TODO: exponential backoff async fn call_transpose(query: &str, api_key: &str) -> Option { // get a new logger let logger = Logger::default(); diff --git a/core/Cargo.toml b/core/Cargo.toml index f5621d5e..6f63cff9 100644 --- a/core/Cargo.toml +++ b/core/Cargo.toml @@ -28,3 +28,5 @@ strsim = "0.10.0" tokio = {version = "1", features = ["full"]} tui = "0.19" derive_builder = "0.12.0" +async-convert = "1.0.0" +futures = "0.3.28" diff --git a/core/src/error.rs b/core/src/error.rs index 5ad3021e..a7827801 100644 --- a/core/src/error.rs +++ b/core/src/error.rs @@ -12,4 +12,6 @@ pub enum Error { RpcError(String), #[error("Error: {0}")] GenericError(String), + #[error("TransposeError: {0}")] + TransposeError(String), } diff --git a/core/src/inspect/core/contracts.rs b/core/src/inspect/core/contracts.rs new file mode 100644 index 00000000..c487bd71 --- /dev/null +++ b/core/src/inspect/core/contracts.rs @@ -0,0 +1,75 @@ +use std::collections::{HashMap, HashSet}; + +use ethers::types::Address; +use futures::future::try_join_all; + +use crate::{error::Error, inspect::InspectArgs}; +use heimdall_common::{resources::transpose::get_label, utils::hex::ToLowerHex}; + +#[derive(Debug, Clone)] +pub struct Contracts { + pub contracts: HashMap, + transpose_api_key: Option, +} + +impl Contracts { + pub fn new(args: &InspectArgs) -> Self { + Self { contracts: HashMap::new(), transpose_api_key: args.transpose_api_key.clone() } + } + + pub async fn add(&mut self, address: Address) -> Result<(), Error> { + // if alias already exists, don't overwrite + if self.contracts.contains_key(&address) { + return Ok(()); + } + + if let Some(transpose_api_key) = &self.transpose_api_key { + self.contracts.insert( + address, + get_label(&address.to_lower_hex(), transpose_api_key) + .await + .unwrap_or(address.to_lower_hex()), + ); + } else { + self.contracts.insert(address, address.to_lower_hex()); + } + + Ok(()) + } + + pub async fn extend(&mut self, addresses: HashSet
) -> Result<(), Error> { + // for each address, get the label + if let Some(transpose_api_key) = &self.transpose_api_key { + let handles: Vec<_> = addresses + .clone() + .into_iter() + .map(move |address| { + let transpose_api_key = transpose_api_key.clone(); + tokio::spawn(async move { + get_label(&address.to_lower_hex(), &transpose_api_key).await + }) + }) + .collect(); + + let labels = + try_join_all(handles).await.map_err(|e| Error::TransposeError(e.to_string()))?; + + self.contracts.extend( + addresses + .into_iter() + .zip(labels.into_iter()) + .map(|(address, label)| (address, label.unwrap_or(address.to_lower_hex()))), + ); + // replace None + } else { + self.contracts + .extend(addresses.into_iter().map(|address| (address, address.to_lower_hex()))); + } + + Ok(()) + } + + pub fn get(&self, address: Address) -> Option<&String> { + self.contracts.get(&address) + } +} diff --git a/core/src/inspect/core/mod.rs b/core/src/inspect/core/mod.rs index 5ee45fbc..ed583ff9 100644 --- a/core/src/inspect/core/mod.rs +++ b/core/src/inspect/core/mod.rs @@ -1 +1,2 @@ +pub(crate) mod contracts; pub mod tracing; diff --git a/core/src/inspect/core/tracing.rs b/core/src/inspect/core/tracing.rs index cbfb854a..d4da4517 100644 --- a/core/src/inspect/core/tracing.rs +++ b/core/src/inspect/core/tracing.rs @@ -1,219 +1,292 @@ -use std::collections::HashMap; - -use ethers::types::{Transaction, TransactionTrace}; -use heimdall_common::{ - resources::transpose::get_label, - utils::{ - hex::ToLowerHex, - io::{logging::TraceFactory, types::Parameterize}, +use std::collections::{HashSet, VecDeque}; + +use ethers::{ + abi::Token, + types::{ + ActionType, Address, Bytes, Call, CallResult, CallType, Create, CreateResult, Reward, + Suicide, TransactionTrace, U256, }, }; +use heimdall_common::ether::signatures::ResolvedFunction; +use serde::{Deserialize, Serialize}; -use crate::{decode::DecodeArgsBuilder, error::Error, inspect::InspectArgs}; - -/// Converts raw [`TransactionTrace`]s to human-readable [`TraceFactory`] -pub async fn build_trace_display( - args: &InspectArgs, - transaction: &Transaction, - transaction_traces: Vec, - address_labels: &mut HashMap>, -) -> Result { - let mut trace = TraceFactory::default(); - let decode_call = trace.add_call_with_extra( - 0, - transaction.gas.as_u32(), // panicky - "heimdall".to_string(), - "inspect".to_string(), - vec![transaction.hash.to_lower_hex()], - "()".to_string(), - vec![format!("{} wei", transaction.value)], - ); - - let mut trace_indices = HashMap::new(); - - for transaction_trace in transaction_traces { - let trace_address = transaction_trace - .trace_address - .iter() - .map(|address| address.to_string()) - .collect::>() - .join("."); - let parent_address = trace_address - .split('.') - .take(trace_address.split('.').count() - 1) - .collect::>() - .join("."); - - // get trace index from parent_address - let parent_index = trace_indices.get(&parent_address).unwrap_or(&decode_call); - - // get result - let mut result_str = "()".to_string(); - if let Some(result) = transaction_trace.result { - result_str = match result { - ethers::types::Res::Call(res) => { - // we can attempt to decode this as if it is calldata, we just need to add some - // 4byte prefix. - let output = - format!("0x00000000{}", res.output.to_string().replacen("0x", "", 1)); - let result = crate::decode::decode( - DecodeArgsBuilder::new() - .target(output) - .skip_resolving(true) - .build() - .map_err(|_e| Error::DecodeError)?, - ) - .await?; - - // get first result - if let Some(resolved_function) = result.first() { - resolved_function - .decoded_inputs - .clone() - .unwrap_or_default() - .iter() - .map(|token| token.parameterize()) - .collect::>() - .join(", ") - } else { - res.output.to_string() - } - } - ethers::types::Res::Create(res) => res.address.to_lower_hex(), - ethers::types::Res::None => "()".to_string(), +use crate::{decode::DecodeArgsBuilder, error::Error}; +use async_convert::{async_trait, TryFrom}; +use futures::future::try_join_all; + +/// Decoded Trace +#[derive(Debug, PartialEq, Clone, Deserialize, Serialize)] +pub struct DecodedTransactionTrace { + pub trace_address: Vec, + pub action: DecodedAction, + pub action_type: ActionType, + pub result: Option, + pub error: Option, + pub subtraces: Vec, +} + +/// Decoded Action +#[derive(Debug, PartialEq, Clone, Serialize, Deserialize)] +#[serde(untagged, rename_all = "lowercase")] +pub enum DecodedAction { + /// Decoded Call + Call(DecodedCall), + /// Create + Create(Create), + /// Suicide + Suicide(Suicide), + /// Reward + Reward(Reward), +} + +/// Decoded Call +#[derive(Debug, Clone, PartialEq, Default, Deserialize, Serialize)] +pub struct DecodedCall { + /// Sender + pub from: Address, + /// Recipient + pub to: Address, + /// Transferred Value + pub value: U256, + /// Gas + pub gas: U256, + /// Input data + pub input: Bytes, + /// The type of the call. + #[serde(rename = "callType")] + pub call_type: CallType, + /// Potential resolved function + #[serde(rename = "resolvedFunction")] + pub resolved_function: Option, + /// Decoded inputs + #[serde(rename = "decodedInputs")] + pub decoded_inputs: Vec, +} + +/// Decoded Response +#[derive(Debug, Default, Clone, PartialEq, Deserialize, Serialize)] +#[serde(untagged)] +pub enum DecodedRes { + /// Call + Call(DecodedCallResult), + /// Create + Create(CreateResult), + /// None + #[default] + None, +} + +/// Call Result +#[derive(Debug, Clone, PartialEq, Default, Deserialize, Serialize)] +pub struct DecodedCallResult { + /// Gas used + #[serde(rename = "gasUsed")] + pub gas_used: U256, + /// Output bytes + pub output: Bytes, + /// Decoded outputs + #[serde(rename = "decodedOutputs")] + pub decoded_outputs: Vec, +} + +#[async_trait] +impl TryFrom> for DecodedTransactionTrace { + type Error = crate::error::Error; + + async fn try_from(value: Vec) -> Result { + // convert each [`TransactionTrace`] to a [`DecodedTransactionTrace`] + let handles = value.into_iter().map(|trace| { + >::try_from(trace) + }); + let mut decoded_transaction_traces = VecDeque::from(try_join_all(handles).await?); + + // get the first trace, this will be the one we are building. + let mut decoded_transaction_trace = + decoded_transaction_traces.pop_front().ok_or(Error::DecodeError)?; + assert!(decoded_transaction_trace.trace_address.is_empty()); // sanity check + + for decoded_trace in decoded_transaction_traces { + // trace_address is the index of the trace in the decoded_transaction_trace. for + // example, if trace_address is `[0]`, it'll be added to + // `decoded_transaction_trace.subtraces` at index 0. if trace_address is `[0, 0]`, it'll + // be added to `decoded_transaction_trace.subtraces[0].subtraces` at index 0. + let mut current_trace = &mut decoded_transaction_trace; + let trace_address = &decoded_trace.trace_address; + + // Iterate through the trace address, navigating through subtraces + for &index in trace_address.iter().take(trace_address.len() - 1) { + current_trace = current_trace.subtraces.get_mut(index).ok_or(Error::DecodeError)?; + // You might need to define this error + } + + // Insert the decoded trace into the correct position + if let Some(last_index) = trace_address.last() { + current_trace.subtraces.insert(*last_index, decoded_trace); + } else { + return Err(Error::DecodeError); // Trace address cannot be empty here } } - if result_str.replacen("0x", "", 1).is_empty() { - result_str = "()".to_string(); + + Ok(decoded_transaction_trace) + } +} + +#[async_trait] +impl TryFrom for DecodedTransactionTrace { + type Error = crate::error::Error; + + async fn try_from(value: TransactionTrace) -> Result { + let action = match value.action { + ethers::types::Action::Call(call) => DecodedAction::Call( + >::try_from(call).await?, + ), + ethers::types::Action::Create(create) => DecodedAction::Create(create), + ethers::types::Action::Suicide(suicide) => DecodedAction::Suicide(suicide), + ethers::types::Action::Reward(reward) => DecodedAction::Reward(reward), + }; + + let result = match value.result { + Some(res) => match res { + ethers::types::Res::Call(call) => Some(DecodedRes::Call( + >::try_from(call) + .await?, + )), + ethers::types::Res::Create(create) => Some(DecodedRes::Create(create)), + ethers::types::Res::None => Some(DecodedRes::None), + }, + None => None, + }; + + Ok(Self { + trace_address: value.trace_address, + action, + action_type: value.action_type, + result, + error: value.error, + subtraces: Vec::new(), // we will build this later + }) + } +} + +#[async_trait] +impl TryFrom for DecodedCall { + type Error = crate::error::Error; + + async fn try_from(value: Call) -> Result { + let calldata = value.input.to_string().replacen("0x", "", 1); + let mut decoded_inputs = Vec::new(); + let mut resolved_function = None; + + if !calldata.is_empty() { + let result = crate::decode::decode( + DecodeArgsBuilder::new() + .target(calldata) + .build() + .map_err(|_e| Error::DecodeError)?, + ) + .await?; + + if let Some(first_result) = result.first() { + decoded_inputs = first_result.decoded_inputs.clone().unwrap_or_default(); + resolved_function = Some(first_result.clone()); + } } - // get action - match transaction_trace.action { - ethers::types::Action::Call(call) => { - // add address label. we will use this to display the address in the trace, if - // available. (requires `transpose_api_key`) - if let Some(transpose_api_key) = &args.transpose_api_key { - if let std::collections::hash_map::Entry::Vacant(e) = - address_labels.entry(call.to.to_lower_hex()) - { - e.insert(get_label(&call.to.to_lower_hex(), transpose_api_key).await); - } - } - let address_label = address_labels - .get(&call.to.to_lower_hex()) - .unwrap_or(&None) - .clone() - .unwrap_or(call.to.to_lower_hex()) - .to_string(); - - // build extra_data, which will be used to display the call type and value transfer - // information - let mut extra_data = vec![]; - let call_type = match call.call_type { - ethers::types::CallType::Call => "call", - ethers::types::CallType::DelegateCall => "delegatecall", - ethers::types::CallType::StaticCall => "staticcall", - ethers::types::CallType::CallCode => "callcode", - ethers::types::CallType::None => "none", + Ok(Self { + from: value.from, + to: value.to, + value: value.value, + gas: value.gas, + input: value.input, + call_type: value.call_type, + resolved_function, + decoded_inputs, + }) + } +} + +#[async_trait] +impl TryFrom for DecodedCallResult { + type Error = crate::error::Error; + + async fn try_from(value: CallResult) -> Result { + // we can attempt to decode this as if it is calldata, we just need to add some + // 4byte prefix. + let output = format!("0x00000000{}", value.output.to_string().replacen("0x", "", 1)); + let result = crate::decode::decode( + DecodeArgsBuilder::new() + .target(output) + .skip_resolving(true) + .build() + .map_err(|_e| Error::DecodeError)?, + ) + .await?; + + // get first result + let decoded_outputs = if let Some(resolved_function) = result.first() { + resolved_function.decoded_inputs.clone().unwrap_or_default() + } else { + vec![] + }; + + Ok(Self { gas_used: value.gas_used, output: value.output, decoded_outputs }) + } +} + +impl DecodedTransactionTrace { + /// Returns a [`HashSet`] of all addresses involved in the traced transaction. if + /// `include_inputs`/`include_outputs` is true, the [`HashSet`] will also include the + /// addresses of the inputs/outputs of the transaction. + pub fn addresses(&self, include_inputs: bool, include_outputs: bool) -> HashSet
{ + let mut addresses = HashSet::new(); + + match &self.action { + DecodedAction::Call(call) => { + addresses.insert(call.from); + addresses.insert(call.to); + + if include_inputs { + let _ = call.decoded_inputs.iter().map(|token| match token { + Token::Address(address) => addresses.insert(*address), + _ => false, + }); } - .to_string(); - extra_data.push(call_type.clone()); - if !call.value.is_zero() { - extra_data.push(format!("{} wei", call.value)); + if include_outputs { + let _ = self.result.iter().map(|result| { + if let DecodedRes::Call(call_result) = result { + let _ = call_result.decoded_outputs.iter().map(|token| match token { + Token::Address(address) => addresses.insert(*address), + _ => false, + }); + } + }); } + } + DecodedAction::Create(create) => { + addresses.insert(create.from); - // attempt to decode calldata - let calldata = call.input.to_string(); - if !calldata.replacen("0x", "", 1).is_empty() { - let result = crate::decode::decode( - DecodeArgsBuilder::new() - .target(calldata) - .build() - .map_err(|_e| Error::DecodeError)?, - ) - .await?; - - // get first result - if let Some(resolved_function) = result.first() { - // convert decoded inputs Option> to Vec - let decoded_inputs = - resolved_function.decoded_inputs.clone().unwrap_or_default(); - - // get index of parent - let parent_index = trace.add_call_with_extra( - *parent_index, - call.gas.as_u32(), // panicky - address_label, - resolved_function.name.clone(), - vec![decoded_inputs - .iter() - .map(|token| token.parameterize()) - .collect::>() - .join(", ")], - result_str, - extra_data, - ); - - // add trace_address to trace_indices - trace_indices.insert(trace_address.clone(), parent_index); - } else { - // get index of parent - let parent_index = trace.add_call_with_extra( - *parent_index, - call.gas.as_u32(), // panicky - address_label, - "unknown".to_string(), - vec![format!("bytes: {}", call.input.to_string())], - result_str, - extra_data, - ); - - // add trace_address to trace_indices - trace_indices.insert(trace_address.clone(), parent_index); - } - } else { - // value transfer - trace.add_call_with_extra( - *parent_index, - call.gas.as_u32(), // panicky - call.to.to_lower_hex(), - "fallback".to_string(), - vec![], - result_str, - extra_data, - ); + if include_outputs { + let _ = self.result.iter().map(|result| { + if let DecodedRes::Create(create_result) = result { + addresses.insert(create_result.address); + } + }); } } - ethers::types::Action::Create(create) => { - // add address label. we will use this to display the address in the trace, if - // available. (requires `transpose_api_key`) - if let Some(transpose_api_key) = &args.transpose_api_key { - if !address_labels.contains_key(&result_str) { - address_labels.insert( - result_str.clone(), - get_label(&result_str, transpose_api_key).await, - ); - } - } - let address_label = address_labels - .get(&result_str) - .unwrap_or(&None) - .clone() - .unwrap_or("NewContract".to_string()) - .to_string(); - - trace.add_creation( - *parent_index, - create.gas.as_u32(), - address_label, - result_str, - create.init.len().try_into().map_err(|_e| Error::DecodeError)?, - ); + DecodedAction::Suicide(suicide) => { + addresses.insert(suicide.address); + addresses.insert(suicide.refund_address); + } + DecodedAction::Reward(reward) => { + addresses.insert(reward.author); } - ethers::types::Action::Suicide(_suicide) => {} - ethers::types::Action::Reward(_) => todo!(), + }; + + // add all addresses found in subtraces + for subtrace in &self.subtraces { + addresses.extend(subtrace.addresses(include_inputs, include_outputs)) } - } - Ok(trace) + addresses + } } diff --git a/core/src/inspect/mod.rs b/core/src/inspect/mod.rs index 0a5181fc..daf4a19f 100644 --- a/core/src/inspect/mod.rs +++ b/core/src/inspect/mod.rs @@ -1,17 +1,18 @@ mod core; -use std::collections::HashMap; - use clap::{AppSettings, Parser}; use derive_builder::Builder; +use ethers::types::TransactionTrace; use heimdall_common::{ ether::rpc::{get_trace, get_transaction}, utils::io::logging::Logger, }; -use crate::{error::Error, inspect::core::tracing::build_trace_display}; +use crate::error::Error; + +use self::core::{contracts::Contracts, tracing::DecodedTransactionTrace}; #[derive(Debug, Clone, Parser, Builder)] #[clap( @@ -54,14 +55,15 @@ impl InspectArgsBuilder { } } +#[derive(Debug, Clone)] +pub struct InspectResult { + pub decoded_trace: Option, +} /// The entrypoint for the inspect module. This function will analyze the given transaction and /// provide a detailed inspection of the transaction, including calldata & trace decoding, log /// visualization, and more. #[allow(deprecated)] -pub async fn inspect(args: InspectArgs) -> Result<(), Error> { - // define - let mut address_labels = HashMap::new(); - +pub async fn inspect(args: InspectArgs) -> Result { // set logger environment variable if not already set // TODO: abstract this to a heimdall_common util if std::env::var("RUST_LOG").is_err() { @@ -75,13 +77,13 @@ pub async fn inspect(args: InspectArgs) -> Result<(), Error> { } // get a new logger and trace - let (_logger, _trace) = Logger::new(match args.verbose.log_level() { + let (logger, _trace) = Logger::new(match args.verbose.log_level() { Some(level) => level.as_str(), None => "SILENT", }); // get calldata from RPC - let transaction = get_transaction(&args.target, &args.rpc_url) + let _transaction = get_transaction(&args.target, &args.rpc_url) .await .map_err(|e| Error::RpcError(e.to_string()))?; @@ -89,18 +91,29 @@ pub async fn inspect(args: InspectArgs) -> Result<(), Error> { let block_trace = get_trace(&args.target, &args.rpc_url).await.map_err(|e| Error::RpcError(e.to_string()))?; - // build displayable trace - let transaction_trace_display = build_trace_display( - &args, - &transaction, - block_trace.trace.unwrap_or_default(), - &mut address_labels, - ) - .await?; + let decoded_trace = + match block_trace.trace { + Some(trace) => , + >>::try_from(trace) + .await + .ok(), + None => None, + }; + if decoded_trace.is_none() { + logger.warn("no trace found for transaction"); + } - transaction_trace_display.display(); + // get contracts client and extend with addresses from trace + let mut contracts = Contracts::new(&args); + if let Some(decoded_trace) = decoded_trace.clone() { + contracts + .extend(decoded_trace.addresses(true, true).into_iter().collect()) + .await + .map_err(|e| Error::GenericError(e.to_string()))?; + }; - println!("{:#?}", block_trace.state_diff); + println!("{:#?}", contracts); - Ok(()) + Ok(InspectResult { decoded_trace }) }