fix(cfg): link fallback block from entry node in control flow graph

crates/cfg/src/core/graph.rs

@@ -6,7 +6,7 @@ use heimdall_vm::{
     ext::exec::VMTrace,
 };
 use petgraph::{matrix_graph::NodeIndex, Graph};
-use std::collections::HashSet;
+use std::collections::HashMap;
 
 /// convert a symbolic execution [`VMTrace`] into a [`Graph`] of blocks, illustrating the
 /// control-flow graph found by the symbolic execution engine.
@@ -15,7 +15,7 @@ pub(crate) fn build_cfg(
     contract_cfg: &mut Graph<String, String>,
     parent_node: Option<NodeIndex<u32>>,
     jump_taken: bool,
-    seen_nodes: &mut HashSet<String>,
+    seen_nodes: &mut HashMap<String, NodeIndex<u32>>,
 ) -> Result<()> {
     let mut cfg_node: String = String::new();
     let mut parent_node = parent_node;
@@ -46,14 +46,19 @@ pub(crate) fn build_cfg(
         cfg_node.push_str(&format!("{}\n", &assembly));
     }
 
-    // check if this node has been seen before
-    if seen_nodes.contains(&cfg_node) {
+    // if this node has been seen before, we still need to link the current parent to it,
+    // otherwise edges into already-visited blocks (e.g. a shared fallback handler) are lost.
+    // we don't recurse again, since the block's children have already been mapped.
+    if let Some(&node_index) = seen_nodes.get(&cfg_node) {
+        if let Some(parent_node) = parent_node {
+            contract_cfg.update_edge(parent_node, node_index, jump_taken.to_string());
+        }
         return Ok(());
     }
-    seen_nodes.insert(cfg_node.clone());
 
     // add the node to the graph
-    let node_index = contract_cfg.add_node(cfg_node);
+    let node_index = contract_cfg.add_node(cfg_node.clone());
+    seen_nodes.insert(cfg_node, node_index);
     if let Some(parent_node) = parent_node {
         contract_cfg.update_edge(parent_node, node_index, jump_taken.to_string());
     }

crates/cfg/src/core/mod.rs

@@ -4,7 +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 std::collections::HashMap;
 
 use petgraph::{dot::Dot, Graph};
 use std::time::{Duration, Instant};
@@ -104,7 +104,7 @@ 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();
-    let mut seen_nodes: HashSet<String> = HashSet::new();
+    let mut seen_nodes = HashMap::new();
     build_cfg(&map, &mut contract_cfg, None, false, &mut seen_nodes)?;
     debug!("building cfg took {:?}", start_cfg_time.elapsed());
 

crates/core/tests/test_cfg.rs

@@ -6,7 +6,7 @@ mod integration_tests {
     use std::path::PathBuf;
 
     use heimdall_cfg::{cfg, CfgArgs, CfgArgsBuilder, HardFork};
-    use petgraph::dot::Dot;
+    use petgraph::{dot::Dot, graph::NodeIndex, visit::EdgeRef};
     use serde_json::Value;
 
     #[tokio::test]
@@ -68,6 +68,50 @@ mod integration_tests {
         }
     }
 
+    #[tokio::test]
+    async fn test_cfg_links_fallback_block() {
+        // regression test for https://github.com/Jon-Becker/heimdall-rs/issues/627
+        // the entry node dispatches to the fallback handler when `CALLDATASIZE < 4`, but the
+        // edge into that block was previously dropped when the block had already been visited
+        // through another path (e.g. the dispatcher's fall-through).
+        let rpc_url = std::env::var("RPC_URL").unwrap_or_else(|_| {
+            println!("RPC_URL not set, skipping test");
+            std::process::exit(0);
+        });
+
+        // WETH (0xc02a...cc2) exhibits the missing fallback edge from the entry node.
+        let result = heimdall_cfg::cfg(CfgArgs {
+            target: String::from("0xc02aaa39b223fe8d0a0e5c4f27ead9083c756cc2"),
+            rpc_url,
+            default: true,
+            color_edges: false,
+            output: String::from(""),
+            name: String::from(""),
+            timeout: 10000,
+            hardfork: HardFork::Latest,
+            etherscan_api_key: String::from(""),
+        })
+        .await
+        .expect("failed to generate cfg");
+
+        // the entry node is always the first node added to the graph.
+        let entry = NodeIndex::new(0);
+        assert!(
+            result.graph[entry].contains("CALLDATASIZE"),
+            "entry node should contain the calldatasize dispatch check"
+        );
+
+        // the entry node must link to the fallback handler at 0xaf.
+        let links_to_fallback = result
+            .graph
+            .edges(entry)
+            .any(|edge| result.graph[edge.target()].starts_with("0xaf JUMPDEST"));
+        assert!(
+            links_to_fallback,
+            "entry node should link to the fallback block (0xaf) when calldatasize < 4"
+        );
+    }
+
     #[tokio::test]
     async fn test_cfg_auto_hardfork() {
         let rpc_url = std::env::var("RPC_URL").unwrap_or_else(|_| {