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

Author: chen4903

crates/cfg/src/core/graph.rs

@@ -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 - 1)), // start from 0x00
+            encode_hex_reduced(U256::from(opcode_offset)),
             opcode_name,
             if opcode_name.contains("PUSH") {
                 encode_hex_reduced(
@@ -42,6 +46,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);
@@ -63,6 +73,7 @@ pub fn build_cfg(
                 .last_instruction
                 .opcode ==
                 JUMPDEST,
+                seen_nodes,
         )?;
     }
 

crates/cfg/src/core/mod.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());