Merge pull request #86 from Jon-Becker/perf/decompile

✨ feat: `PUSH0` support, events & errors in .sol header

cache/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "heimdall-cache"
-version = "0.4.2"
+version = "0.4.3"
 edition = "2021"
 license = "MIT"
 readme = "README.md"

common/Cargo.toml

@@ -5,7 +5,7 @@ keywords = ["ethereum", "web3", "decompiler", "evm", "crypto"]
 license = "MIT"
 name = "heimdall-common"
 readme = "README.md"
-version = "0.4.2"
+version = "0.4.3"
 
 [dependencies]
 async-openai = "0.10.0"

common/src/ether/compiler.rs

@@ -0,0 +1,84 @@
+// returns the compiler version used to compile the contract.
+// for example: (solc, 0.8.10) or (vyper, 0.2.16)
+pub fn detect_compiler(bytecode: String) -> (String, String) {
+
+    let mut compiler = "unknown".to_string();
+    let mut version = "unknown".to_string();
+
+    // perfom prefix check for rough version matching
+    if bytecode.starts_with("363d3d373d3d3d363d73") {
+        compiler = "proxy".to_string();
+        version = "minimal".to_string();
+    }
+    else if bytecode.starts_with("366000600037611000600036600073") {
+        compiler = "proxy".to_string();
+        version = "vyper".to_string();
+    }
+    else if bytecode.starts_with("6004361015") {
+        compiler = "vyper".to_string();
+        version = "0.2.0-0.2.4,0.2.11-0.3.3".to_string();
+    }
+    else if bytecode.starts_with("341561000a") {
+        compiler = "vyper".to_string();
+        version = "0.2.5-0.2.8".to_string();
+    }
+    else if bytecode.starts_with("731bf797") {
+        compiler = "solc".to_string();
+        version = "0.4.10-0.4.24".to_string();
+    }
+    else if bytecode.starts_with("6080604052") {
+        compiler = "solc".to_string();
+        version = "0.4.22+".to_string();
+    }
+    else if bytecode.starts_with("6060604052") {
+        compiler = "solc".to_string();
+        version = "0.4.11-0.4.21".to_string();
+    }
+    else if bytecode.contains("7679706572") {
+        compiler = "vyper".to_string();
+    }
+    else if bytecode.contains("736f6c63") {
+        compiler = "solc".to_string();
+    }
+
+    // perform metadata check
+    if compiler == "solc" {
+        let compiler_version = bytecode.split("736f6c6343").collect::<Vec<&str>>();
+
+        if compiler_version.len() > 1 {
+            if let Some(encoded_version) = compiler_version[1].get(0..6) {
+                let version_array = encoded_version.chars()
+                    .collect::<Vec<char>>()
+                    .chunks(2)
+                    .map(|c| c.iter().collect::<String>())
+                    .collect::<Vec<String>>();
+
+                version = String::new();
+                for version_part in version_array {
+                    version.push_str(&format!("{}.", u8::from_str_radix(&version_part, 16).unwrap()));
+                }
+            }
+        }
+    }
+    else if compiler == "vyper" {
+        let compiler_version = bytecode.split("767970657283").collect::<Vec<&str>>();
+
+        if compiler_version.len() > 1 {
+            if let Some(encoded_version) = compiler_version[1].get(0..6) {
+                let version_array = encoded_version.chars()
+                    .collect::<Vec<char>>()
+                    .chunks(2)
+                    .map(|c| c.iter().collect::<String>())
+                    .collect::<Vec<String>>();
+
+                version = String::new();
+                for version_part in version_array {
+                    version.push_str(&format!("{}.", u8::from_str_radix(&version_part, 16).unwrap()));
+                }
+            }
+        }
+    }
+
+
+    (compiler, version.trim_end_matches('.').to_string())
+}

common/src/ether/evm/opcodes.rs

@@ -76,6 +76,7 @@ pub fn opcode(code: &str) -> Opcode {
         "59" => Opcode { name: "MSIZE", mingas: 2, inputs: 0, outputs: 1 },
         "5a" => Opcode { name: "GAS", mingas: 2, inputs: 0, outputs: 1 },
         "5b" => Opcode { name: "JUMPDEST", mingas: 1, inputs: 0, outputs: 0 },
+        "5f" => Opcode { name: "PUSH0", mingas: 3, inputs: 0, outputs: 1 },
         "60" => Opcode { name: "PUSH1", mingas: 3, inputs: 0, outputs: 1 },
         "61" => Opcode { name: "PUSH2", mingas: 3, inputs: 0, outputs: 1 },
         "62" => Opcode { name: "PUSH3", mingas: 3, inputs: 0, outputs: 1 },

common/src/ether/evm/vm.rs

@@ -1383,6 +1383,14 @@ impl VM {
                     );
                 }
 
+                // PUSH0
+                if op == 0x5f {
+                    self.stack.push(
+                        U256::from(0u8).encode_hex().as_str(),
+                        operation.clone(),
+                    );
+                }
+
                 // PUSH1 -> PUSH32
                 if (0x60..=0x7F).contains(&op) {
                     // Get the number of bytes to push

common/src/ether/mod.rs

@@ -1,4 +1,6 @@
 pub mod signatures;
+pub mod selectors;
 pub mod solidity;
+pub mod compiler;
 pub mod evm;
 pub mod yul;

common/src/ether/selectors.rs

@@ -0,0 +1,104 @@
+use std::{collections::HashMap, sync::{Arc, Mutex}, time::Duration, thread};
+
+use indicatif::ProgressBar;
+
+use crate::io::logging::Logger;
+
+use super::{evm::vm::VM, signatures::{resolve_function_signature, ResolvedFunction}};
+
+// find all function selectors in the given EVM assembly.
+pub fn find_function_selectors(assembly: String) -> Vec<String> {
+    let mut function_selectors = Vec::new();
+
+    // search through assembly for PUSH4 instructions, optimistically assuming that they are function selectors
+    let assembly: Vec<String> = assembly
+        .split('\n')
+        .map(|line| line.trim().to_string())
+        .collect();
+    for line in assembly.iter() {
+        let instruction_args: Vec<String> = line.split(' ').map(|arg| arg.to_string()).collect();
+
+        if instruction_args.len() >= 2 {
+            let instruction = instruction_args[1].clone();
+
+            if instruction == "PUSH4" {
+                let function_selector = instruction_args[2].clone();
+                function_selectors.push(function_selector);
+            }
+        }
+    }
+    function_selectors.sort();
+    function_selectors.dedup();
+    function_selectors
+}
+
+// resolve a selector's function entry point from the EVM bytecode
+pub fn resolve_entry_point(evm: &VM, selector: String) -> u128 {
+    let mut vm = evm.clone();
+
+    // execute the EVM call to find the entry point for the given selector
+    vm.calldata = selector.clone();
+    while vm.bytecode.len() >= (vm.instruction * 2 + 2) as usize {
+        let call = vm.step();
+
+        // if the opcode is an JUMPI and it matched the selector, the next jumpi is the entry point
+        if call.last_instruction.opcode == "57" {
+            let jump_condition = call.last_instruction.input_operations[1].solidify();
+            let jump_taken = call.last_instruction.inputs[1].try_into().unwrap_or(1);
+
+            if jump_condition.contains(&selector) &&
+               jump_condition.contains("msg.data[0]") &&
+               jump_condition.contains(" == ") &&
+               jump_taken == 1
+            {
+                return call.last_instruction.inputs[0].try_into().unwrap_or(0)
+            }
+        }
+
+        if vm.exitcode != 255 || !vm.returndata.is_empty() {
+            break;
+        }
+    }
+
+    0
+}
+
+// resolve a function signature from the given selectors
+pub fn resolve_function_selectors(
+    selectors: Vec<String>,
+    logger: &Logger,
+) -> HashMap<String, Vec<ResolvedFunction>> {
+    let resolved_functions: Arc<Mutex<HashMap<String, Vec<ResolvedFunction>>>> = Arc::new(Mutex::new(HashMap::new()));
+    let resolve_progress: Arc<Mutex<ProgressBar>> = Arc::new(Mutex::new(ProgressBar::new_spinner()));
+
+    let mut threads = Vec::new();
+
+    resolve_progress.lock().unwrap().enable_steady_tick(Duration::from_millis(100));
+    resolve_progress.lock().unwrap().set_style(logger.info_spinner());
+
+    for selector in selectors {
+        let function_clone = resolved_functions.clone();
+        let resolve_progress = resolve_progress.clone();
+
+        // create a new thread for each selector
+        threads.push(thread::spawn(move || {
+            if let Some(function) = resolve_function_signature(&selector) {
+                let mut _resolved_functions = function_clone.lock().unwrap();
+                let mut _resolve_progress = resolve_progress.lock().unwrap();
+                _resolve_progress.set_message(format!("resolved {} selectors...", _resolved_functions.len()));
+                _resolved_functions.insert(selector, function);
+            }
+        }));
+
+    }
+
+    // wait for all threads to finish
+    for thread in threads {
+        thread.join().unwrap();
+    }
+
+    resolve_progress.lock().unwrap().finish_and_clear();
+
+    let x = resolved_functions.lock().unwrap().clone();
+    x
+}

common/src/ether/solidity.rs

@@ -492,8 +492,10 @@ impl WrappedOpcode {
             "RETURNDATASIZE" => {
                 solidified_wrapped_opcode.push_str("ret0.length");
             }
+            "PUSH0" => {
+                solidified_wrapped_opcode.push_str("0");
+            }
             opcode => {
-
                 if opcode.starts_with("PUSH") {
                     solidified_wrapped_opcode.push_str(self.inputs[0]._solidify().as_str());
                 }
@@ -577,6 +579,15 @@ mod tests {
     use super::*;
     use ethers::types::U256;
 
+    #[test]
+    fn test_push0() {
+
+        // wraps an ADD operation with 2 raw inputs
+        let add_operation_wrapped = WrappedOpcode::new(0x5f, vec![]);
+        assert_eq!(add_operation_wrapped.solidify(), "0");
+
+    }
+
     #[test]
     fn test_solidify_add() {
 

common/src/ether/yul.rs

@@ -6,7 +6,10 @@ impl WrappedOpcode {
 
     // Returns a WrappedOpcode's yul representation.
     pub fn yulify(&self) -> String {
-        if self.opcode.name.starts_with("PUSH") {
+        if self.opcode.name == "PUSH0" {
+            "0".to_string()
+        }
+        else if self.opcode.name.starts_with("PUSH") {
             self.inputs[0]._yulify()
         }
         else {
@@ -40,6 +43,14 @@ mod tests {
     use super::*;
     use ethers::types::U256;
 
+    #[test]
+    fn test_push0() {
+
+        // wraps an ADD operation with 2 raw inputs
+        let add_operation_wrapped = WrappedOpcode::new(0x5f, vec![]);
+        assert_eq!(add_operation_wrapped.yulify(), "0");
+    }
+
     #[test]
     fn test_yulify_add() {
 

config/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "heimdall-config"
-version = "0.4.2"
+version = "0.4.3"
 edition = "2021"
 license = "MIT"
 readme = "README.md"

heimdall/Cargo.toml

@@ -5,7 +5,7 @@ keywords = ["ethereum", "web3", "decompiler", "evm", "crypto"]
 license = "MIT"
 name = "heimdall"
 readme = "README.md"
-version = "0.4.2"
+version = "0.4.3"
 
 [dependencies]
 backtrace = "0.3"

heimdall/src/cfg/mod.rs

@@ -6,6 +6,8 @@ pub mod util;
 
 use heimdall_cache::read_cache;
 use heimdall_cache::store_cache;
+use heimdall_common::ether::compiler::detect_compiler;
+use heimdall_common::ether::selectors::find_function_selectors;
 use indicatif::ProgressBar;
 use std::env;
 use std::fs;
@@ -27,9 +29,7 @@ use heimdall_common::{
 use petgraph::Graph;
 
 use crate::cfg::output::build_output;
-use crate::cfg::util::detect_compiler;
-use crate::cfg::util::find_function_selectors;
-use crate::cfg::util::map_contract;
+use crate::cfg::util::{map_contract};
 
 #[derive(Debug, Clone, Parser)]
 #[clap(about = "Generate a visual control flow graph for EVM bytecode",