diff --git a/Cargo.lock b/Cargo.lock index b0cdb2be..3b84f9b8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2225,6 +2225,7 @@ dependencies = [ "lazy_static", "petgraph", "thiserror 1.0.69", + "tokio", "tracing", ] diff --git a/crates/cfg/Cargo.toml b/crates/cfg/Cargo.toml index d4773955..a0c96ac0 100644 --- a/crates/cfg/Cargo.toml +++ b/crates/cfg/Cargo.toml @@ -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"] } \ No newline at end of file diff --git a/crates/cfg/src/core/graph.rs b/crates/cfg/src/core/graph.rs index 702d31a4..d4b43f53 100644 --- a/crates/cfg/src/core/graph.rs +++ b/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, parent_node: Option>, jump_taken: bool, + seen_nodes: &mut HashSet, ) -> 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> { + let args = CfgArgsBuilder::new() + .target("0x6080604052348015600e575f80fd5b50600436106030575f3560e01c80632125b65b146034578063b69ef8a8146044575b5f80fd5b6044603f3660046046565b505050565b005b5f805f606084860312156057575f80fd5b833563ffffffff811681146069575f80fd5b925060208401356001600160a01b03811681146083575f80fd5b915060408401356001600160e01b0381168114609d575f80fd5b80915050925092509256".to_string()) + .build()?; + + let result = cfg(args).await?; + + println!("Contract Cfg: {:#?}", result); + + Ok(()) + } +} diff --git a/crates/cfg/src/core/mod.rs b/crates/cfg/src/core/mod.rs index 054cbf96..f78f57fd 100644 --- a/crates/cfg/src/core/mod.rs +++ b/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 { 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 = 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()); diff --git a/crates/common/src/ether/types.rs b/crates/common/src/ether/types.rs index 233cb7ce..46069963 100644 --- a/crates/common/src/ether/types.rs +++ b/crates/common/src/ether/types.rs @@ -69,10 +69,7 @@ fn extract_types_from_string(string: &str) -> Result> { let array_range = find_balanced_encapsulator(split, ('[', ']'))?; let size = split[array_range].to_string(); - array_size = match size.parse::() { - Ok(size) => Some(size), - Err(_) => None, - }; + array_size = size.parse::().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::() { - Ok(size) => Some(size), - Err(_) => None, - }); + array_size.push_back(size.parse::().ok()); string = string.replacen(&format!("[{}]", &size), "", 1); } diff --git a/crates/core/tests/test_cfg.rs b/crates/core/tests/test_cfg.rs index 50a4adad..4087a639 100644 --- a/crates/core/tests/test_cfg.rs +++ b/crates/core/tests/test_cfg.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)) } } diff --git a/crates/vm/src/core/types.rs b/crates/vm/src/core/types.rs index 2d67566c..883de195 100644 --- a/crates/vm/src/core/types.rs +++ b/crates/vm/src/core/types.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::() { - Ok(size) => Some(size), - Err(_) => None, - }); + array_size.push_back(size.parse::().ok()); string = string.replacen(&format!("[{}]", &size), "", 1); } diff --git a/crates/vm/src/core/vm.rs b/crates/vm/src/core/vm.rs index f7bc6893..fd0cd5aa 100644 --- a/crates/vm/src/core/vm.rs +++ b/crates/vm/src/core/vm.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);