diff --git a/common/src/ether/bytecode.rs b/common/src/ether/bytecode.rs new file mode 100644 index 00000000..e577d018 --- /dev/null +++ b/common/src/ether/bytecode.rs @@ -0,0 +1,87 @@ +use super::rpc::get_code; +use crate::{ + constants::{ADDRESS_REGEX, BYTECODE_REGEX}, + debug_max, + utils::io::logging::Logger, +}; +use std::fs; + +pub async fn get_bytecode_from_target( + target: &str, + rpc_url: &str, +) -> Result> { + let (logger, _) = Logger::new(""); + + if ADDRESS_REGEX.is_match(target)? { + // Target is a contract address, so we need to fetch the bytecode from the RPC provider. + get_code(target, rpc_url).await + } else if BYTECODE_REGEX.is_match(target)? { + debug_max!("using provided bytecode for snapshotting."); + + // Target is already a bytecode, so we just need to remove 0x from the begining + Ok(target.replacen("0x", "", 1)) + } else { + debug_max!("using provided file for snapshotting."); + + // Target is a file path, so we need to read the bytecode from the file. + match fs::read_to_string(target) { + Ok(contents) => { + let _contents = contents.replace('\n', ""); + if BYTECODE_REGEX.is_match(&_contents)? && _contents.len() % 2 == 0 { + Ok(_contents.replacen("0x", "", 1)) + } else { + logger.error(&format!("file '{}' doesn't contain valid bytecode.", &target)); + std::process::exit(1) + } + } + Err(_) => { + logger.error(&format!("failed to open file '{}' .", &target)); + std::process::exit(1) + } + } + } +} +#[cfg(test)] +mod tests { + use super::*; + use std::fs; + + #[tokio::test] + async fn test_get_bytecode_when_target_is_address() { + let bytecode = get_bytecode_from_target( + "0x9f00c43700bc0000Ff91bE00841F8e04c0495000", + "https://rpc.ankr.com/eth", + ) + .await + .unwrap(); + + assert!(BYTECODE_REGEX.is_match(&bytecode).unwrap()); + } + + #[tokio::test] + async fn test_get_bytecode_when_target_is_bytecode() { + let bytecode = get_bytecode_from_target( + "0x00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000001", + "https://rpc.ankr.com/eth", + ) + .await + .unwrap(); + + assert!(BYTECODE_REGEX.is_match(&bytecode).unwrap()); + } + + #[tokio::test] + async fn test_get_bytecode_when_target_is_file_path() { + let file_path = "./mock-file.txt"; + let mock_bytecode = "0x00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000001"; + + fs::write(file_path, mock_bytecode).unwrap(); + + let bytecode = + get_bytecode_from_target(file_path, "https://rpc.ankr.com/eth").await.unwrap(); + + assert!(BYTECODE_REGEX.is_match(&bytecode).unwrap()); + + fs::remove_file(file_path).unwrap(); + } +} diff --git a/common/src/ether/mod.rs b/common/src/ether/mod.rs index dc484173..d6e3ac26 100644 --- a/common/src/ether/mod.rs +++ b/common/src/ether/mod.rs @@ -1,3 +1,4 @@ +pub mod bytecode; pub mod compiler; pub mod evm; pub mod lexers; diff --git a/common/src/ether/rpc.rs b/common/src/ether/rpc.rs index 8ff61491..a81f795c 100644 --- a/common/src/ether/rpc.rs +++ b/common/src/ether/rpc.rs @@ -147,7 +147,7 @@ pub async fn get_code( ) .map_err(|_| logger.error(&format!("failed to cache bytecode for contract: {:?}", &contract_address))); - Ok(bytecode_as_bytes.to_string()) + Ok(bytecode_as_bytes.to_string().replacen("0x", "", 1)) }) .await .map_err(|_| Box::from("failed to fetch bytecode")) diff --git a/common/src/ether/selectors.rs b/common/src/ether/selectors.rs index 0de14478..d171a724 100644 --- a/common/src/ether/selectors.rs +++ b/common/src/ether/selectors.rs @@ -9,9 +9,41 @@ use tokio::task; use crate::utils::{io::logging::Logger, strings::decode_hex}; -use super::{evm::core::vm::VM, signatures::ResolveSelector}; +use super::{ + evm::core::vm::VM, + signatures::{ResolveSelector, ResolvedFunction}, +}; use crate::debug_max; +// Find all function selectors and all the data associated to this function, represented by +// [`ResolvedFunction`] +pub async fn get_resolved_selectors( + disassembled_bytecode: &str, + skip_resolving: &bool, + evm: &VM, +) -> Result< + (HashMap, HashMap>), + Box, +> { + let selectors = find_function_selectors(evm, disassembled_bytecode); + + let mut resolved_selectors = HashMap::new(); + if !skip_resolving { + resolved_selectors = + resolve_selectors::(selectors.keys().cloned().collect()).await; + + debug_max!(&format!( + "resolved {} possible functions from {} detected selectors.", + resolved_selectors.len(), + selectors.len() + )); + } else { + debug_max!(&format!("found {} possible function selectors.", selectors.len())); + } + + Ok((selectors, resolved_selectors)) +} + /// find all function selectors in the given EVM assembly. pub fn find_function_selectors(evm: &VM, assembly: &str) -> HashMap { let mut function_selectors = HashMap::new(); diff --git a/common/src/utils/io/logging.rs b/common/src/utils/io/logging.rs index af325b19..70d12d6d 100644 --- a/common/src/utils/io/logging.rs +++ b/common/src/utils/io/logging.rs @@ -673,11 +673,33 @@ impl Logger { } } +/// Set `RUST_LOG` variable to env if does not exist +/// +/// ``` +/// use heimdall_common::utils::io::logging::set_logger_env; +/// +/// let verbosity = clap_verbosity_flag::Verbosity::new(-1, 0); +/// set_logger_env(&verbosity); +/// ``` +pub fn set_logger_env(verbosity: &clap_verbosity_flag::Verbosity) { + let env_not_set = std::env::var("RUST_LOG").is_err(); + + if env_not_set { + let log_level = match verbosity.log_level() { + Some(level) => level.as_str(), + None => "SILENT", + }; + + std::env::set_var("RUST_LOG", log_level); + } +} + #[cfg(test)] mod tests { use std::time::Instant; use super::*; + use std::env; #[test] fn test_raw_trace() { @@ -956,4 +978,15 @@ mod tests { let (_logger, _) = Logger::new("MAX"); debug_max!("log"); } + + #[test] + fn test_set_logger_env_default() { + env::remove_var("RUST_LOG"); + + let verbosity = clap_verbosity_flag::Verbosity::new(-1, 0); + + set_logger_env(&verbosity); + + assert_eq!(env::var("RUST_LOG").unwrap(), "SILENT"); + } } diff --git a/common/src/utils/strings.rs b/common/src/utils/strings.rs index 168093d4..16423e83 100644 --- a/common/src/utils/strings.rs +++ b/common/src/utils/strings.rs @@ -352,6 +352,26 @@ pub fn classify_token(token: &str) -> TokenType { TokenType::Function } +/// Returns a collapsed version of a string if this string is greater than 66 characters in length. +/// The collapsed string consists of the first 66 characters, followed by an ellipsis ("..."), and +/// then the last 16 characters of the original string. ``` +/// use heimdall_common::utils::strings::get_shortned_target; +/// +/// let long_target = "0".repeat(80); +/// let shortened_target = get_shortned_target(&long_target); +/// ``` +pub fn get_shortned_target(target: &str) -> String { + let mut shortened_target = target.to_string(); + + if shortened_target.len() > 66 { + shortened_target = shortened_target.chars().take(66).collect::() + + "..." + + &shortened_target.chars().skip(shortened_target.len() - 16).collect::(); + } + + shortened_target +} + #[cfg(test)] mod tests { use ethers::types::{I256, U256}; @@ -691,4 +711,20 @@ mod tests { assert_eq!(classification, TokenType::Function); } } + + #[test] + fn test_shorten_long_target() { + let long_target = "0".repeat(80); + let shortened_target = get_shortned_target(&long_target); + + assert_eq!(shortened_target.len(), 85); + } + + #[test] + fn test_shorten_short_target() { + let short_target = "0".repeat(66); + let shortened_target = get_shortned_target(&short_target); + + assert_eq!(shortened_target.len(), 66); + } } diff --git a/core/src/snapshot/mod.rs b/core/src/snapshot/mod.rs index edb787bb..14b838a9 100644 --- a/core/src/snapshot/mod.rs +++ b/core/src/snapshot/mod.rs @@ -8,24 +8,22 @@ use heimdall_common::debug_max; use std::{ collections::{HashMap, HashSet}, - fs, time::Duration, }; use clap::{AppSettings, Parser}; use derive_builder::Builder; use heimdall_common::{ - constants::{ADDRESS_REGEX, BYTECODE_REGEX}, ether::{ + bytecode::get_bytecode_from_target, compiler::detect_compiler, evm::core::vm::VM, - rpc::get_code, - selectors::{find_function_selectors, resolve_selectors}, - signatures::{score_signature, ResolvedError, ResolvedFunction, ResolvedLog}, + selectors::get_resolved_selectors, + signatures::{ResolvedError, ResolvedFunction, ResolvedLog}, }, utils::{ io::logging::*, - strings::{decode_hex, encode_hex_reduced}, + strings::{decode_hex, get_shortned_target}, }, }; use indicatif::ProgressBar; @@ -34,7 +32,7 @@ use crate::{ disassemble::{disassemble, DisassemblerArgs}, snapshot::{ analyze::snapshot_trace, - resolve::match_parameters, + resolve::resolve_signatures, structures::snapshot::{GasUsed, Snapshot}, util::tui, }, @@ -107,33 +105,15 @@ pub struct SnapshotResult { /// signatures, access control, gas consumption, storage accesses, event emissions, and more. pub async fn snapshot(args: SnapshotArgs) -> Result> { use std::time::Instant; - let now = Instant::now(); - // set logger environment variable if not already set - if std::env::var("RUST_LOG").is_err() { - std::env::set_var( - "RUST_LOG", - match args.verbose.log_level() { - Some(level) => level.as_str(), - None => "SILENT", - }, - ); - } + set_logger_env(&args.verbose); + let now = Instant::now(); let (logger, mut trace) = Logger::new(match args.verbose.log_level() { Some(level) => level.as_str(), None => "SILENT", }); - let mut all_resolved_events: HashMap = HashMap::new(); - let mut all_resolved_errors: HashMap = HashMap::new(); - - // truncate target for prettier display - let mut shortened_target = args.target.clone(); - if shortened_target.len() > 66 { - shortened_target = shortened_target.chars().take(66).collect::() + - "..." + - &shortened_target.chars().skip(shortened_target.len() - 16).collect::(); - } + let shortened_target = get_shortned_target(&args.target); let snapshot_call = trace.add_call( 0, line!(), @@ -143,54 +123,7 @@ pub async fn snapshot(args: SnapshotArgs) -> Result { - let _contents = contents.replace('\n', ""); - if BYTECODE_REGEX.is_match(&_contents)? && _contents.len() % 2 == 0 { - _contents.replacen("0x", "", 1) - } else { - logger - .error(&format!("file '{}' doesn't contain valid bytecode.", &args.target)); - std::process::exit(1) - } - } - Err(_) => { - logger.error(&format!("failed to open file '{}' .", &args.target)); - std::process::exit(1) - } - }; - } - - // disassemble the bytecode - let disassembled_bytecode = disassemble(DisassemblerArgs { - target: contract_bytecode.clone(), - verbose: args.verbose.clone(), - rpc_url: args.rpc_url, - decimal_counter: false, - name: args.name, - output: String::new(), - }) - .await?; - trace.add_call( - snapshot_call, - line!(), - "heimdall".to_string(), - "disassemble".to_string(), - vec![format!("{} bytes", contract_bytecode.len() / 2usize)], - "()".to_string(), - ); + let contract_bytecode = get_bytecode_from_target(&args.target, &args.rpc_url).await?; // perform versioning and compiler heuristics let (compiler, version) = detect_compiler(&contract_bytecode); @@ -210,7 +143,6 @@ pub async fn snapshot(args: SnapshotArgs) -> Result Result 66 { - shortened_target = shortened_target.chars().take(66).collect::() + - "..." + - &shortened_target.chars().skip(shortened_target.len() - 16).collect::(); - } + let shortened_target = get_shortned_target(&contract_bytecode); let vm_trace = trace.add_creation( snapshot_call, line!(), @@ -234,40 +161,83 @@ pub async fn snapshot(args: SnapshotArgs) -> Result(selectors.keys().cloned().collect()).await; + let disassembled_bytecode = disassemble(DisassemblerArgs { + rpc_url: args.rpc_url.clone(), + verbose: args.verbose.clone(), + target: args.target.clone(), + name: args.name.clone(), + decimal_counter: false, + output: String::new(), + }) + .await?; - // if resolved selectors are empty, we can't perform symbolic execution - if resolved_selectors.is_empty() { - logger.error(&format!( - "failed to resolve any function selectors from '{shortened_target}' .", - shortened_target = shortened_target - )); - } + let (selectors, resolved_selectors) = + get_resolved_selectors(&disassembled_bytecode, &args.skip_resolving, &evm).await?; + + let (snapshots, all_resolved_errors, all_resolved_events) = get_snapshots( + selectors, + resolved_selectors, + &contract_bytecode, + &logger, + &mut trace, + vm_trace, + &evm, + &args, + ) + .await?; - logger.info(&format!( - "resolved {} possible functions from {} detected selectors.", - resolved_selectors.len(), - selectors.len() - )); - } else { - logger.info(&format!("found {} possible function selectors.", selectors.len())); + logger.info("symbolic execution completed."); + logger.debug(&format!("snapshot completed in {:?}.", now.elapsed())); + + // open the tui + if !args.no_tui { + tui::handle( + snapshots.clone(), + &all_resolved_errors, + &all_resolved_events, + if args.target.len() > 64 { &shortened_target } else { args.target.as_str() }, + (compiler, &version), + ) } - logger.info(&format!("performing symbolic execution on '{shortened_target}' .")); + trace.display(); + Ok(SnapshotResult { + snapshots, + resolved_errors: all_resolved_errors, + resolved_events: all_resolved_events, + }) +} - // get a new progress bar +async fn get_snapshots( + selectors: HashMap, + resolved_selectors: HashMap>, + contract_bytecode: &str, + logger: &Logger, + trace: &mut TraceFactory, + vm_trace: u32, + evm: &VM, + args: &SnapshotArgs, +) -> Result< + (Vec, HashMap, HashMap), + Box, +> { + let mut all_resolved_errors: HashMap = HashMap::new(); + let mut all_resolved_events: HashMap = HashMap::new(); + let mut snapshots: Vec = Vec::new(); let mut snapshot_progress = ProgressBar::new_spinner(); + snapshot_progress.enable_steady_tick(Duration::from_millis(100)); snapshot_progress.set_style(logger.info_spinner()); - // perform EVM analysis - let mut snapshots: Vec = Vec::new(); for (selector, function_entry_point) in selectors { snapshot_progress.set_message(format!("executing '0x{selector}'")); @@ -288,7 +258,7 @@ pub async fn snapshot(args: SnapshotArgs) -> Result Result Result func.clone(), - None => { - trace.add_warn( - func_analysis_trace, - line!(), - "failed to resolve function signature", - ); - Vec::new() - } - }; - - let mut matched_resolved_functions = match_parameters(resolved_functions, &snapshot); - - trace.br(func_analysis_trace); - if matched_resolved_functions.is_empty() { - trace.add_warn( - func_analysis_trace, - line!(), - "no resolved signatures matched this function's parameters", - ); - } else { - let mut selected_function_index: u8 = 0; - - // sort matches by signature using score heuristic from `score_signature` - matched_resolved_functions.sort_by(|a, b| { - let a_score = score_signature(&a.signature); - let b_score = score_signature(&b.signature); - b_score.cmp(&a_score) - }); - - if matched_resolved_functions.len() > 1 { - snapshot_progress.suspend(|| { - selected_function_index = logger.option( - "warn", - "multiple possible matches found. select an option below", - matched_resolved_functions - .iter() - .map(|x| x.signature.clone()) - .collect(), - Some(0u8), - args.default, - ); - }); - } - - let selected_match = - match matched_resolved_functions.get(selected_function_index as usize) { - Some(selected_match) => selected_match, - None => continue, - }; - - snapshot.resolved_function = Some(selected_match.clone()); - - let match_trace = trace.add_info( - func_analysis_trace, - line!(), - &format!( - "{} resolved signature{} matched this function's parameters", - matched_resolved_functions.len(), - if matched_resolved_functions.len() > 1 { "s" } else { "" } - ) - .to_string(), - ); - - for resolved_function in matched_resolved_functions { - trace.add_message(match_trace, line!(), vec![resolved_function.signature]); - } - } - - snapshot_progress.finish_and_clear(); - - // resolve custom error signatures - let mut resolved_counter = 0; - let resolved_errors: HashMap> = resolve_selectors( - snapshot - .errors - .keys() - .map(|error_selector| encode_hex_reduced(*error_selector).replacen("0x", "", 1)) - .collect(), - ) - .await; - for (error_selector, _) in snapshot.errors.clone() { - let error_selector_str = encode_hex_reduced(error_selector).replacen("0x", "", 1); - let mut selected_error_index: u8 = 0; - let mut resolved_error_selectors = match resolved_errors.get(&error_selector_str) { - Some(func) => func.clone(), - None => Vec::new(), - }; - - // sort matches by signature using score heuristic from `score_signature` - resolved_error_selectors.sort_by(|a, b| { - let a_score = score_signature(&a.signature); - let b_score = score_signature(&b.signature); - b_score.cmp(&a_score) - }); - - if resolved_error_selectors.len() > 1 { - snapshot_progress.suspend(|| { - selected_error_index = logger.option( - "warn", - "multiple possible matches found. select an option below", - resolved_error_selectors.iter().map(|x| x.signature.clone()).collect(), - Some(0u8), - args.default, - ); - }); - } - - let selected_match = - match resolved_error_selectors.get(selected_error_index as usize) { - Some(selected_match) => selected_match, - None => continue, - }; - - resolved_counter += 1; - snapshot.errors.insert(error_selector, Some(selected_match.clone())); - all_resolved_errors.insert(error_selector_str, selected_match.clone()); - } - - if resolved_counter > 0 { - trace.br(func_analysis_trace); - let error_trace = trace.add_info( - func_analysis_trace, - line!(), - &format!( - "resolved {} error signatures from {} selectors.", - resolved_counter, - snapshot.errors.len() - ) - .to_string(), - ); - - for resolved_error in all_resolved_errors.values() { - trace.add_message(error_trace, line!(), vec![resolved_error.signature.clone()]); - } - } - - // resolve custom event signatures - resolved_counter = 0; - let resolved_events: HashMap> = resolve_selectors( - snapshot - .events - .keys() - .map(|event_selector| encode_hex_reduced(*event_selector).replacen("0x", "", 1)) - .collect(), + resolve_signatures( + &mut snapshot, + &mut all_resolved_errors, + &mut all_resolved_events, + &mut snapshot_progress, + trace, + &selector, + &resolved_selectors, + func_analysis_trace, + args.default, ) - .await; - for (event_selector, (_, raw_event)) in snapshot.events.clone() { - let mut selected_event_index: u8 = 0; - let event_selector_str = encode_hex_reduced(event_selector).replacen("0x", "", 1); - let mut resolved_event_selectors = match resolved_events.get(&event_selector_str) { - Some(func) => func.clone(), - None => Vec::new(), - }; - - // sort matches by signature using score heuristic from `score_signature` - resolved_event_selectors.sort_by(|a, b| { - let a_score = score_signature(&a.signature); - let b_score = score_signature(&b.signature); - b_score.cmp(&a_score) - }); - - if resolved_event_selectors.len() > 1 { - snapshot_progress.suspend(|| { - selected_event_index = logger.option( - "warn", - "multiple possible matches found. select an option below", - resolved_event_selectors.iter().map(|x| x.signature.clone()).collect(), - Some(0u8), - args.default, - ); - }); - } - - let selected_match = - match resolved_event_selectors.get(selected_event_index as usize) { - Some(selected_match) => selected_match, - None => continue, - }; - - resolved_counter += 1; - snapshot.events.insert(event_selector, (Some(selected_match.clone()), raw_event)); - all_resolved_events.insert(event_selector_str, selected_match.clone()); - } - - if resolved_counter > 0 { - let event_trace = trace.add_info( - func_analysis_trace, - line!(), - &format!( - "resolved {} event signatures from {} selectors.", - resolved_counter, - snapshot.events.len() - ), - ); - - for resolved_event in all_resolved_events.values() { - trace.add_message(event_trace, line!(), vec![resolved_event.signature.clone()]); - } - } + .await?; } - // push snapshots.push(snapshot); - // get a new progress bar snapshot_progress = ProgressBar::new_spinner(); snapshot_progress.enable_steady_tick(Duration::from_millis(100)); snapshot_progress.set_style(logger.info_spinner()); } - snapshot_progress.finish_and_clear(); - logger.info("symbolic execution completed."); - logger.debug(&format!("snapshot completed in {:?}.", now.elapsed())); - // open the tui - if !args.no_tui { - tui::handle( - snapshots.clone(), - &all_resolved_errors, - &all_resolved_events, - if args.target.len() > 64 { &shortened_target } else { args.target.as_str() }, - (compiler, &version), - ) - } + snapshot_progress.finish_and_clear(); - trace.display(); - Ok(SnapshotResult { - snapshots, - resolved_errors: all_resolved_errors, - resolved_events: all_resolved_events, - }) + Ok((snapshots, all_resolved_errors, all_resolved_events)) } diff --git a/core/src/snapshot/resolve.rs b/core/src/snapshot/resolve.rs index 8ddff5c7..67a71a7d 100644 --- a/core/src/snapshot/resolve.rs +++ b/core/src/snapshot/resolve.rs @@ -1,6 +1,18 @@ -use heimdall_common::{debug_max, ether::signatures::ResolvedFunction}; +use std::collections::HashMap; use super::structures::snapshot::Snapshot; +use heimdall_common::{ + debug_max, + ether::{ + selectors::resolve_selectors, + signatures::{score_signature, ResolvedError, ResolvedFunction, ResolvedLog}, + }, + utils::{ + io::logging::{Logger, TraceFactory}, + strings::encode_hex_reduced, + }, +}; +use indicatif::ProgressBar; /// Given a list of potential [`ResolvedFunction`]s and a [`Snapshot`], return a list of /// [`ResolvedFunction`]s (that is, resolved signatures that were found on a 4byte directory) that @@ -79,3 +91,274 @@ pub fn match_parameters( matched_functions } + +// Given a [`Snapshot`], resolve all the errors, functions and events signatures +pub async fn resolve_signatures( + snapshot: &mut Snapshot, + all_resolved_errors: &mut HashMap, + all_resolved_events: &mut HashMap, + snapshot_progress: &mut ProgressBar, + trace: &mut TraceFactory, + selector: &str, + resolved_selectors: &HashMap>, + func_analysis_trace: u32, + default: bool, +) -> Result<(), Box> { + let resolved_functions = match resolved_selectors.get(selector) { + Some(func) => func.clone(), + None => { + trace.add_warn(func_analysis_trace, line!(), "failed to resolve function signature"); + Vec::new() + } + }; + + let mut matched_resolved_functions = match_parameters(resolved_functions, snapshot); + + trace.br(func_analysis_trace); + if matched_resolved_functions.is_empty() { + trace.add_warn( + func_analysis_trace, + line!(), + "no resolved signatures matched this function's parameters", + ); + } else { + resolve_function_signatures( + &mut matched_resolved_functions, + snapshot, + snapshot_progress, + default, + &func_analysis_trace, + trace, + ) + .await?; + } + + snapshot_progress.finish_and_clear(); + + let mut resolved_counter = 0; + resolve_error_signatures( + snapshot, + all_resolved_errors, + snapshot_progress, + &mut resolved_counter, + default, + ) + .await?; + + if resolved_counter > 0 { + trace.br(func_analysis_trace); + let error_trace = trace.add_info( + func_analysis_trace, + line!(), + &format!( + "resolved {} error signatures from {} selectors.", + resolved_counter, + snapshot.errors.len() + ) + .to_string(), + ); + + for resolved_error in all_resolved_errors.values() { + trace.add_message(error_trace, line!(), vec![resolved_error.signature.clone()]); + } + } + + resolved_counter = 0; + resolve_event_signatures( + snapshot, + all_resolved_events, + snapshot_progress, + &mut resolved_counter, + default, + ) + .await?; + + if resolved_counter > 0 { + let event_trace = trace.add_info( + func_analysis_trace, + line!(), + &format!( + "resolved {} event signatures from {} selectors.", + resolved_counter, + snapshot.events.len() + ), + ); + + for resolved_event in all_resolved_events.values() { + trace.add_message(event_trace, line!(), vec![resolved_event.signature.clone()]); + } + } + + Ok(()) +} + +async fn resolve_function_signatures( + matched_resolved_functions: &mut Vec, + snapshot: &mut Snapshot, + snapshot_progress: &mut ProgressBar, + default: bool, + func_analysis_trace: &u32, + trace: &mut TraceFactory, +) -> Result<(), Box> { + let (logger, _) = Logger::new(""); + let mut selected_function_index: u8 = 0; + + // sort matches by signature using score heuristic from `score_signature` + matched_resolved_functions.sort_by(|a, b| { + let a_score = score_signature(&a.signature); + let b_score = score_signature(&b.signature); + b_score.cmp(&a_score) + }); + + if matched_resolved_functions.len() > 1 { + snapshot_progress.suspend(|| { + selected_function_index = logger.option( + "warn", + "multiple possible matches found. select an option below", + matched_resolved_functions.iter().map(|x| x.signature.clone()).collect(), + Some(0u8), + default, + ); + }); + } + + let selected_match = match matched_resolved_functions.get(selected_function_index as usize) { + Some(selected_match) => selected_match, + None => panic!(), + }; + + snapshot.resolved_function = Some(selected_match.clone()); + + let match_trace = trace.add_info( + *func_analysis_trace, + line!(), + &format!( + "{} resolved signature{} matched this function's parameters", + matched_resolved_functions.len(), + if matched_resolved_functions.len() > 1 { "s" } else { "" } + ) + .to_string(), + ); + + for resolved_function in matched_resolved_functions { + trace.add_message(match_trace, line!(), vec![resolved_function.signature.clone()]); + } + + Ok(()) +} + +async fn resolve_error_signatures( + snapshot: &mut Snapshot, + all_resolved_errors: &mut HashMap, + snapshot_progress: &mut ProgressBar, + resolved_counter: &mut i32, + default: bool, +) -> Result<(), Box> { + let (logger, _) = Logger::new(""); + + let resolved_errors: HashMap> = resolve_selectors( + snapshot + .errors + .keys() + .map(|error_selector| encode_hex_reduced(*error_selector).replacen("0x", "", 1)) + .collect(), + ) + .await; + for (error_selector, _) in snapshot.errors.clone() { + let error_selector_str = encode_hex_reduced(error_selector).replacen("0x", "", 1); + let mut selected_error_index: u8 = 0; + let mut resolved_error_selectors = match resolved_errors.get(&error_selector_str) { + Some(func) => func.clone(), + None => Vec::new(), + }; + + // sort matches by signature using score heuristic from `score_signature` + resolved_error_selectors.sort_by(|a, b| { + let a_score = score_signature(&a.signature); + let b_score = score_signature(&b.signature); + b_score.cmp(&a_score) + }); + + if resolved_error_selectors.len() > 1 { + snapshot_progress.suspend(|| { + selected_error_index = logger.option( + "warn", + "multiple possible matches found. select an option below", + resolved_error_selectors.iter().map(|x| x.signature.clone()).collect(), + Some(0u8), + default, + ); + }); + } + + let selected_match = match resolved_error_selectors.get(selected_error_index as usize) { + Some(selected_match) => selected_match, + None => continue, + }; + + *resolved_counter += 1; + + snapshot.errors.insert(error_selector, Some(selected_match.clone())); + all_resolved_errors.insert(error_selector_str, selected_match.clone()); + } + + Ok(()) +} + +async fn resolve_event_signatures( + snapshot: &mut Snapshot, + all_resolved_events: &mut HashMap, + snapshot_progress: &mut ProgressBar, + resolved_counter: &mut i32, + default: bool, +) -> Result<(), Box> { + let (logger, _) = Logger::new(""); + + let resolved_events: HashMap> = resolve_selectors( + snapshot + .events + .keys() + .map(|event_selector| encode_hex_reduced(*event_selector).replacen("0x", "", 1)) + .collect(), + ) + .await; + + for (event_selector, (_, raw_event)) in snapshot.events.clone() { + let mut selected_event_index: u8 = 0; + let event_selector_str = encode_hex_reduced(event_selector).replacen("0x", "", 1); + let mut resolved_event_selectors = match resolved_events.get(&event_selector_str) { + Some(func) => func.clone(), + None => Vec::new(), + }; + + // sort matches by signature using score heuristic from `score_signature` + resolved_event_selectors.sort_by(|a, b| { + let a_score = score_signature(&a.signature); + let b_score = score_signature(&b.signature); + b_score.cmp(&a_score) + }); + + if resolved_event_selectors.len() > 1 { + snapshot_progress.suspend(|| { + selected_event_index = logger.option( + "warn", + "multiple possible matches found. select an option below", + resolved_event_selectors.iter().map(|x| x.signature.clone()).collect(), + Some(0u8), + default, + ); + }); + } + + let selected_match = match resolved_event_selectors.get(selected_event_index as usize) { + Some(selected_match) => selected_match, + None => continue, + }; + + *resolved_counter += 1; + snapshot.events.insert(event_selector, (Some(selected_match.clone()), raw_event)); + all_resolved_events.insert(event_selector_str, selected_match.clone()); + } + + Ok(()) +} diff --git a/scripts/test b/scripts/test index b5714126..fd864a17 100644 --- a/scripts/test +++ b/scripts/test @@ -1,3 +1,6 @@ +# clear cache +heimdall cache clean || true + # if --cov is passed, then we want to run the coverage command if [ "$1" = "--cov" ]; then