From 769a33f99f8eab76e3bd3b79ba9a0207a7bd7e98 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ian=20Guimar=C3=A3es?= Date: Wed, 20 Dec 2023 15:27:34 -0300 Subject: [PATCH] chore: `heimdall-common` replacements (#241) * refactor: replace duplicated code with get_bytecode_from_target * refactor: replace duplicated logger env set up with set_logger_env * refactor: replace duplicated target shortning with get_shortned_address * refactor: remove duplicated log env set up * style: code format * chore: replace logger setup with set_logger_env function for dump mod * chore: replace logger setup with set_logger_env function for inspect mod --- core/src/cfg/mod.rs | 55 +++++-------------------------- core/src/decode/mod.rs | 16 +++------ core/src/decompile/mod.rs | 59 ++++----------------------------- core/src/disassemble/mod.rs | 46 +++----------------------- core/src/dump/mod.rs | 11 +------ core/src/heimdall.rs | 66 ------------------------------------- core/src/inspect/mod.rs | 13 ++------ 7 files changed, 27 insertions(+), 239 deletions(-) diff --git a/core/src/cfg/mod.rs b/core/src/cfg/mod.rs index d6a92451..06a10303 100644 --- a/core/src/cfg/mod.rs +++ b/core/src/cfg/mod.rs @@ -3,17 +3,16 @@ pub mod output; use derive_builder::Builder; use heimdall_common::{ debug_max, - ether::{compiler::detect_compiler, rpc::get_code, selectors::find_function_selectors}, + ether::{ + bytecode::get_bytecode_from_target, compiler::detect_compiler, + selectors::find_function_selectors, + }, }; use indicatif::ProgressBar; -use std::{fs, time::Duration}; +use std::time::Duration; use clap::{AppSettings, Parser}; -use heimdall_common::{ - constants::{ADDRESS_REGEX, BYTECODE_REGEX}, - ether::evm::core::vm::VM, - utils::io::logging::*, -}; +use heimdall_common::{ether::evm::core::vm::VM, utils::io::logging::*}; use petgraph::Graph; use crate::{ @@ -79,16 +78,7 @@ pub async fn cfg(args: CFGArgs) -> Result, Box level.as_str(), - None => "SILENT", - }, - ); - } + set_logger_env(&args.verbose); let (logger, mut trace) = Logger::new(match args.verbose.log_level() { Some(level) => level.as_str(), @@ -113,36 +103,7 @@ pub async fn cfg(args: CFGArgs) -> Result, Box { - let _contents = contents.replace('\n', ""); - if BYTECODE_REGEX.is_match(&_contents).unwrap() && _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) - } - }; - } + let contract_bytecode = get_bytecode_from_target(&args.target, &args.rpc_url).await?; // disassemble the bytecode let disassembled_bytecode = disassemble(DisassemblerArgs { diff --git a/core/src/decode/mod.rs b/core/src/decode/mod.rs index 70ff2091..dfe8d0c9 100644 --- a/core/src/decode/mod.rs +++ b/core/src/decode/mod.rs @@ -21,7 +21,10 @@ use heimdall_common::{ signatures::{score_signature, ResolveSelector, ResolvedFunction}, }, utils::{ - io::{logging::Logger, types::display}, + io::{ + logging::{set_logger_env, Logger}, + types::display, + }, strings::decode_hex, }, }; @@ -94,16 +97,7 @@ impl DecodeArgsBuilder { /// calldata, without the ABI of the target contract. #[allow(deprecated)] pub async fn decode(args: DecodeArgs) -> Result, Error> { - // set logger environment variable if not already set - if std::env::var("RUST_LOG").is_err() { - std::env::set_var( - "RUST_LOG", - match args.verbose.log_level() { - Some(level) => level.as_str(), - None => "SILENT", - }, - ); - } + set_logger_env(&args.verbose); // get a new logger and trace let (logger, mut trace) = Logger::new(match args.verbose.log_level() { diff --git a/core/src/decompile/mod.rs b/core/src/decompile/mod.rs index 22d63837..a598db33 100644 --- a/core/src/decompile/mod.rs +++ b/core/src/decompile/mod.rs @@ -4,7 +4,9 @@ pub mod out; pub mod precompile; pub mod resolve; pub mod util; -use heimdall_common::debug_max; +use heimdall_common::{ + debug_max, ether::bytecode::get_bytecode_from_target, utils::strings::get_shortned_target, +}; use crate::{ decompile::{ @@ -20,17 +22,15 @@ use derive_builder::Builder; use heimdall_common::{ ether::{ compiler::detect_compiler, - rpc::get_code, selectors::{find_function_selectors, resolve_selectors}, }, utils::strings::encode_hex_reduced, }; use indicatif::ProgressBar; -use std::{collections::HashMap, fs, time::Duration}; +use std::{collections::HashMap, time::Duration}; use clap::{AppSettings, Parser}; use heimdall_common::{ - constants::{ADDRESS_REGEX, BYTECODE_REGEX}, ether::{evm::core::vm::VM, signatures::*}, utils::io::logging::*, }; @@ -110,16 +110,7 @@ pub async fn decompile( use std::time::Instant; let now = Instant::now(); - // set logger environment variable if not already set - if std::env::var("RUST_LOG").is_err() { - std::env::set_var( - "RUST_LOG", - match args.verbose.log_level() { - Some(level) => level.as_str(), - None => "SILENT", - }, - ); - } + set_logger_env(&args.verbose); // get a new logger let (logger, mut trace) = Logger::new(match args.verbose.log_level() { @@ -136,13 +127,7 @@ pub async fn decompile( std::process::exit(1); } - // truncate target for prettier display - let mut shortened_target = args.target.clone(); - if shortened_target.len() > 66 { - shortened_target = shortened_target.chars().take(66).collect::() + - "..." + - &shortened_target.chars().skip(shortened_target.len() - 16).collect::(); - } + let shortened_target = get_shortned_target(&args.target); let decompile_call = trace.add_call( 0, line!(), @@ -152,37 +137,7 @@ pub async fn decompile( "()".to_string(), ); - // parse the various formats that are accepted as targets - // i.e, file, bytecode, contract address - let contract_bytecode: String; - if ADDRESS_REGEX.is_match(&args.target)? { - // We are decompiling a contract address, so we need to fetch the bytecode from the RPC - // provider - contract_bytecode = get_code(&args.target, &args.rpc_url).await?; - } else if BYTECODE_REGEX.is_match(&args.target)? { - debug_max!("using provided bytecode for decompilation"); - contract_bytecode = args.target.clone().replacen("0x", "", 1); - } else { - debug_max!("using provided file for decompilation."); - - // We are decompiling a file, so we need to read the bytecode from the file. - contract_bytecode = match fs::read_to_string(&args.target) { - Ok(contents) => { - 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) - } - }; - } + let contract_bytecode = get_bytecode_from_target(&args.target, &args.rpc_url).await?; // disassemble the bytecode let disassembled_bytecode = disassemble(DisassemblerArgs { diff --git a/core/src/disassemble/mod.rs b/core/src/disassemble/mod.rs index 02ffd1a3..cd6fbb76 100644 --- a/core/src/disassemble/mod.rs +++ b/core/src/disassemble/mod.rs @@ -1,12 +1,9 @@ -use std::fs; - use clap::{AppSettings, Parser}; use derive_builder::Builder; use heimdall_common::{ - constants::{ADDRESS_REGEX, BYTECODE_REGEX}, - ether::{evm::core::opcodes::Opcode, rpc::get_code}, + ether::{bytecode::get_bytecode_from_target, evm::core::opcodes::Opcode}, utils::{ - io::logging::Logger, + io::logging::{set_logger_env, Logger}, strings::{decode_hex, encode_hex}, }, }; @@ -60,16 +57,7 @@ pub async fn disassemble(args: DisassemblerArgs) -> Result level.as_str(), - None => "SILENT", - }, - ); - } + set_logger_env(&args.verbose); // get a new logger let (logger, _) = Logger::new(match args.verbose.log_level() { @@ -77,39 +65,13 @@ pub async fn disassemble(args: DisassemblerArgs) -> Result "SILENT", }); - let contract_bytecode: String; - if ADDRESS_REGEX.is_match(&args.target)? { - // We are disassembling a contract address, so we need to fetch the bytecode from the RPC - // provider. - contract_bytecode = get_code(&args.target, &args.rpc_url).await?; - } else if BYTECODE_REGEX.is_match(&args.target)? { - contract_bytecode = args.target; - } else { - // We are disassembling a file, so we need to read the bytecode from the file. - contract_bytecode = match fs::read_to_string(&args.target) { - Ok(contents) => { - let _contents = contents.replace('\n', ""); - if BYTECODE_REGEX.is_match(&_contents)? && _contents.len() % 2 == 0 { - _contents - } 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) - } - }; - } + let contract_bytecode = get_bytecode_from_target(&args.target, &args.rpc_url).await?; let mut program_counter = 0; let mut output: String = String::new(); // Iterate over the bytecode, disassembling each instruction. let byte_array = decode_hex(&contract_bytecode.replacen("0x", "", 1))?; - while program_counter < byte_array.len() { let operation = Opcode::new(byte_array[program_counter]); let mut pushed_bytes: String = String::new(); diff --git a/core/src/dump/mod.rs b/core/src/dump/mod.rs index 2d5654ad..905ddb15 100644 --- a/core/src/dump/mod.rs +++ b/core/src/dump/mod.rs @@ -94,16 +94,7 @@ impl DumpArgsBuilder { /// entry point for the dump module. Will fetch all storage slots accessed by the target contract, /// and dump them to a CSV file or the TUI. pub async fn dump(args: DumpArgs) -> Result, Box> { - // set logger environment variable if not already set - if std::env::var("RUST_LOG").is_err() { - std::env::set_var( - "RUST_LOG", - match args.verbose.log_level() { - Some(level) => level.as_str(), - None => "SILENT", - }, - ); - } + set_logger_env(&args.verbose); let (logger, _) = Logger::new(match args.verbose.log_level() { Some(level) => level.as_str(), diff --git a/core/src/heimdall.rs b/core/src/heimdall.rs index 0b9ba6a9..39cd94d4 100644 --- a/core/src/heimdall.rs +++ b/core/src/heimdall.rs @@ -104,17 +104,6 @@ fn main() { cmd.rpc_url = configuration.rpc_url; } - // set logger environment variable if not already set - if std::env::var("RUST_LOG").is_err() { - std::env::set_var( - "RUST_LOG", - match cmd.verbose.log_level() { - Some(level) => level.as_str(), - None => "SILENT", - }, - ); - } - disassemble(cmd); } @@ -124,17 +113,6 @@ fn main() { cmd.rpc_url = configuration.rpc_url; } - // set logger environment variable if not already set - if std::env::var("RUST_LOG").is_err() { - std::env::set_var( - "RUST_LOG", - match cmd.verbose.log_level() { - Some(level) => level.as_str(), - None => "SILENT", - }, - ); - } - decompile(cmd); } @@ -149,17 +127,6 @@ fn main() { cmd.openai_api_key = configuration.openai_api_key; } - // set logger environment variable if not already set - if std::env::var("RUST_LOG").is_err() { - std::env::set_var( - "RUST_LOG", - match cmd.verbose.log_level() { - Some(level) => level.as_str(), - None => "SILENT", - }, - ); - } - decode(cmd); } @@ -169,17 +136,6 @@ fn main() { cmd.rpc_url = configuration.rpc_url; } - // set logger environment variable if not already set - if std::env::var("RUST_LOG").is_err() { - std::env::set_var( - "RUST_LOG", - match cmd.verbose.log_level() { - Some(level) => level.as_str(), - None => "SILENT", - }, - ); - } - cfg(cmd); } @@ -194,17 +150,6 @@ fn main() { cmd.transpose_api_key = configuration.transpose_api_key; } - // set logger environment variable if not already set - if std::env::var("RUST_LOG").is_err() { - std::env::set_var( - "RUST_LOG", - match cmd.verbose.log_level() { - Some(level) => level.as_str(), - None => "SILENT", - }, - ); - } - dump(cmd); } @@ -214,17 +159,6 @@ fn main() { cmd.rpc_url = configuration.rpc_url; } - // set logger environment variable if not already set - if std::env::var("RUST_LOG").is_err() { - std::env::set_var( - "RUST_LOG", - match cmd.verbose.log_level() { - Some(level) => level.as_str(), - None => "SILENT", - }, - ); - } - snapshot(cmd); } Subcommands::Config(cmd) => { diff --git a/core/src/inspect/mod.rs b/core/src/inspect/mod.rs index 2bc37821..82399ecb 100644 --- a/core/src/inspect/mod.rs +++ b/core/src/inspect/mod.rs @@ -14,7 +14,7 @@ use heimdall_common::{ utils::{ env::set_env, hex::ToLowerHex, - io::logging::{Logger, TraceFactory}, + io::logging::{set_logger_env, Logger, TraceFactory}, }, }; @@ -87,16 +87,7 @@ pub struct InspectResult { /// visualization, and more. #[allow(deprecated)] pub async fn inspect(args: InspectArgs) -> Result { - // set logger environment variable if not already set - if std::env::var("RUST_LOG").is_err() { - std::env::set_var( - "RUST_LOG", - match args.verbose.log_level() { - Some(level) => level.as_str(), - None => "SILENT", - }, - ); - } + set_logger_env(&args.verbose); // set skip_resolving env variable // TODO: create a trait that can be added to a struct to set env variables