refactor(snapshot): code split

common/src/ether/bytecode.rs

@@ -0,0 +1,84 @@
+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,
+    rpc_url: &str,
+) -> Result<String, Box<dyn std::error::Error>> {
+    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.
+        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 std::fs;
+
+    #[tokio::test]
+    async fn test_get_bytecode_when_target_is_address() {
+        let bytecode = get_contract_bytecode(
+            "0x9f00c43700bc0000Ff91bE00841F8e04c0495000",
+            "https://rpc.ankr.com/eth",
+        )
+        .await
+        .unwrap();
+
+        assert!(BYTECODE_REGEX.is_match(&bytecode).unwrap());
+    }
+
+    #[tokio::test]
+    async fn test_get_bytecode_when_target_is_bytecode() {
+        let bytecode = get_contract_bytecode(
+            "0x00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000001",
+            "https://rpc.ankr.com/eth",
+        )
+        .await
+        .unwrap();
+
+        assert!(BYTECODE_REGEX.is_match(&bytecode).unwrap());
+    }
+
+    #[tokio::test]
+    async fn test_get_bytecode_when_target_is_file_path() {
+        let file_path = "./mock-file.txt";
+        let mock_bytecode = "0x00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000001";
+
+        fs::write(file_path, mock_bytecode).unwrap();
+
+        let bytecode = get_contract_bytecode(file_path, "https://rpc.ankr.com/eth").await.unwrap();
+
+        assert!(BYTECODE_REGEX.is_match(&bytecode).unwrap());
+
+        fs::remove_file(file_path).unwrap();
+    }
+}

common/src/ether/mod.rs

@@ -1,3 +1,4 @@
+pub mod bytecode;
 pub mod compiler;
 pub mod evm;
 pub mod lexers;

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"))

common/src/ether/selectors.rs

@@ -9,9 +9,51 @@ use tokio::task;
 
 use crate::utils::{io::logging::Logger, strings::decode_hex};
 
-use super::{evm::core::vm::VM, signatures::ResolveSelector};
+use super::{
+    evm::core::vm::VM,
+    signatures::{ResolveSelector, ResolvedFunction},
+};
 use crate::debug_max;
 
+// Find all function selectors and all the data associated to this function, represented by
+// [`ResolvedFunction`]
+pub async fn get_resolved_selectors(
+    disassembled_bytecode: &str,
+    skip_resolving: &bool,
+    evm: &VM,
+    shortened_target: &str,
+) -> Result<
+    (HashMap<String, u128>, HashMap<String, Vec<ResolvedFunction>>),
+    Box<dyn std::error::Error>,
+> {
+    let selectors = find_function_selectors(evm, disassembled_bytecode);
+
+    let mut resolved_selectors = HashMap::new();
+    if !skip_resolving {
+        resolved_selectors =
+            resolve_selectors::<ResolvedFunction>(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(),
+            selectors.len()
+        ));
+    } else {
+        debug_max!(&format!("found {} possible function selectors.", selectors.len()));
+    }
+
+    debug_max!(&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<String, u128> {
     let mut function_selectors = HashMap::new();

common/src/utils/io/logging.rs

@@ -604,11 +604,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() {
@@ -887,4 +924,15 @@ mod tests {
         let (_logger, _) = Logger::new("MAX");
         debug_max!("log");
     }
+
+    #[test]
+    fn test_set_logger_env_default() {
+        env::remove_var("RUST_LOG");
+
+        let verbosity = clap_verbosity_flag::Verbosity::new(-1, 0);
+
+        set_logger_env(&verbosity);
+
+        assert_eq!(env::var("RUST_LOG").unwrap(), "SILENT");
+    }
 }

common/src/utils/strings.rs

@@ -352,6 +352,26 @@
     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::<String>() +
+            "..." +
+            &shortened_target.chars().skip(shortened_target.len() - 16).collect::<String>();
+    }
+
+    shortened_target
+}
+
 #[cfg(test)]
 mod tests {
     use ethers::types::{I256, U256};
@@ -691,4 +711,20 @@
             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);
+    }
 }