feat(decompile): abi, solidity, and yul analyzers implemented

Author: Jon-Becker

crates/decompile/src/core/analyze.rs

@@ -28,6 +28,7 @@ use crate::{
 
 /// The type of analyzer to use. This will determine which heuristics are used when analyzing a
 /// [`VMTrace`] generated by symbolic execution.
+#[derive(Debug, Clone, Eq, PartialEq)]
 pub enum AnalyzerType {
     /// Analyze the trace using Solidity heuristics, which will generate high-level Solidity code
     Solidity,
@@ -61,12 +62,14 @@ impl Display for AnalyzerType {
 }
 
 /// State shared between heuristics
-#[derive(Default, Debug, Clone)]
+#[derive(Debug, Clone)]
 pub(crate) struct AnalyzerState {
     /// If we reach a JUMPI, this will hold the conditional for scope tracking
     pub jumped_conditional: Option<String>,
     /// Tracks a stack of conditionals, used for scope tracking
     pub conditional_stack: Vec<String>,
+    /// Tracks which analyzer type we are using
+    pub analyzer_type: AnalyzerType,
 }
 
 /// The analyzer, which will analyze a [`VMTrace`] generated by symbolic execution and build an
@@ -142,7 +145,11 @@ impl Analyzer {
     /// Inner analysis implementation
     fn analyze_inner(&mut self, branch: &VMTrace) -> Result<(), Error> {
         // get the analyzer state
-        let mut analyzer_state = AnalyzerState::default();
+        let mut analyzer_state = AnalyzerState {
+            jumped_conditional: None,
+            conditional_stack: Vec::new(),
+            analyzer_type: self.typ.clone(),
+        };
 
         // for each operation in the current trace branch, peform analysis with registerred
         // heuristics

crates/decompile/src/core/mod.rs

@@ -177,7 +177,7 @@ pub async fn decompile(args: DecompilerArgs) -> Result<DecompileResult, Error> {
 
             // analyze the symbolic execution trace
             let analyzed_function = analyzer.analyze()?;
-            println!("{:#?}", analyzed_function.arguments);
+            println!("{:#?}", analyzed_function.logic);
 
             Ok::<_, Error>(analyzed_function)
         })
@@ -186,6 +186,7 @@ pub async fn decompile(args: DecompilerArgs) -> Result<DecompileResult, Error> {
     debug!("analyzing symbolic execution results took {:?}", start_analysis_time.elapsed());
     info!("analyzed {} symbolic execution traces", analyzed_functions.len());
 
+    // resolve event and error selectors
     if !args.skip_resolving {
         // resolve error selectors
         let start_error_resolving_time = Instant::now();
@@ -252,6 +253,7 @@ pub async fn decompile(args: DecompilerArgs) -> Result<DecompileResult, Error> {
         all_resolved_events.extend(resolved_events);
     }
 
+    // match analyzed parameters with resolved signatures for each function
     analyzed_functions.iter_mut().for_each(|f| {
         let resolve_function_signatures =
             resolved_selectors.get(&f.selector).unwrap_or(&Vec::new()).to_owned();

crates/decompile/src/core/out/abi.rs

@@ -136,9 +136,6 @@ pub fn build_abi(
         abi.functions.insert(name, vec![function]);
     });
 
-    // pretty print abi
-    println!("{}", serde_json::to_string_pretty(&abi).expect("failed to serialize abi"));
-
     debug!("constructing abi took {:?}", start_time.elapsed());
 
     Ok(abi)

crates/decompile/src/utils/heuristics/arguments.rs

@@ -2,19 +2,23 @@ use std::collections::HashSet;
 
 use ethers::types::U256;
 
-use heimdall_common::ether::evm::core::{types::convert_bitmask, vm::State};
+use heimdall_common::ether::evm::core::{
+    types::{byte_size_to_type, convert_bitmask},
+    vm::State,
+};
 use tracing::{debug, trace};
 
 use crate::{
-    core::analyze::AnalyzerState,
+    core::analyze::{AnalyzerState, AnalyzerType},
     interfaces::{AnalyzedFunction, CalldataFrame, TypeHeuristic},
+    utils::constants::{AND_BITMASK_REGEX, AND_BITMASK_REGEX_2},
     Error,
 };
 
 pub fn argument_heuristic(
     function: &mut AnalyzedFunction,
     state: &State,
-    _: &mut AnalyzerState,
+    analyzer_state: &mut AnalyzerState,
 ) -> Result<(), Error> {
     match state.last_instruction.opcode {
         // CALLDATALOAD
@@ -76,6 +80,84 @@ pub fn argument_heuristic(
             }
         }
 
+        // RETURN
+        0xf3 => {
+            // Safely convert U256 to usize
+            let size: usize = state.last_instruction.inputs[1].try_into().unwrap_or(0);
+
+            let return_memory_operations = function.get_memory_range(
+                state.last_instruction.inputs[0],
+                state.last_instruction.inputs[1],
+            );
+            let return_memory_operations_solidified = return_memory_operations
+                .iter()
+                .map(|x| x.operations.solidify())
+                .collect::<Vec<String>>()
+                .join(", ");
+
+            // add the return statement to the function logic
+            if analyzer_state.analyzer_type == AnalyzerType::Solidity {
+                if return_memory_operations.len() <= 1 {
+                    function.logic.push(format!("return {return_memory_operations_solidified};"));
+                } else {
+                    function.logic.push(format!(
+                        "return abi.encodePacked({return_memory_operations_solidified});"
+                    ));
+                }
+            } else if analyzer_state.analyzer_type == AnalyzerType::Yul {
+                function.logic.push(format!(
+                    "return({}, {})",
+                    state.last_instruction.input_operations[0].yulify(),
+                    state.last_instruction.input_operations[1].yulify()
+                ));
+            }
+
+            // if we've already determined a return type, we don't want to do it again.
+            // we use bytes32 as a default return type
+            if function.returns != Some(String::from("bytes32")) {
+                return Ok(());
+            }
+
+            // if the any input op is ISZERO(x), this is a boolean return
+            if return_memory_operations.iter().any(|x| x.operations.opcode.name == "ISZERO") {
+                function.returns = Some(String::from("bool"));
+            }
+            // if the size of returndata is > 32, it must be a bytes memory return.
+            // it could be a struct, but we cant really determine that from the bytecode
+            else if size > 32 {
+                function.returns = Some(String::from("bytes memory"));
+            } else {
+                // attempt to find a return type within the return memory operations
+                let byte_size = match AND_BITMASK_REGEX
+                    .find(&return_memory_operations_solidified)
+                    .ok()
+                    .flatten()
+                {
+                    Some(bitmask) => {
+                        let cast = bitmask.as_str();
+
+                        cast.matches("ff").count()
+                    }
+                    None => match AND_BITMASK_REGEX_2
+                        .find(&return_memory_operations_solidified)
+                        .ok()
+                        .flatten()
+                    {
+                        Some(bitmask) => {
+                            let cast = bitmask.as_str();
+
+                            cast.matches("ff").count()
+                        }
+                        None => 32,
+                    },
+                };
+
+                // convert the cast size to a string
+                let (_, cast_types) = byte_size_to_type(byte_size);
+                function.returns = Some(cast_types[0].to_string());
+            }
+        }
+
         // integer type heuristics
         0x02 | 0x04 | 0x05 | 0x06 | 0x07 | 0x08 | 0x09 | 0x0b | 0x10 | 0x11 | 0x12 | 0x13 => {
             // check if this instruction is operating on a known argument.