From 8c18e89e34e05f282794f776b3d34416ed7cad42 Mon Sep 17 00:00:00 2001 From: Ian Guimaraes Date: Thu, 23 Nov 2023 09:46:18 -0300 Subject: [PATCH 01/23] refactor: extract logger functions --- core/src/snapshot/mod.rs | 58 +++++++++++++++++++++++++++++----------- 1 file changed, 42 insertions(+), 16 deletions(-) diff --git a/core/src/snapshot/mod.rs b/core/src/snapshot/mod.rs index cbe80bba..b538761e 100644 --- a/core/src/snapshot/mod.rs +++ b/core/src/snapshot/mod.rs @@ -101,25 +101,13 @@ 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", - }, - ); - } - - let (logger, mut trace) = Logger::new(match args.verbose.log_level() { - Some(level) => level.as_str(), - None => "SILENT", - }); + let now = Instant::now(); let mut all_resolved_events: HashMap = HashMap::new(); let mut all_resolved_errors: HashMap = HashMap::new(); + let (logger, mut trace) = get_logger_and_trace(&args.verbose); + + set_logger_env(&args.verbose); // truncate target for prettier display let mut shortened_target = args.target.clone(); @@ -563,3 +551,41 @@ pub async fn snapshot(args: SnapshotArgs) -> Result level.as_str(), + None => "SILENT", + }; + + std::env::set_var("RUST_LOG", log_level); + } +} + +fn get_logger_and_trace(verbosity: &clap_verbosity_flag::Verbosity) -> (Logger, TraceFactory) { + Logger::new(match verbosity.log_level() { + Some(level) => level.as_str(), + None => "SILENT", + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::env; + + #[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"); + } +} + From b49da544eebc397c7288b7b2d8d6eae898d28446 Mon Sep 17 00:00:00 2001 From: Ian Guimaraes Date: Thu, 23 Nov 2023 10:25:34 -0300 Subject: [PATCH 02/23] refactor: extract shortned address to function --- common/src/utils/io/logging.rs | 6 ++--- core/src/decompile/resolve.rs | 2 +- core/src/snapshot/mod.rs | 41 +++++++++++++++++++++++++--------- core/src/snapshot/resolve.rs | 2 +- 4 files changed, 36 insertions(+), 15 deletions(-) diff --git a/common/src/utils/io/logging.rs b/common/src/utils/io/logging.rs index 8d587ecf..f813aadd 100644 --- a/common/src/utils/io/logging.rs +++ b/common/src/utils/io/logging.rs @@ -108,7 +108,7 @@ impl TraceFactory { "{} {} {}", replace_last(prefix, "│ ", " ├─").bold().bright_white(), format!("[{}]", trace.instruction).bold().bright_white(), - trace.message.get(0).expect("Failed to build trace.") + trace.message.first().expect("Failed to build trace.") ); // print the children @@ -136,7 +136,7 @@ impl TraceFactory { println!( "{} emit {}", replace_last(prefix, "│ ", " ├─").bold().bright_white(), - trace.message.get(0).expect("Failed to build trace.") + trace.message.first().expect("Failed to build trace.") ); } TraceCategory::LogUnknown => { @@ -211,7 +211,7 @@ impl TraceFactory { "{} {} create → {}", replace_last(prefix, "│ ", " ├─").bold().bright_white(), format!("[{}]", trace.instruction).bold().bright_white(), - trace.message.get(0).expect("Failed to build trace.") + trace.message.first().expect("Failed to build trace.") ); // print the children diff --git a/core/src/decompile/resolve.rs b/core/src/decompile/resolve.rs index bb06c495..a00fa138 100644 --- a/core/src/decompile/resolve.rs +++ b/core/src/decompile/resolve.rs @@ -21,7 +21,7 @@ pub fn match_parameters( &function .arguments .values() - .map(|(_, types)| types.get(0).unwrap().clone()) + .map(|(_, types)| types.first().unwrap().clone()) .collect::>() .join(",") )); diff --git a/core/src/snapshot/mod.rs b/core/src/snapshot/mod.rs index b538761e..2cdd7f35 100644 --- a/core/src/snapshot/mod.rs +++ b/core/src/snapshot/mod.rs @@ -102,20 +102,14 @@ pub struct SnapshotResult { pub async fn snapshot(args: SnapshotArgs) -> Result> { use std::time::Instant; + set_logger_env(&args.verbose); + let now = Instant::now(); let mut all_resolved_events: HashMap = HashMap::new(); let mut all_resolved_errors: HashMap = HashMap::new(); let (logger, mut trace) = get_logger_and_trace(&args.verbose); + let shortened_target = get_shortned_target(&args.target); - set_logger_env(&args.verbose); - - // 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 snapshot_call = trace.add_call( 0, line!(), @@ -572,6 +566,18 @@ fn get_logger_and_trace(verbosity: &clap_verbosity_flag::Verbosity) -> (Logger, }) } +fn get_shortned_target(target: &String) -> String { + let mut shortened_target = target.clone(); + + 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 super::*; @@ -587,5 +593,20 @@ mod tests { assert_eq!(env::var("RUST_LOG").unwrap(), "SILENT"); } -} + #[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/resolve.rs b/core/src/snapshot/resolve.rs index a4c6508d..62863da9 100644 --- a/core/src/snapshot/resolve.rs +++ b/core/src/snapshot/resolve.rs @@ -22,7 +22,7 @@ pub fn match_parameters( &function .arguments .values() - .map(|(_, types)| types.get(0).unwrap().clone()) + .map(|(_, types)| types.first().unwrap().clone()) .collect::>() .join(",") )); From 8ca1b0a8b03ad831d0450aedd19f6988d7259b31 Mon Sep 17 00:00:00 2001 From: Ian Guimaraes Date: Sat, 25 Nov 2023 13:04:38 -0300 Subject: [PATCH 03/23] chore: merge with origin branch --- core/src/snapshot/mod.rs | 124 +++++++++++++++++++++++++++------------ 1 file changed, 88 insertions(+), 36 deletions(-) diff --git a/core/src/snapshot/mod.rs b/core/src/snapshot/mod.rs index 2cdd7f35..e4bcf81a 100644 --- a/core/src/snapshot/mod.rs +++ b/core/src/snapshot/mod.rs @@ -119,37 +119,8 @@ 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(), @@ -158,6 +129,7 @@ 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!(), @@ -578,9 +545,44 @@ fn get_shortned_target(target: &String) -> String { shortened_target } +async fn get_contract_bytecode( + target: &str, + rpc_url: &str, + logger: &Logger, +) -> Result> { + if ADDRESS_REGEX.is_match(target)? { + // We are snapshotting 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)? { + logger.debug_max("using provided bytecode for snapshotting."); + Ok(target.replacen("0x", "", 1)) + } else { + logger.debug_max("using provided file for snapshotting."); + + // We are snapshotting a file, 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 fancy_regex::Regex; use std::env; #[test] @@ -609,4 +611,54 @@ mod tests { assert_eq!(shortened_target.len(), 66); } + + #[tokio::test] + async fn test_get_bytecode_when_target_is_address() { + let bytecode_regex = Regex::new(r"^[0-9a-fA-F]{0,50000}$").unwrap(); + let (logger, _) = get_logger_and_trace(&clap_verbosity_flag::Verbosity::new(-1, 0)); + let bytecode = get_contract_bytecode( + "0x9f00c43700bc0000Ff91bE00841F8e04c0495000", + "https://eth.llamarpc.com", + &logger, + ) + .await + .unwrap(); + + assert!(bytecode_regex.is_match(&bytecode).unwrap()); + // Not possible to express with regex since fancy_regex + // doesn't support look-arounds + assert!(!bytecode.starts_with("0x")); + } + + #[tokio::test] + async fn test_get_bytecode_when_target_is_bytecode() { + let bytecode_regex = Regex::new(r"^[0-9a-fA-F]{0,50000}$").unwrap(); + let (logger, _) = get_logger_and_trace(&clap_verbosity_flag::Verbosity::new(-1, 0)); + let bytecode = get_contract_bytecode( + "0x00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000001", + "https://eth.llamarpc.com", + &logger, + ).await.unwrap(); + + assert!(bytecode_regex.is_match(&bytecode).unwrap()); + assert!(!bytecode.starts_with("0x")); + } + + #[tokio::test] + async fn test_get_bytecode_when_target_is_file_path() { + let (logger, _) = get_logger_and_trace(&clap_verbosity_flag::Verbosity::new(-1, 0)); + let bytecode_regex = Regex::new(r"^[0-9a-fA-F]{0,50000}$").unwrap(); + let file_path = "./mock-file.txt"; + let mock_bytecode = "0x00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000001"; + + fs::write(file_path, mock_bytecode).unwrap(); + + let bytecode = + get_contract_bytecode(file_path, "https://eth.llamarpc.com", &logger).await.unwrap(); + + assert!(bytecode_regex.is_match(&bytecode).unwrap()); + assert!(!bytecode.starts_with("0x")); + + fs::remove_file(file_path).unwrap(); + } } From 46a04bd2ded02ff8d98fdf5af926b7b06911e206 Mon Sep 17 00:00:00 2001 From: Ian Guimaraes Date: Sat, 25 Nov 2023 17:33:29 -0300 Subject: [PATCH 04/23] chore: remove duplicated function declaration --- core/src/snapshot/mod.rs | 718 +++++++++++++++++++++++---------------- 1 file changed, 433 insertions(+), 285 deletions(-) diff --git a/core/src/snapshot/mod.rs b/core/src/snapshot/mod.rs index e4bcf81a..ced90564 100644 --- a/core/src/snapshot/mod.rs +++ b/core/src/snapshot/mod.rs @@ -12,6 +12,7 @@ use std::{ }; use clap::{AppSettings, Parser}; +use clap_verbosity_flag::Verbosity; use derive_builder::Builder; use heimdall_common::{ constants::{ADDRESS_REGEX, BYTECODE_REGEX}, @@ -121,24 +122,6 @@ pub async fn snapshot(args: SnapshotArgs) -> Result Result 64 { &shortened_target } else { args.target.as_str() }, + (compiler, &version), + ) + } + + trace.display(); + Ok(SnapshotResult { + snapshots, + resolved_errors: all_resolved_errors, + resolved_events: all_resolved_events, + }) +} + +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); + } +} + +fn get_logger_and_trace(verbosity: &clap_verbosity_flag::Verbosity) -> (Logger, TraceFactory) { + Logger::new(match verbosity.log_level() { + Some(level) => level.as_str(), + None => "SILENT", + }) +} + +fn get_shortned_target(target: &String) -> String { + let mut shortened_target = target.clone(); + + if shortened_target.len() > 66 { + shortened_target = shortened_target.chars().take(66).collect::() + + "..." + + &shortened_target.chars().skip(shortened_target.len() - 16).collect::(); + } + + shortened_target +} + +async fn get_contract_bytecode( + target: &str, + rpc_url: &str, + logger: &Logger, +) -> Result> { + if ADDRESS_REGEX.is_match(target)? { + // We are snapshotting 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)? { + logger.debug_max("using provided bytecode for snapshotting."); + Ok(target.replacen("0x", "", 1)) + } else { + logger.debug_max("using provided file for snapshotting."); + + // We are snapshotting a file, 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) + } + } + } +} + +async fn get_selectors( + contract_bytecode: &str, + skip_resolving: bool, + logger: &Logger, + evm: &VM, + shortened_target: &str, + rpc_url: String, + verbose: Verbosity, + trace: &mut TraceFactory, + snapshot_call: u32, +) -> Result< + (HashMap, HashMap>), + Box, +> { + trace.add_call( + snapshot_call, + line!(), + "heimdall".to_string(), + "disassemble".to_string(), + vec![format!("{} bytes", contract_bytecode.len() / 2usize)], + "()".to_string(), + ); + // find and resolve all selectors in the bytecode - let selectors = find_function_selectors(&evm, &disassembled_bytecode); + let disassembled_bytecode = disassemble(DisassemblerArgs { + rpc_url, + verbose, + target: contract_bytecode.to_string(), + decimal_counter: false, + output: String::new(), + }) + .await?; + let selectors = find_function_selectors(evm, &disassembled_bytecode); let mut resolved_selectors = HashMap::new(); - if !args.skip_resolving { + if !skip_resolving { resolved_selectors = resolve_selectors::(selectors.keys().cloned().collect()).await; @@ -203,6 +332,21 @@ pub async fn snapshot(args: SnapshotArgs) -> Result, + resolved_selectors: HashMap>, + contract_bytecode: &str, + logger: &Logger, + trace: &mut TraceFactory, + vm_trace: u32, + evm: &VM, + args: &SnapshotArgs, + all_resolved_events: &mut HashMap, + all_resolved_errors: &mut HashMap, +) -> Result, Box> { // get a new progress bar let mut snapshot_progress = ProgressBar::new_spinner(); snapshot_progress.enable_steady_tick(Duration::from_millis(100)); @@ -230,7 +374,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); + resolve_signatures( + &mut snapshot, + &selector, + &resolved_selectors, + trace, + func_analysis_trace, // TODO: not clone + &mut snapshot_progress, + logger, + args.default, + all_resolved_events, + all_resolved_errors, + ) + .await?; + } - 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, - ); - }); - } + // push + snapshots.push(snapshot); - 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(), - ); + // 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()); + } - for resolved_function in matched_resolved_functions { - trace.add_message(match_trace, line!(), vec![resolved_function.signature]); - } - } + snapshot_progress.finish_and_clear(); - snapshot_progress.finish_and_clear(); + Ok(snapshots) +} - // 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, - ); - }); - } +async fn resolve_signatures( + snapshot: &mut Snapshot, + selector: &str, + resolved_selectors: &HashMap>, + trace: &mut TraceFactory, + func_analysis_trace: u32, + snapshot_progress: &mut ProgressBar, + logger: &Logger, + default: bool, + all_resolved_events: &mut HashMap, + all_resolved_errors: &mut HashMap, +) -> 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 selected_match = - match resolved_error_selectors.get(selected_error_index as usize) { - Some(selected_match) => selected_match, - None => continue, - }; + let mut matched_resolved_functions = match_parameters(resolved_functions, snapshot); - resolved_counter += 1; - snapshot.errors.insert(error_selector, Some(selected_match.clone())); - all_resolved_errors.insert(error_selector_str, selected_match.clone()); - } + 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, + logger, + default, + &func_analysis_trace, + trace, + ) + .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(), - ); + snapshot_progress.finish_and_clear(); - for resolved_error in all_resolved_errors.values() { - trace.add_message(error_trace, line!(), vec![resolved_error.signature.clone()]); - } - } + // resolve custom error signatures + let mut resolved_counter = 0; + resolve_error_signatures( + snapshot, + snapshot_progress, + logger, + &mut resolved_counter, + all_resolved_errors, + default, + ) + .await?; - // 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(), + 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() ) - .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, - ); - }); - } + .to_string(), + ); - let selected_match = - match resolved_event_selectors.get(selected_event_index as usize) { - Some(selected_match) => selected_match, - None => continue, - }; + for resolved_error in all_resolved_errors.values() { + trace.add_message(error_trace, line!(), vec![resolved_error.signature.clone()]); + } + } - resolved_counter += 1; - snapshot.events.insert(event_selector, (Some(selected_match.clone()), raw_event)); - all_resolved_events.insert(event_selector_str, selected_match.clone()); - } + resolved_counter = 0; + resolve_custom_events_signatures( + snapshot, + snapshot_progress, + logger, + &mut resolved_counter, + all_resolved_events, + 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() - ), - ); + 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()]); - } - } + for resolved_event in all_resolved_events.values() { + trace.add_message(event_trace, line!(), vec![resolved_event.signature.clone()]); } + } - // push - snapshots.push(snapshot); + Ok(()) +} - // 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()); +async fn resolve_function_signatures( + matched_resolved_functions: &mut Vec, + snapshot: &mut Snapshot, + snapshot_progress: &mut ProgressBar, + logger: &Logger, + default: bool, + func_analysis_trace: &u32, + trace: &mut TraceFactory, +) -> Result<(), Box> { + 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, + ); + }); } - 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), + let selected_match = match matched_resolved_functions.get(selected_function_index as usize) { + Some(selected_match) => selected_match, + None => panic!(), // TODO: can this function panic? It used to be a 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.clone()]); } - trace.display(); - Ok(SnapshotResult { - snapshots, - resolved_errors: all_resolved_errors, - resolved_events: all_resolved_events, - }) + Ok(()) } -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", +async fn resolve_error_signatures( + snapshot: &mut Snapshot, + snapshot_progress: &mut ProgressBar, + logger: &Logger, + resolved_counter: &mut i32, + all_resolved_errors: &mut HashMap, + default: bool, +) -> Result<(), Box> { + 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(), }; - std::env::set_var("RUST_LOG", log_level); - } -} - -fn get_logger_and_trace(verbosity: &clap_verbosity_flag::Verbosity) -> (Logger, TraceFactory) { - Logger::new(match verbosity.log_level() { - Some(level) => level.as_str(), - None => "SILENT", - }) -} + // 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, + ); + }); + } -fn get_shortned_target(target: &String) -> String { - let mut shortened_target = target.clone(); + let selected_match = match resolved_error_selectors.get(selected_error_index as usize) { + Some(selected_match) => selected_match, + None => continue, + }; - if shortened_target.len() > 66 { - shortened_target = shortened_target.chars().take(66).collect::() + - "..." + - &shortened_target.chars().skip(shortened_target.len() - 16).collect::(); + *resolved_counter += 1; + snapshot.errors.insert(error_selector, Some(selected_match.clone())); + all_resolved_errors.insert(error_selector_str, selected_match.clone()); } - shortened_target + Ok(()) } -async fn get_contract_bytecode( - target: &str, - rpc_url: &str, +async fn resolve_custom_events_signatures( + snapshot: &mut Snapshot, + snapshot_progress: &mut ProgressBar, logger: &Logger, -) -> Result> { - if ADDRESS_REGEX.is_match(target)? { - // We are snapshotting 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)? { - logger.debug_max("using provided bytecode for snapshotting."); - Ok(target.replacen("0x", "", 1)) - } else { - logger.debug_max("using provided file for snapshotting."); + resolved_counter: &mut i32, + all_resolved_events: &mut HashMap, + default: bool, +) -> Result<(), Box> { + 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(), + }; - // We are snapshotting a file, 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) - } + // 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(()) } #[cfg(test)] @@ -661,4 +806,7 @@ mod tests { fs::remove_file(file_path).unwrap(); } + + #[tokio::test] + async fn test_get_selectors() {} } From 63ad6e625ee318fa60999b8bea30ad08bfdce7dc Mon Sep 17 00:00:00 2001 From: Ian Guimaraes Date: Tue, 28 Nov 2023 08:55:47 -0300 Subject: [PATCH 05/23] chore: move get_shortnet_target and set_logger_env to common --- common/src/utils/io/logging.rs | 48 ++++++++++++++++++++++++++ common/src/utils/strings.rs | 36 ++++++++++++++++++++ core/src/snapshot/mod.rs | 62 +--------------------------------- 3 files changed, 85 insertions(+), 61 deletions(-) diff --git a/common/src/utils/io/logging.rs b/common/src/utils/io/logging.rs index f813aadd..f670b71f 100644 --- a/common/src/utils/io/logging.rs +++ b/common/src/utils/io/logging.rs @@ -597,11 +597,48 @@ 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); + } +} + +/// Returns a new instance of `Logger` and `TraceFactory` given a log level +/// +/// ``` +/// use heimdall_common::utils::io::logging::get_logger_and_trace; +/// +/// let verbosity = clap_verbosity_flag::Verbosity::new(-1, 0); +/// get_logger_and_trace(&verbosity); +/// ``` +pub fn get_logger_and_trace(verbosity: &clap_verbosity_flag::Verbosity) -> (Logger, TraceFactory) { + Logger::new(match verbosity.log_level() { + Some(level) => level.as_str(), + None => "SILENT", + }) +} + #[cfg(test)] mod tests { use std::time::Instant; use super::*; + use std::env; #[test] fn test_raw_trace() { @@ -879,4 +916,15 @@ mod tests { let (logger, _) = Logger::new("MAX"); logger.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 feb5595f..ed767f5f 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: &String) -> String { + let mut shortened_target = target.clone(); + + 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 ced90564..dfb2eef7 100644 --- a/core/src/snapshot/mod.rs +++ b/core/src/snapshot/mod.rs @@ -25,7 +25,7 @@ use heimdall_common::{ }, utils::{ io::logging::*, - strings::{decode_hex, encode_hex_reduced}, + strings::{decode_hex, encode_hex_reduced, get_shortned_target}, }, }; use indicatif::ProgressBar; @@ -208,38 +208,6 @@ pub async fn snapshot(args: SnapshotArgs) -> Result level.as_str(), - None => "SILENT", - }; - - std::env::set_var("RUST_LOG", log_level); - } -} - -fn get_logger_and_trace(verbosity: &clap_verbosity_flag::Verbosity) -> (Logger, TraceFactory) { - Logger::new(match verbosity.log_level() { - Some(level) => level.as_str(), - None => "SILENT", - }) -} - -fn get_shortned_target(target: &String) -> String { - let mut shortened_target = target.clone(); - - if shortened_target.len() > 66 { - shortened_target = shortened_target.chars().take(66).collect::() + - "..." + - &shortened_target.chars().skip(shortened_target.len() - 16).collect::(); - } - - shortened_target -} - async fn get_contract_bytecode( target: &str, rpc_url: &str, @@ -728,34 +696,6 @@ async fn resolve_custom_events_signatures( mod tests { use super::*; use fancy_regex::Regex; - use std::env; - - #[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"); - } - - #[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); - } #[tokio::test] async fn test_get_bytecode_when_target_is_address() { From 40be8fd122227cf280fbf5f18a215aa63e7cc7a0 Mon Sep 17 00:00:00 2001 From: Ian Guimaraes Date: Fri, 1 Dec 2023 15:21:04 -0300 Subject: [PATCH 06/23] refactor: extract function selector resolvers to resolve.rs --- core/src/snapshot/mod.rs | 272 ---------------------------------- core/src/snapshot/resolve.rs | 277 ++++++++++++++++++++++++++++++++++- 2 files changed, 276 insertions(+), 273 deletions(-) diff --git a/core/src/snapshot/mod.rs b/core/src/snapshot/mod.rs index dfb2eef7..d52551da 100644 --- a/core/src/snapshot/mod.rs +++ b/core/src/snapshot/mod.rs @@ -420,278 +420,6 @@ async fn get_snapshots( Ok(snapshots) } -async fn resolve_signatures( - snapshot: &mut Snapshot, - selector: &str, - resolved_selectors: &HashMap>, - trace: &mut TraceFactory, - func_analysis_trace: u32, - snapshot_progress: &mut ProgressBar, - logger: &Logger, - default: bool, - all_resolved_events: &mut HashMap, - all_resolved_errors: &mut HashMap, -) -> 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, - logger, - default, - &func_analysis_trace, - trace, - ) - .await?; - } - - snapshot_progress.finish_and_clear(); - - // resolve custom error signatures - let mut resolved_counter = 0; - resolve_error_signatures( - snapshot, - snapshot_progress, - logger, - &mut resolved_counter, - all_resolved_errors, - 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_custom_events_signatures( - snapshot, - snapshot_progress, - logger, - &mut resolved_counter, - all_resolved_events, - 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, - logger: &Logger, - default: bool, - func_analysis_trace: &u32, - trace: &mut TraceFactory, -) -> Result<(), Box> { - 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!(), // TODO: can this function panic? It used to be a 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.clone()]); - } - - Ok(()) -} - -async fn resolve_error_signatures( - snapshot: &mut Snapshot, - snapshot_progress: &mut ProgressBar, - logger: &Logger, - resolved_counter: &mut i32, - all_resolved_errors: &mut HashMap, - default: bool, -) -> Result<(), Box> { - 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_custom_events_signatures( - snapshot: &mut Snapshot, - snapshot_progress: &mut ProgressBar, - logger: &Logger, - resolved_counter: &mut i32, - all_resolved_events: &mut HashMap, - default: bool, -) -> Result<(), Box> { - 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(()) -} - #[cfg(test)] mod tests { use super::*; diff --git a/core/src/snapshot/resolve.rs b/core/src/snapshot/resolve.rs index 62863da9..4c2c4031 100644 --- a/core/src/snapshot/resolve.rs +++ b/core/src/snapshot/resolve.rs @@ -1,4 +1,7 @@ -use heimdall_common::{ether::signatures::ResolvedFunction, utils::io::logging::Logger}; +use std::collections::HashMap; + +use heimdall_common::{ether::{signatures::{ResolvedFunction, ResolvedLog, ResolvedError, score_signature}, selectors::resolve_selectors}, utils::{io::logging::{Logger, TraceFactory}, strings::encode_hex_reduced}}; +use indicatif::ProgressBar; use super::structures::snapshot::Snapshot; @@ -86,3 +89,275 @@ pub fn match_parameters( matched_functions } + +pub async fn resolve_signatures( + snapshot: &mut Snapshot, + selector: &str, + resolved_selectors: &HashMap>, + trace: &mut TraceFactory, + func_analysis_trace: u32, + snapshot_progress: &mut ProgressBar, + logger: &Logger, + default: bool, + all_resolved_events: &mut HashMap, + all_resolved_errors: &mut HashMap, +) -> 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, + logger, + default, + &func_analysis_trace, + trace, + ) + .await?; + } + + snapshot_progress.finish_and_clear(); + + // resolve custom error signatures + let mut resolved_counter = 0; + resolve_error_signatures( + snapshot, + snapshot_progress, + logger, + &mut resolved_counter, + all_resolved_errors, + 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_custom_events_signatures( + snapshot, + snapshot_progress, + logger, + &mut resolved_counter, + all_resolved_events, + 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, + logger: &Logger, + default: bool, + func_analysis_trace: &u32, + trace: &mut TraceFactory, +) -> Result<(), Box> { + 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!(), // TODO: can this function panic? It used to be a 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.clone()]); + } + + Ok(()) +} + +async fn resolve_error_signatures( + snapshot: &mut Snapshot, + snapshot_progress: &mut ProgressBar, + logger: &Logger, + resolved_counter: &mut i32, + all_resolved_errors: &mut HashMap, + default: bool, +) -> Result<(), Box> { + 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_custom_events_signatures( + snapshot: &mut Snapshot, + snapshot_progress: &mut ProgressBar, + logger: &Logger, + resolved_counter: &mut i32, + all_resolved_events: &mut HashMap, + default: bool, +) -> Result<(), Box> { + 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(()) +} From 0dc9274b2c50c4b38d85939b8b0ea7d1e7641c0d Mon Sep 17 00:00:00 2001 From: Ian Guimaraes Date: Fri, 1 Dec 2023 15:22:15 -0300 Subject: [PATCH 07/23] refactor: move get_contract_bytecode to a new bytecode module --- common/src/ether/bytecode.rs | 96 ++++++++++++++++++++++++++++++++++++ common/src/ether/mod.rs | 1 + core/src/snapshot/mod.rs | 87 -------------------------------- 3 files changed, 97 insertions(+), 87 deletions(-) create mode 100644 common/src/ether/bytecode.rs diff --git a/common/src/ether/bytecode.rs b/common/src/ether/bytecode.rs new file mode 100644 index 00000000..a883fc03 --- /dev/null +++ b/common/src/ether/bytecode.rs @@ -0,0 +1,96 @@ +use std::fs; +use crate::{constants::{ADDRESS_REGEX, BYTECODE_REGEX}, utils::io::logging::Logger}; +use super::rpc::get_code; + +pub async fn get_contract_bytecode( + target: &str, + rpc_url: &str, + logger: &Logger, +) -> Result> { + if ADDRESS_REGEX.is_match(target)? { + // We are snapshotting 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)? { + logger.debug_max("using provided bytecode for snapshotting."); + Ok(target.replacen("0x", "", 1)) + } else { + logger.debug_max("using provided file for snapshotting."); + + // We are snapshotting a file, 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 std::fs; + + use crate::utils::io::logging::get_logger_and_trace; + + use super::*; + use fancy_regex::Regex; + + #[tokio::test] + async fn test_get_bytecode_when_target_is_address() { + let bytecode_regex = Regex::new(r"^[0-9a-fA-F]{0,50000}$").unwrap(); + let (logger, _) = get_logger_and_trace(&clap_verbosity_flag::Verbosity::new(-1, 0)); + let bytecode = get_contract_bytecode( + "0x9f00c43700bc0000Ff91bE00841F8e04c0495000", + "https://eth.llamarpc.com", + &logger, + ) + .await + .unwrap(); + + assert!(bytecode_regex.is_match(&bytecode).unwrap()); + // Not possible to express with regex since fancy_regex + // doesn't support look-arounds + assert!(!bytecode.starts_with("0x")); + } + + #[tokio::test] + async fn test_get_bytecode_when_target_is_bytecode() { + let bytecode_regex = Regex::new(r"^[0-9a-fA-F]{0,50000}$").unwrap(); + let (logger, _) = get_logger_and_trace(&clap_verbosity_flag::Verbosity::new(-1, 0)); + let bytecode = get_contract_bytecode( + "0x00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000001", + "https://eth.llamarpc.com", + &logger, + ).await.unwrap(); + + assert!(bytecode_regex.is_match(&bytecode).unwrap()); + assert!(!bytecode.starts_with("0x")); + } + + #[tokio::test] + async fn test_get_bytecode_when_target_is_file_path() { + let (logger, _) = get_logger_and_trace(&clap_verbosity_flag::Verbosity::new(-1, 0)); + let bytecode_regex = Regex::new(r"^[0-9a-fA-F]{0,50000}$").unwrap(); + let file_path = "./mock-file.txt"; + let mock_bytecode = "0x00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000001"; + + fs::write(file_path, mock_bytecode).unwrap(); + + let bytecode = + get_contract_bytecode(file_path, "https://eth.llamarpc.com", &logger).await.unwrap(); + + assert!(bytecode_regex.is_match(&bytecode).unwrap()); + assert!(!bytecode.starts_with("0x")); + + fs::remove_file(file_path).unwrap(); + } +} diff --git a/common/src/ether/mod.rs b/common/src/ether/mod.rs index dc484173..b92f2e91 100644 --- a/common/src/ether/mod.rs +++ b/common/src/ether/mod.rs @@ -4,3 +4,4 @@ pub mod lexers; pub mod rpc; pub mod selectors; pub mod signatures; +pub mod bytecode; diff --git a/core/src/snapshot/mod.rs b/core/src/snapshot/mod.rs index d52551da..34060554 100644 --- a/core/src/snapshot/mod.rs +++ b/core/src/snapshot/mod.rs @@ -208,40 +208,6 @@ pub async fn snapshot(args: SnapshotArgs) -> Result Result> { - if ADDRESS_REGEX.is_match(target)? { - // We are snapshotting 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)? { - logger.debug_max("using provided bytecode for snapshotting."); - Ok(target.replacen("0x", "", 1)) - } else { - logger.debug_max("using provided file for snapshotting."); - - // We are snapshotting a file, 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) - } - } - } -} - async fn get_selectors( contract_bytecode: &str, skip_resolving: bool, @@ -422,59 +388,6 @@ async fn get_snapshots( #[cfg(test)] mod tests { - use super::*; - use fancy_regex::Regex; - - #[tokio::test] - async fn test_get_bytecode_when_target_is_address() { - let bytecode_regex = Regex::new(r"^[0-9a-fA-F]{0,50000}$").unwrap(); - let (logger, _) = get_logger_and_trace(&clap_verbosity_flag::Verbosity::new(-1, 0)); - let bytecode = get_contract_bytecode( - "0x9f00c43700bc0000Ff91bE00841F8e04c0495000", - "https://eth.llamarpc.com", - &logger, - ) - .await - .unwrap(); - - assert!(bytecode_regex.is_match(&bytecode).unwrap()); - // Not possible to express with regex since fancy_regex - // doesn't support look-arounds - assert!(!bytecode.starts_with("0x")); - } - - #[tokio::test] - async fn test_get_bytecode_when_target_is_bytecode() { - let bytecode_regex = Regex::new(r"^[0-9a-fA-F]{0,50000}$").unwrap(); - let (logger, _) = get_logger_and_trace(&clap_verbosity_flag::Verbosity::new(-1, 0)); - let bytecode = get_contract_bytecode( - "0x00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000001", - "https://eth.llamarpc.com", - &logger, - ).await.unwrap(); - - assert!(bytecode_regex.is_match(&bytecode).unwrap()); - assert!(!bytecode.starts_with("0x")); - } - - #[tokio::test] - async fn test_get_bytecode_when_target_is_file_path() { - let (logger, _) = get_logger_and_trace(&clap_verbosity_flag::Verbosity::new(-1, 0)); - let bytecode_regex = Regex::new(r"^[0-9a-fA-F]{0,50000}$").unwrap(); - let file_path = "./mock-file.txt"; - let mock_bytecode = "0x00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000001"; - - fs::write(file_path, mock_bytecode).unwrap(); - - let bytecode = - get_contract_bytecode(file_path, "https://eth.llamarpc.com", &logger).await.unwrap(); - - assert!(bytecode_regex.is_match(&bytecode).unwrap()); - assert!(!bytecode.starts_with("0x")); - - fs::remove_file(file_path).unwrap(); - } - #[tokio::test] async fn test_get_selectors() {} } From f9b9cfae311eb8a75b79cfdfcd21f03934e3bf70 Mon Sep 17 00:00:00 2001 From: Ian Guimaraes Date: Fri, 1 Dec 2023 15:23:12 -0300 Subject: [PATCH 08/23] refactor: move get_selectors to selectors module --- common/src/ether/selectors.rs | 60 ++++++++++++++++++++++++++++++++++- core/src/snapshot/mod.rs | 60 +++-------------------------------- 2 files changed, 63 insertions(+), 57 deletions(-) diff --git a/common/src/ether/selectors.rs b/common/src/ether/selectors.rs index 9bb537fc..dde82bd5 100644 --- a/common/src/ether/selectors.rs +++ b/common/src/ether/selectors.rs @@ -9,7 +9,65 @@ 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}}; + +// TODO: rename +// TODO: doc +pub async fn get_resolved_selectors( + contract_bytecode: &str, + snapshot_call: u32, +) -> Result< + (HashMap, HashMap>), + Box, +> { + let logger = Logger::default(); + + trace.add_call( + snapshot_call, + line!(), + "heimdall".to_string(), + "disassemble".to_string(), + vec![format!("{} bytes", contract_bytecode.len() / 2usize)], + "()".to_string(), + ); + + // find and resolve all selectors in the bytecode + let disassembled_bytecode = disassemble(DisassemblerArgs { + rpc_url, + verbose, + target: contract_bytecode.to_string(), + decimal_counter: false, + output: String::new(), + }) + .await?; + 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; + + // 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 + )); + } + + 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(&format!("performing symbolic execution on '{shortened_target}' .")); + + Ok((selectors, resolved_selectors)) +} /// find all function selectors in the given EVM assembly. pub fn find_function_selectors(evm: &VM, assembly: &str) -> HashMap { diff --git a/core/src/snapshot/mod.rs b/core/src/snapshot/mod.rs index 34060554..a93aee56 100644 --- a/core/src/snapshot/mod.rs +++ b/core/src/snapshot/mod.rs @@ -7,7 +7,6 @@ pub mod util; use std::{ collections::{HashMap, HashSet}, - fs, time::Duration, }; @@ -15,17 +14,15 @@ use clap::{AppSettings, Parser}; use clap_verbosity_flag::Verbosity; use derive_builder::Builder; use heimdall_common::{ - constants::{ADDRESS_REGEX, BYTECODE_REGEX}, ether::{ compiler::detect_compiler, evm::core::vm::VM, - rpc::get_code, selectors::{find_function_selectors, resolve_selectors}, - signatures::{score_signature, ResolvedError, ResolvedFunction, ResolvedLog}, + signatures::{ResolvedError, ResolvedFunction, ResolvedLog}, bytecode::get_contract_bytecode, }, utils::{ io::logging::*, - strings::{decode_hex, encode_hex_reduced, get_shortned_target}, + strings::{decode_hex, get_shortned_target}, }, }; use indicatif::ProgressBar; @@ -34,7 +31,7 @@ use crate::{ disassemble::{disassemble, DisassemblerArgs}, snapshot::{ analyze::snapshot_trace, - resolve::match_parameters, + resolve::resolve_signatures, structures::snapshot::{GasUsed, Snapshot}, util::tui, }, @@ -218,56 +215,7 @@ async fn get_selectors( verbose: Verbosity, trace: &mut TraceFactory, snapshot_call: u32, -) -> Result< - (HashMap, HashMap>), - Box, -> { - trace.add_call( - snapshot_call, - line!(), - "heimdall".to_string(), - "disassemble".to_string(), - vec![format!("{} bytes", contract_bytecode.len() / 2usize)], - "()".to_string(), - ); - - // find and resolve all selectors in the bytecode - let disassembled_bytecode = disassemble(DisassemblerArgs { - rpc_url, - verbose, - target: contract_bytecode.to_string(), - decimal_counter: false, - output: String::new(), - }) - .await?; - 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; - - // 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 - )); - } - - 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(&format!("performing symbolic execution on '{shortened_target}' .")); - - Ok((selectors, resolved_selectors)) -} +) async fn get_snapshots( selectors: HashMap, From 07a985ded3391929c7d8272d0b871ef3033e5266 Mon Sep 17 00:00:00 2001 From: Ian Guimaraes Date: Fri, 1 Dec 2023 16:57:45 -0300 Subject: [PATCH 09/23] refactor: make bytecode disambled external from get_resolved_selectors --- common/src/ether/selectors.rs | 30 +++++------------------- core/src/snapshot/mod.rs | 43 ++++++++++++++++++----------------- 2 files changed, 28 insertions(+), 45 deletions(-) diff --git a/common/src/ether/selectors.rs b/common/src/ether/selectors.rs index dde82bd5..e2e55d3c 100644 --- a/common/src/ether/selectors.rs +++ b/common/src/ether/selectors.rs @@ -11,35 +11,18 @@ use crate::utils::{io::logging::Logger, strings::decode_hex}; use super::{evm::core::vm::VM, signatures::{ResolveSelector, ResolvedFunction}}; -// TODO: rename -// TODO: doc +// Find all function selectors and all the data associated to this function, represented by +// [`ResolvedFunction`] pub async fn get_resolved_selectors( - contract_bytecode: &str, - snapshot_call: u32, + disassembled_bytecode: &str, + skip_resolving: &bool, + evm: &VM, + shortened_target: &str, ) -> Result< (HashMap, HashMap>), Box, > { let logger = Logger::default(); - - trace.add_call( - snapshot_call, - line!(), - "heimdall".to_string(), - "disassemble".to_string(), - vec![format!("{} bytes", contract_bytecode.len() / 2usize)], - "()".to_string(), - ); - - // find and resolve all selectors in the bytecode - let disassembled_bytecode = disassemble(DisassemblerArgs { - rpc_url, - verbose, - target: contract_bytecode.to_string(), - decimal_counter: false, - output: String::new(), - }) - .await?; let selectors = find_function_selectors(evm, &disassembled_bytecode); let mut resolved_selectors = HashMap::new(); @@ -51,7 +34,6 @@ pub async fn get_resolved_selectors( if resolved_selectors.is_empty() { logger.error(&format!( "failed to resolve any function selectors from '{shortened_target}' .", - shortened_target = shortened_target )); } diff --git a/core/src/snapshot/mod.rs b/core/src/snapshot/mod.rs index a93aee56..abf61f16 100644 --- a/core/src/snapshot/mod.rs +++ b/core/src/snapshot/mod.rs @@ -17,7 +17,7 @@ use heimdall_common::{ ether::{ compiler::detect_compiler, evm::core::vm::VM, - selectors::{find_function_selectors, resolve_selectors}, + selectors::{find_function_selectors, resolve_selectors, get_resolved_selectors}, signatures::{ResolvedError, ResolvedFunction, ResolvedLog}, bytecode::get_contract_bytecode, }, utils::{ @@ -156,16 +156,28 @@ pub async fn snapshot(args: SnapshotArgs) -> Result Result, resolved_selectors: HashMap>, @@ -256,6 +256,7 @@ async fn get_snapshots( // get a map of possible jump destinations let (map, jumpdest_count) = + // TODO: maybe it doesn't need to be cloned evm.clone().symbolic_exec_selector(&selector, function_entry_point); trace.add_debug( From 95b4ef2e82595e60e707cc8c036bfb3212c4c5dd Mon Sep 17 00:00:00 2001 From: Ian Guimaraes Date: Fri, 1 Dec 2023 17:06:08 -0300 Subject: [PATCH 10/23] refactor: remove logger from params of resolve_signature --- core/src/snapshot/mod.rs | 1 - core/src/snapshot/resolve.rs | 13 ++++++------- 2 files changed, 6 insertions(+), 8 deletions(-) diff --git a/core/src/snapshot/mod.rs b/core/src/snapshot/mod.rs index abf61f16..24f97beb 100644 --- a/core/src/snapshot/mod.rs +++ b/core/src/snapshot/mod.rs @@ -313,7 +313,6 @@ async fn get_snapshots( trace, func_analysis_trace, // TODO: not clone &mut snapshot_progress, - logger, args.default, all_resolved_events, all_resolved_errors, diff --git a/core/src/snapshot/resolve.rs b/core/src/snapshot/resolve.rs index 4c2c4031..ea907697 100644 --- a/core/src/snapshot/resolve.rs +++ b/core/src/snapshot/resolve.rs @@ -97,7 +97,6 @@ pub async fn resolve_signatures( trace: &mut TraceFactory, func_analysis_trace: u32, snapshot_progress: &mut ProgressBar, - logger: &Logger, default: bool, all_resolved_events: &mut HashMap, all_resolved_errors: &mut HashMap, @@ -124,7 +123,6 @@ pub async fn resolve_signatures( &mut matched_resolved_functions, snapshot, snapshot_progress, - logger, default, &func_analysis_trace, trace, @@ -139,7 +137,6 @@ pub async fn resolve_signatures( resolve_error_signatures( snapshot, snapshot_progress, - logger, &mut resolved_counter, all_resolved_errors, default, @@ -168,7 +165,6 @@ pub async fn resolve_signatures( resolve_custom_events_signatures( snapshot, snapshot_progress, - logger, &mut resolved_counter, all_resolved_events, default, @@ -198,11 +194,12 @@ async fn resolve_function_signatures( matched_resolved_functions: &mut Vec, snapshot: &mut Snapshot, snapshot_progress: &mut ProgressBar, - logger: &Logger, 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` @@ -252,11 +249,12 @@ async fn resolve_function_signatures( async fn resolve_error_signatures( snapshot: &mut Snapshot, snapshot_progress: &mut ProgressBar, - logger: &Logger, resolved_counter: &mut i32, all_resolved_errors: &mut HashMap, default: bool, ) -> Result<(), Box> { + let (logger, _) = Logger::new(""); + let resolved_errors: HashMap> = resolve_selectors( snapshot .errors @@ -308,11 +306,12 @@ async fn resolve_error_signatures( async fn resolve_custom_events_signatures( snapshot: &mut Snapshot, snapshot_progress: &mut ProgressBar, - logger: &Logger, resolved_counter: &mut i32, all_resolved_events: &mut HashMap, default: bool, ) -> Result<(), Box> { + let (logger, _) = Logger::new(""); + let resolved_events: HashMap> = resolve_selectors( snapshot .events From 273086dc37aa90531610122b7e95ecdd93738b7b Mon Sep 17 00:00:00 2001 From: Ian Guimaraes Date: Sun, 3 Dec 2023 23:02:45 -0300 Subject: [PATCH 11/23] refactor: rework snapshot mod variable and args declaration order --- common/src/ether/bytecode.rs | 9 ++++-- common/src/ether/mod.rs | 2 +- common/src/ether/selectors.rs | 7 +++-- core/src/snapshot/mod.rs | 52 ++++++++++++++--------------------- core/src/snapshot/resolve.rs | 33 ++++++++++++++-------- 5 files changed, 53 insertions(+), 50 deletions(-) diff --git a/common/src/ether/bytecode.rs b/common/src/ether/bytecode.rs index a883fc03..a9831ab4 100644 --- a/common/src/ether/bytecode.rs +++ b/common/src/ether/bytecode.rs @@ -1,6 +1,9 @@ -use std::fs; -use crate::{constants::{ADDRESS_REGEX, BYTECODE_REGEX}, utils::io::logging::Logger}; use super::rpc::get_code; +use crate::{ + constants::{ADDRESS_REGEX, BYTECODE_REGEX}, + utils::io::logging::Logger, +}; +use std::fs; pub async fn get_contract_bytecode( target: &str, @@ -35,7 +38,7 @@ pub async fn get_contract_bytecode( } } } - #[cfg(test)] +#[cfg(test)] mod tests { use std::fs; diff --git a/common/src/ether/mod.rs b/common/src/ether/mod.rs index b92f2e91..d6e3ac26 100644 --- a/common/src/ether/mod.rs +++ b/common/src/ether/mod.rs @@ -1,7 +1,7 @@ +pub mod bytecode; pub mod compiler; pub mod evm; pub mod lexers; pub mod rpc; pub mod selectors; pub mod signatures; -pub mod bytecode; diff --git a/common/src/ether/selectors.rs b/common/src/ether/selectors.rs index e2e55d3c..95c0e49a 100644 --- a/common/src/ether/selectors.rs +++ b/common/src/ether/selectors.rs @@ -9,7 +9,10 @@ use tokio::task; use crate::utils::{io::logging::Logger, strings::decode_hex}; -use super::{evm::core::vm::VM, signatures::{ResolveSelector, ResolvedFunction}}; +use super::{ + evm::core::vm::VM, + signatures::{ResolveSelector, ResolvedFunction}, +}; // Find all function selectors and all the data associated to this function, represented by // [`ResolvedFunction`] @@ -23,7 +26,7 @@ pub async fn get_resolved_selectors( Box, > { let logger = Logger::default(); - let selectors = find_function_selectors(evm, &disassembled_bytecode); + let selectors = find_function_selectors(evm, disassembled_bytecode); let mut resolved_selectors = HashMap::new(); if !skip_resolving { diff --git a/core/src/snapshot/mod.rs b/core/src/snapshot/mod.rs index 24f97beb..2d016504 100644 --- a/core/src/snapshot/mod.rs +++ b/core/src/snapshot/mod.rs @@ -11,14 +11,14 @@ use std::{ }; use clap::{AppSettings, Parser}; -use clap_verbosity_flag::Verbosity; use derive_builder::Builder; use heimdall_common::{ ether::{ + bytecode::get_contract_bytecode, compiler::detect_compiler, evm::core::vm::VM, - selectors::{find_function_selectors, resolve_selectors, get_resolved_selectors}, - signatures::{ResolvedError, ResolvedFunction, ResolvedLog}, bytecode::get_contract_bytecode, + selectors::get_resolved_selectors, + signatures::{ResolvedError, ResolvedFunction, ResolvedLog}, }, utils::{ io::logging::*, @@ -103,8 +103,6 @@ pub async fn snapshot(args: SnapshotArgs) -> Result = HashMap::new(); - let mut all_resolved_errors: HashMap = HashMap::new(); let (logger, mut trace) = get_logger_and_trace(&args.verbose); let shortened_target = get_shortned_target(&args.target); @@ -137,7 +135,6 @@ pub async fn snapshot(args: SnapshotArgs) -> Result Result Result Result, - all_resolved_errors: &mut HashMap, -) -> Result, Box> { - // get a new progress bar +) -> 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}'")); @@ -256,7 +254,6 @@ async fn get_snapshots( // get a map of possible jump destinations let (map, jumpdest_count) = - // TODO: maybe it doesn't need to be cloned evm.clone().symbolic_exec_selector(&selector, function_entry_point); trace.add_debug( @@ -304,26 +301,23 @@ async fn get_snapshots( func_analysis_trace, ); - // resolve signatures if !args.skip_resolving { resolve_signatures( &mut snapshot, + &mut all_resolved_errors, + &mut all_resolved_events, + &mut snapshot_progress, + trace, &selector, &resolved_selectors, - trace, - func_analysis_trace, // TODO: not clone - &mut snapshot_progress, + func_analysis_trace, args.default, - all_resolved_events, - all_resolved_errors, ) .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()); @@ -331,11 +325,5 @@ async fn get_snapshots( snapshot_progress.finish_and_clear(); - Ok(snapshots) -} - -#[cfg(test)] -mod tests { - #[tokio::test] - async fn test_get_selectors() {} + Ok((snapshots, all_resolved_errors, all_resolved_events)) } diff --git a/core/src/snapshot/resolve.rs b/core/src/snapshot/resolve.rs index ea907697..7045c94d 100644 --- a/core/src/snapshot/resolve.rs +++ b/core/src/snapshot/resolve.rs @@ -1,6 +1,15 @@ use std::collections::HashMap; -use heimdall_common::{ether::{signatures::{ResolvedFunction, ResolvedLog, ResolvedError, score_signature}, selectors::resolve_selectors}, utils::{io::logging::{Logger, TraceFactory}, strings::encode_hex_reduced}}; +use heimdall_common::{ + ether::{ + selectors::resolve_selectors, + signatures::{score_signature, ResolvedError, ResolvedFunction, ResolvedLog}, + }, + utils::{ + io::logging::{Logger, TraceFactory}, + strings::encode_hex_reduced, + }, +}; use indicatif::ProgressBar; use super::structures::snapshot::Snapshot; @@ -90,16 +99,17 @@ 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>, - trace: &mut TraceFactory, func_analysis_trace: u32, - snapshot_progress: &mut ProgressBar, default: bool, - all_resolved_events: &mut HashMap, - all_resolved_errors: &mut HashMap, ) -> Result<(), Box> { let resolved_functions = match resolved_selectors.get(selector) { Some(func) => func.clone(), @@ -132,13 +142,12 @@ pub async fn resolve_signatures( snapshot_progress.finish_and_clear(); - // resolve custom error signatures let mut resolved_counter = 0; resolve_error_signatures( snapshot, + all_resolved_errors, snapshot_progress, &mut resolved_counter, - all_resolved_errors, default, ) .await?; @@ -164,9 +173,9 @@ pub async fn resolve_signatures( resolved_counter = 0; resolve_custom_events_signatures( snapshot, + all_resolved_events, snapshot_progress, &mut resolved_counter, - all_resolved_events, default, ) .await?; @@ -199,7 +208,6 @@ async fn resolve_function_signatures( 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` @@ -223,7 +231,7 @@ async fn resolve_function_signatures( let selected_match = match matched_resolved_functions.get(selected_function_index as usize) { Some(selected_match) => selected_match, - None => panic!(), // TODO: can this function panic? It used to be a continue + None => panic!(), }; snapshot.resolved_function = Some(selected_match.clone()); @@ -248,9 +256,9 @@ async fn resolve_function_signatures( async fn resolve_error_signatures( snapshot: &mut Snapshot, + all_resolved_errors: &mut HashMap, snapshot_progress: &mut ProgressBar, resolved_counter: &mut i32, - all_resolved_errors: &mut HashMap, default: bool, ) -> Result<(), Box> { let (logger, _) = Logger::new(""); @@ -296,6 +304,7 @@ async fn resolve_error_signatures( }; *resolved_counter += 1; + snapshot.errors.insert(error_selector, Some(selected_match.clone())); all_resolved_errors.insert(error_selector_str, selected_match.clone()); } @@ -305,9 +314,9 @@ async fn resolve_error_signatures( async fn resolve_custom_events_signatures( snapshot: &mut Snapshot, + all_resolved_events: &mut HashMap, snapshot_progress: &mut ProgressBar, resolved_counter: &mut i32, - all_resolved_events: &mut HashMap, default: bool, ) -> Result<(), Box> { let (logger, _) = Logger::new(""); From f023cd6bb56f2abc662ba783a19db122b1fa6769 Mon Sep 17 00:00:00 2001 From: Ian Guimaraes Date: Tue, 5 Dec 2023 08:44:32 -0300 Subject: [PATCH 12/23] refactor: remove logger from params list and initialize internally --- common/src/ether/bytecode.rs | 18 ++++++------------ 1 file changed, 6 insertions(+), 12 deletions(-) diff --git a/common/src/ether/bytecode.rs b/common/src/ether/bytecode.rs index a9831ab4..4f0feb2a 100644 --- a/common/src/ether/bytecode.rs +++ b/common/src/ether/bytecode.rs @@ -8,8 +8,9 @@ use std::fs; pub async fn get_contract_bytecode( target: &str, rpc_url: &str, - logger: &Logger, ) -> Result> { + let (logger, _) = Logger::new(""); + if ADDRESS_REGEX.is_match(target)? { // We are snapshotting a contract address, so we need to fetch the bytecode from the RPC // provider. @@ -41,20 +42,15 @@ pub async fn get_contract_bytecode( #[cfg(test)] mod tests { use std::fs; - - use crate::utils::io::logging::get_logger_and_trace; - use super::*; use fancy_regex::Regex; #[tokio::test] async fn test_get_bytecode_when_target_is_address() { let bytecode_regex = Regex::new(r"^[0-9a-fA-F]{0,50000}$").unwrap(); - let (logger, _) = get_logger_and_trace(&clap_verbosity_flag::Verbosity::new(-1, 0)); let bytecode = get_contract_bytecode( "0x9f00c43700bc0000Ff91bE00841F8e04c0495000", "https://eth.llamarpc.com", - &logger, ) .await .unwrap(); @@ -68,12 +64,12 @@ mod tests { #[tokio::test] async fn test_get_bytecode_when_target_is_bytecode() { let bytecode_regex = Regex::new(r"^[0-9a-fA-F]{0,50000}$").unwrap(); - let (logger, _) = get_logger_and_trace(&clap_verbosity_flag::Verbosity::new(-1, 0)); let bytecode = get_contract_bytecode( "0x00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000001", "https://eth.llamarpc.com", - &logger, - ).await.unwrap(); + ) + .await + .unwrap(); assert!(bytecode_regex.is_match(&bytecode).unwrap()); assert!(!bytecode.starts_with("0x")); @@ -81,15 +77,13 @@ mod tests { #[tokio::test] async fn test_get_bytecode_when_target_is_file_path() { - let (logger, _) = get_logger_and_trace(&clap_verbosity_flag::Verbosity::new(-1, 0)); let bytecode_regex = Regex::new(r"^[0-9a-fA-F]{0,50000}$").unwrap(); let file_path = "./mock-file.txt"; let mock_bytecode = "0x00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000001"; fs::write(file_path, mock_bytecode).unwrap(); - let bytecode = - get_contract_bytecode(file_path, "https://eth.llamarpc.com", &logger).await.unwrap(); + let bytecode = get_contract_bytecode(file_path, "https://eth.llamarpc.com").await.unwrap(); assert!(bytecode_regex.is_match(&bytecode).unwrap()); assert!(!bytecode.starts_with("0x")); From 2739e82e017e92bdd3497a33550f4bd43a750d0b Mon Sep 17 00:00:00 2001 From: Ian Guimaraes Date: Tue, 5 Dec 2023 08:51:57 -0300 Subject: [PATCH 13/23] fix: remove extrar logger param from get_contract_bytecode call --- common/src/ether/bytecode.rs | 2 +- core/src/snapshot/mod.rs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/common/src/ether/bytecode.rs b/common/src/ether/bytecode.rs index 4f0feb2a..5789ed59 100644 --- a/common/src/ether/bytecode.rs +++ b/common/src/ether/bytecode.rs @@ -41,9 +41,9 @@ pub async fn get_contract_bytecode( } #[cfg(test)] mod tests { - use std::fs; use super::*; use fancy_regex::Regex; + use std::fs; #[tokio::test] async fn test_get_bytecode_when_target_is_address() { diff --git a/core/src/snapshot/mod.rs b/core/src/snapshot/mod.rs index 7bd229cd..bd9414ee 100644 --- a/core/src/snapshot/mod.rs +++ b/core/src/snapshot/mod.rs @@ -121,7 +121,7 @@ pub async fn snapshot(args: SnapshotArgs) -> Result Date: Fri, 8 Dec 2023 11:09:34 -0300 Subject: [PATCH 14/23] chore: change rpc from llama to ankr --- common/src/ether/bytecode.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/common/src/ether/bytecode.rs b/common/src/ether/bytecode.rs index 5789ed59..4f3e8a05 100644 --- a/common/src/ether/bytecode.rs +++ b/common/src/ether/bytecode.rs @@ -50,7 +50,7 @@ mod tests { let bytecode_regex = Regex::new(r"^[0-9a-fA-F]{0,50000}$").unwrap(); let bytecode = get_contract_bytecode( "0x9f00c43700bc0000Ff91bE00841F8e04c0495000", - "https://eth.llamarpc.com", + "https://rpc.ankr.com/eth", ) .await .unwrap(); @@ -66,7 +66,7 @@ mod tests { let bytecode_regex = Regex::new(r"^[0-9a-fA-F]{0,50000}$").unwrap(); let bytecode = get_contract_bytecode( "0x00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000001", - "https://eth.llamarpc.com", + "https://rpc.ankr.com/eth", ) .await .unwrap(); @@ -83,7 +83,7 @@ mod tests { fs::write(file_path, mock_bytecode).unwrap(); - let bytecode = get_contract_bytecode(file_path, "https://eth.llamarpc.com").await.unwrap(); + let bytecode = get_contract_bytecode(file_path, "https://rpc.ankr.com/eth").await.unwrap(); assert!(bytecode_regex.is_match(&bytecode).unwrap()); assert!(!bytecode.starts_with("0x")); From 4f731f30f996a8317934e781f1e9cbe203ae8410 Mon Sep 17 00:00:00 2001 From: Jon-Becker Date: Fri, 8 Dec 2023 22:55:01 -0500 Subject: [PATCH 15/23] fix(tests): remove unnecessary `0x` prefix constraint --- common/src/ether/bytecode.rs | 15 +++------------ common/src/ether/rpc.rs | 2 +- scripts/test | 3 +++ 3 files changed, 7 insertions(+), 13 deletions(-) diff --git a/common/src/ether/bytecode.rs b/common/src/ether/bytecode.rs index 4f3e8a05..3e9a1d6d 100644 --- a/common/src/ether/bytecode.rs +++ b/common/src/ether/bytecode.rs @@ -42,12 +42,10 @@ pub async fn get_contract_bytecode( #[cfg(test)] mod tests { use super::*; - use fancy_regex::Regex; use std::fs; #[tokio::test] async fn test_get_bytecode_when_target_is_address() { - let bytecode_regex = Regex::new(r"^[0-9a-fA-F]{0,50000}$").unwrap(); let bytecode = get_contract_bytecode( "0x9f00c43700bc0000Ff91bE00841F8e04c0495000", "https://rpc.ankr.com/eth", @@ -55,15 +53,11 @@ mod tests { .await .unwrap(); - assert!(bytecode_regex.is_match(&bytecode).unwrap()); - // Not possible to express with regex since fancy_regex - // doesn't support look-arounds - assert!(!bytecode.starts_with("0x")); + assert!(BYTECODE_REGEX.is_match(&bytecode).unwrap()); } #[tokio::test] async fn test_get_bytecode_when_target_is_bytecode() { - let bytecode_regex = Regex::new(r"^[0-9a-fA-F]{0,50000}$").unwrap(); let bytecode = get_contract_bytecode( "0x00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000001", "https://rpc.ankr.com/eth", @@ -71,13 +65,11 @@ mod tests { .await .unwrap(); - assert!(bytecode_regex.is_match(&bytecode).unwrap()); - assert!(!bytecode.starts_with("0x")); + assert!(BYTECODE_REGEX.is_match(&bytecode).unwrap()); } #[tokio::test] async fn test_get_bytecode_when_target_is_file_path() { - let bytecode_regex = Regex::new(r"^[0-9a-fA-F]{0,50000}$").unwrap(); let file_path = "./mock-file.txt"; let mock_bytecode = "0x00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000001"; @@ -85,8 +77,7 @@ mod tests { let bytecode = get_contract_bytecode(file_path, "https://rpc.ankr.com/eth").await.unwrap(); - assert!(bytecode_regex.is_match(&bytecode).unwrap()); - assert!(!bytecode.starts_with("0x")); + assert!(BYTECODE_REGEX.is_match(&bytecode).unwrap()); fs::remove_file(file_path).unwrap(); } diff --git a/common/src/ether/rpc.rs b/common/src/ether/rpc.rs index 74aee27e..a20606fd 100644 --- a/common/src/ether/rpc.rs +++ b/common/src/ether/rpc.rs @@ -144,7 +144,7 @@ pub async fn get_code( None, ); - 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/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 From f336a3ac325b8161fdce1575647566da3fb64fc3 Mon Sep 17 00:00:00 2001 From: Ian Guimaraes Date: Sat, 9 Dec 2023 09:59:09 -0300 Subject: [PATCH 16/23] refactor: rename resolve_custom_event_signatures to resolve_event_signatures --- core/src/snapshot/resolve.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/core/src/snapshot/resolve.rs b/core/src/snapshot/resolve.rs index f1429764..67a71a7d 100644 --- a/core/src/snapshot/resolve.rs +++ b/core/src/snapshot/resolve.rs @@ -164,7 +164,7 @@ pub async fn resolve_signatures( } resolved_counter = 0; - resolve_custom_events_signatures( + resolve_event_signatures( snapshot, all_resolved_events, snapshot_progress, @@ -305,7 +305,7 @@ async fn resolve_error_signatures( Ok(()) } -async fn resolve_custom_events_signatures( +async fn resolve_event_signatures( snapshot: &mut Snapshot, all_resolved_events: &mut HashMap, snapshot_progress: &mut ProgressBar, From 60857e36147f4aeb6c58e946b9431de6c32b4749 Mon Sep 17 00:00:00 2001 From: Ian Guimaraes Date: Sat, 9 Dec 2023 10:06:28 -0300 Subject: [PATCH 17/23] refactor: rename get_contract_bytecode to get_bytecode_from_target --- common/src/ether/bytecode.rs | 9 +++++---- core/src/snapshot/mod.rs | 4 ++-- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/common/src/ether/bytecode.rs b/common/src/ether/bytecode.rs index 3e9a1d6d..d55878fc 100644 --- a/common/src/ether/bytecode.rs +++ b/common/src/ether/bytecode.rs @@ -5,7 +5,7 @@ use crate::{ }; use std::fs; -pub async fn get_contract_bytecode( +pub async fn get_bytecode_from_target( target: &str, rpc_url: &str, ) -> Result> { @@ -46,7 +46,7 @@ mod tests { #[tokio::test] async fn test_get_bytecode_when_target_is_address() { - let bytecode = get_contract_bytecode( + let bytecode = get_bytecode_from_target( "0x9f00c43700bc0000Ff91bE00841F8e04c0495000", "https://rpc.ankr.com/eth", ) @@ -58,7 +58,7 @@ mod tests { #[tokio::test] async fn test_get_bytecode_when_target_is_bytecode() { - let bytecode = get_contract_bytecode( + let bytecode = get_bytecode_from_target( "0x00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000001", "https://rpc.ankr.com/eth", ) @@ -75,10 +75,11 @@ mod tests { fs::write(file_path, mock_bytecode).unwrap(); - let bytecode = get_contract_bytecode(file_path, "https://rpc.ankr.com/eth").await.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/core/src/snapshot/mod.rs b/core/src/snapshot/mod.rs index bd9414ee..57f5f1bb 100644 --- a/core/src/snapshot/mod.rs +++ b/core/src/snapshot/mod.rs @@ -15,7 +15,7 @@ use clap::{AppSettings, Parser}; use derive_builder::Builder; use heimdall_common::{ ether::{ - bytecode::get_contract_bytecode, + bytecode::get_bytecode_from_target, compiler::detect_compiler, evm::core::vm::VM, selectors::get_resolved_selectors, @@ -121,7 +121,7 @@ pub async fn snapshot(args: SnapshotArgs) -> Result Date: Sat, 9 Dec 2023 10:17:29 -0300 Subject: [PATCH 18/23] fix: update comments from get_bytecode_from_target to be module agnostic --- common/src/ether/bytecode.rs | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/common/src/ether/bytecode.rs b/common/src/ether/bytecode.rs index d55878fc..08ce0143 100644 --- a/common/src/ether/bytecode.rs +++ b/common/src/ether/bytecode.rs @@ -12,16 +12,17 @@ pub async fn get_bytecode_from_target( let (logger, _) = Logger::new(""); if ADDRESS_REGEX.is_match(target)? { - // We are snapshotting a contract address, so we need to fetch the bytecode from the RPC - // provider. + // 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)? { logger.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 { logger.debug_max("using provided file for snapshotting."); - // We are snapshotting a file, so we need to read the bytecode from the file. + // 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', ""); From ea577dfcb4551a58d4b714f2b828540989f8fa15 Mon Sep 17 00:00:00 2001 From: Ian Guimaraes Date: Sat, 9 Dec 2023 10:22:01 -0300 Subject: [PATCH 19/23] refactor: switch from logger::debug_max to debug_max --- common/src/ether/bytecode.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/common/src/ether/bytecode.rs b/common/src/ether/bytecode.rs index 08ce0143..c932fb87 100644 --- a/common/src/ether/bytecode.rs +++ b/common/src/ether/bytecode.rs @@ -1,7 +1,7 @@ use super::rpc::get_code; use crate::{ constants::{ADDRESS_REGEX, BYTECODE_REGEX}, - utils::io::logging::Logger, + utils::io::logging::Logger, debug_max, }; use std::fs; @@ -15,12 +15,12 @@ pub async fn get_bytecode_from_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)? { - logger.debug_max("using provided bytecode for snapshotting."); + 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 { - logger.debug_max("using provided file for snapshotting."); + 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) { From e75641657d38aa38ffeb59c5ee92b84e304853b7 Mon Sep 17 00:00:00 2001 From: Ian Guimaraes Date: Sat, 9 Dec 2023 10:30:39 -0300 Subject: [PATCH 20/23] fix: remove out of context log --- common/src/ether/selectors.rs | 2 -- 1 file changed, 2 deletions(-) diff --git a/common/src/ether/selectors.rs b/common/src/ether/selectors.rs index 176a170f..425d1a02 100644 --- a/common/src/ether/selectors.rs +++ b/common/src/ether/selectors.rs @@ -49,8 +49,6 @@ pub async fn get_resolved_selectors( debug_max!(&format!("found {} possible function selectors.", selectors.len())); } - debug_max!(&format!("performing symbolic execution on '{shortened_target}' .")); - Ok((selectors, resolved_selectors)) } From d4b4812f5ba36be459d8114057b1531dd5d097b0 Mon Sep 17 00:00:00 2001 From: Ian Guimaraes Date: Sat, 9 Dec 2023 10:43:59 -0300 Subject: [PATCH 21/23] refactor: remove get_logger_and_trace function --- common/src/utils/io/logging.rs | 15 --------------- core/src/snapshot/mod.rs | 6 ++++-- 2 files changed, 4 insertions(+), 17 deletions(-) diff --git a/common/src/utils/io/logging.rs b/common/src/utils/io/logging.rs index dc7c61ce..45a1754a 100644 --- a/common/src/utils/io/logging.rs +++ b/common/src/utils/io/logging.rs @@ -625,21 +625,6 @@ pub fn set_logger_env(verbosity: &clap_verbosity_flag::Verbosity) { } } -/// Returns a new instance of `Logger` and `TraceFactory` given a log level -/// -/// ``` -/// use heimdall_common::utils::io::logging::get_logger_and_trace; -/// -/// let verbosity = clap_verbosity_flag::Verbosity::new(-1, 0); -/// get_logger_and_trace(&verbosity); -/// ``` -pub fn get_logger_and_trace(verbosity: &clap_verbosity_flag::Verbosity) -> (Logger, TraceFactory) { - Logger::new(match verbosity.log_level() { - Some(level) => level.as_str(), - None => "SILENT", - }) -} - #[cfg(test)] mod tests { use std::time::Instant; diff --git a/core/src/snapshot/mod.rs b/core/src/snapshot/mod.rs index 57f5f1bb..577e61cf 100644 --- a/core/src/snapshot/mod.rs +++ b/core/src/snapshot/mod.rs @@ -109,9 +109,11 @@ pub async fn snapshot(args: SnapshotArgs) -> Result level.as_str(), + None => "SILENT", + }); let shortened_target = get_shortned_target(&args.target); - let snapshot_call = trace.add_call( 0, line!(), From 0633c2e37269c6aa2c15b6d28a1ab546c2f99bbf Mon Sep 17 00:00:00 2001 From: Ian Guimaraes Date: Sat, 9 Dec 2023 10:51:34 -0300 Subject: [PATCH 22/23] style: code format --- common/src/ether/bytecode.rs | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/common/src/ether/bytecode.rs b/common/src/ether/bytecode.rs index c932fb87..e577d018 100644 --- a/common/src/ether/bytecode.rs +++ b/common/src/ether/bytecode.rs @@ -1,7 +1,8 @@ use super::rpc::get_code; use crate::{ constants::{ADDRESS_REGEX, BYTECODE_REGEX}, - utils::io::logging::Logger, debug_max, + debug_max, + utils::io::logging::Logger, }; use std::fs; @@ -76,11 +77,11 @@ mod tests { fs::write(file_path, mock_bytecode).unwrap(); - let bytecode = get_bytecode_from_target(file_path, "https://rpc.ankr.com/eth").await.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(); } } - From ecb59fe7fd417a08f4dc5cbd6b46d09b24d83f1d Mon Sep 17 00:00:00 2001 From: Ian Guimaraes Date: Sat, 9 Dec 2023 12:15:27 -0300 Subject: [PATCH 23/23] refactor: change target param from string reference to string slice --- common/src/ether/selectors.rs | 8 -------- common/src/utils/strings.rs | 4 ++-- core/src/snapshot/mod.rs | 9 ++------- 3 files changed, 4 insertions(+), 17 deletions(-) diff --git a/common/src/ether/selectors.rs b/common/src/ether/selectors.rs index 425d1a02..d171a724 100644 --- a/common/src/ether/selectors.rs +++ b/common/src/ether/selectors.rs @@ -21,7 +21,6 @@ pub async fn get_resolved_selectors( disassembled_bytecode: &str, skip_resolving: &bool, evm: &VM, - shortened_target: &str, ) -> Result< (HashMap, HashMap>), Box, @@ -33,13 +32,6 @@ pub async fn get_resolved_selectors( resolved_selectors = resolve_selectors::(selectors.keys().cloned().collect()).await; - // if resolved selectors are empty, we can't perform symbolic execution - if resolved_selectors.is_empty() { - debug_max!(&format!( - "failed to resolve any function selectors from '{shortened_target}' .", - )); - } - debug_max!(&format!( "resolved {} possible functions from {} detected selectors.", resolved_selectors.len(), diff --git a/common/src/utils/strings.rs b/common/src/utils/strings.rs index e12112dc..16423e83 100644 --- a/common/src/utils/strings.rs +++ b/common/src/utils/strings.rs @@ -360,8 +360,8 @@ pub fn classify_token(token: &str) -> TokenType { /// let long_target = "0".repeat(80); /// let shortened_target = get_shortned_target(&long_target); /// ``` -pub fn get_shortned_target(target: &String) -> String { - let mut shortened_target = target.clone(); +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::() + diff --git a/core/src/snapshot/mod.rs b/core/src/snapshot/mod.rs index 577e61cf..14b838a9 100644 --- a/core/src/snapshot/mod.rs +++ b/core/src/snapshot/mod.rs @@ -180,13 +180,8 @@ pub async fn snapshot(args: SnapshotArgs) -> Result