fix(cfg): graph starts from offset 0x01 instead of 0x00, prune redundant nodes (#559)

* simple fix

* fix: when a node has already appeared, we will not add it to the graph

* move tokio to dev-dependencies

* chore: fix nightly fmt issues, make tests pass

---------

Co-authored-by: Jon-Becker <jonathan@jbecker.dev>

Author: chen4903

Cargo.lock

@@ -29,3 +29,6 @@ alloy = { version = "0.3.3", features = ["full", "rpc-types-debug", "rpc-types-t
 
 heimdall-disassembler.workspace = true
 heimdall-vm.workspace = true
+
+[dev-dependencies]
+tokio = { version = "1", features = ["full"] }

crates/cfg/Cargo.toml

@@ -6,6 +6,7 @@ use heimdall_vm::{
     ext::exec::VMTrace,
 };
 use petgraph::{matrix_graph::NodeIndex, Graph};
+use std::collections::HashSet;
 
 /// convert a symbolic execution [`VMTrace`] into a [`Graph`] of blocks, illustrating the
 /// control-flow graph found by the symbolic execution engine.
@@ -15,6 +16,7 @@ pub fn build_cfg(
     contract_cfg: &mut Graph<String, String>,
     parent_node: Option<NodeIndex<u32>>,
     jump_taken: bool,
+    seen_nodes: &mut HashSet<String>,
 ) -> Result<()> {
     let mut cfg_node: String = String::new();
     let mut parent_node = parent_node;
@@ -23,9 +25,11 @@ pub fn build_cfg(
     for operation in &vm_trace.operations {
         let opcode_name = opcode_name(operation.last_instruction.opcode);
 
+        let opcode_offset = operation.last_instruction.instruction - 1; // start from 0x00
+
         let assembly = format!(
             "{} {} {}",
-            encode_hex_reduced(U256::from(operation.last_instruction.instruction)),
+            encode_hex_reduced(U256::from(opcode_offset)),
             opcode_name,
             if opcode_name.contains("PUSH") {
                 encode_hex_reduced(
@@ -43,6 +47,12 @@ pub fn build_cfg(
         cfg_node.push_str(&format!("{}\n", &assembly));
     }
 
+    // check if this node has been seen before
+    if seen_nodes.contains(&cfg_node) {
+        return Ok(());
+    }
+    seen_nodes.insert(cfg_node.clone());
+
     // add the node to the graph
     let node_index = contract_cfg.add_node(cfg_node);
     if let Some(parent_node) = parent_node {
@@ -63,8 +73,29 @@ pub fn build_cfg(
                 .last_instruction
                 .opcode ==
                 JUMPDEST,
+            seen_nodes,
         )?;
     }
 
     Ok(())
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::{cfg, CfgArgsBuilder};
+    use tokio::test;
+
+    #[test]
+    async fn test_build_cfg() -> Result<(), Box<dyn std::error::Error>> {
+        let args = CfgArgsBuilder::new()
+            .target("0x6080604052348015600e575f80fd5b50600436106030575f3560e01c80632125b65b146034578063b69ef8a8146044575b5f80fd5b6044603f3660046046565b505050565b005b5f805f606084860312156057575f80fd5b833563ffffffff811681146069575f80fd5b925060208401356001600160a01b03811681146083575f80fd5b915060408401356001600160e01b0381168114609d575f80fd5b80915050925092509256".to_string())
+            .build()?;
+
+        let result = cfg(args).await?;
+
+        println!("Contract Cfg: {:#?}", result);
+
+        Ok(())
+    }
+}

crates/cfg/src/core/graph.rs

@@ -4,6 +4,7 @@ use alloy::primitives::Address;
 use eyre::eyre;
 use heimdall_common::{ether::compiler::detect_compiler, utils::strings::StringExt};
 use heimdall_vm::core::vm::VM;
+use std::collections::HashSet;
 
 use petgraph::{dot::Dot, Graph};
 use std::time::{Duration, Instant};
@@ -93,7 +94,8 @@ pub async fn cfg(args: CfgArgs) -> Result<CfgResult, Error> {
     let start_cfg_time = Instant::now();
     info!("building cfg for '{}' from symbolic execution trace", args.target.truncate(64));
     let mut contract_cfg = Graph::new();
-    build_cfg(&map, &mut contract_cfg, None, false)?;
+    let mut seen_nodes: HashSet<String> = HashSet::new();
+    build_cfg(&map, &mut contract_cfg, None, false, &mut seen_nodes)?;
     debug!("building cfg took {:?}", start_cfg_time.elapsed());
 
     debug!("cfg generated in {:?}", start_time.elapsed());

crates/cfg/src/core/mod.rs

@@ -69,10 +69,7 @@ fn extract_types_from_string(string: &str) -> Result<Vec<DynSolType>> {
                     let array_range = find_balanced_encapsulator(split, ('[', ']'))?;
 
                     let size = split[array_range].to_string();
-                    array_size = match size.parse::<usize>() {
-                        Ok(size) => Some(size),
-                        Err(_) => None,
-                    };
+                    array_size = size.parse::<usize>().ok();
                 }
             }
 
@@ -174,10 +171,7 @@ pub fn to_type(string: &str) -> DynSolType {
 
         let size = string[array_range].to_string();
 
-        array_size.push_back(match size.parse::<usize>() {
-            Ok(size) => Some(size),
-            Err(_) => None,
-        });
+        array_size.push_back(size.parse::<usize>().ok());
 
         string = string.replacen(&format!("[{}]", &size), "", 1);
     }

crates/common/src/ether/types.rs

@@ -57,7 +57,7 @@ mod integration_tests {
 
         let output = format!("{}", Dot::with_config(&result.graph, &[]));
 
-        for line in &[String::from("\"0x03a0 JUMPDEST \\l0x03a1 STOP \\l\"")] {
+        for line in &[String::from("\"0x039f JUMPDEST \\l0x03a0 STOP \\l\"")] {
             assert!(output.contains(line))
         }
     }

crates/core/tests/test_cfg.rs

@@ -35,10 +35,7 @@ pub fn to_type(string: &str) -> DynSolType {
 
         let size = string[array_range].to_string();
 
-        array_size.push_back(match size.parse::<usize>() {
-            Ok(size) => Some(size),
-            Err(_) => None,
-        });
+        array_size.push_back(size.parse::<usize>().ok());
 
         string = string.replacen(&format!("[{}]", &size), "", 1);
     }

crates/vm/src/core/types.rs

@@ -744,7 +744,7 @@ impl VM {
                 let result = keccak256(data);
 
                 // consume dynamic gas
-                let minimum_word_size = ((size + 31) / 32) as u128;
+                let minimum_word_size = size.div_ceil(32) as u128;
                 let gas_cost = 6 * minimum_word_size + self.memory.expansion_cost(offset, size);
                 self.consume_gas(gas_cost);
 
@@ -850,7 +850,7 @@ impl VM {
                 }
 
                 // consume dynamic gas
-                let minimum_word_size = ((size + 31) / 32) as u128;
+                let minimum_word_size = size.div_ceil(32) as u128;
                 let gas_cost = 3 * minimum_word_size + self.memory.expansion_cost(offset, size);
                 self.consume_gas(gas_cost);
 
@@ -892,7 +892,7 @@ impl VM {
                 }
 
                 // consume dynamic gas
-                let minimum_word_size = ((size + 31) / 32) as u128;
+                let minimum_word_size = size.div_ceil(32) as u128;
                 let gas_cost = 3 * minimum_word_size + self.memory.expansion_cost(offset, size);
                 self.consume_gas(gas_cost);
 
@@ -940,7 +940,7 @@ impl VM {
                 value.fill(0xff);
 
                 // consume dynamic gas
-                let minimum_word_size = ((size + 31) / 32) as u128;
+                let minimum_word_size = size.div_ceil(32) as u128;
                 let gas_cost =
                     3 * minimum_word_size + self.memory.expansion_cost(dest_offset, size);
                 self.consume_gas(gas_cost);
@@ -979,7 +979,7 @@ impl VM {
                 value.fill(0xff);
 
                 // consume dynamic gas
-                let minimum_word_size = ((size + 31) / 32) as u128;
+                let minimum_word_size = size.div_ceil(32) as u128;
                 let gas_cost =
                     3 * minimum_word_size + self.memory.expansion_cost(dest_offset, size);
                 self.consume_gas(gas_cost);
@@ -1211,7 +1211,7 @@ impl VM {
                 }
 
                 // consume dynamic gas
-                let minimum_word_size = ((size + 31) / 32) as u128;
+                let minimum_word_size = size.div_ceil(32) as u128;
                 let gas_cost = 3 * minimum_word_size + self.memory.expansion_cost(offset, size);
                 self.consume_gas(gas_cost);