Skip to content

Commit

Permalink
🔧 fix: support for all EVM contracts
Browse files Browse the repository at this point in the history
  • Loading branch information
@Jon-Becker
Jon-Becker committed Feb 6, 2023
1 parent b928c44 commit b0c41c6
Show file tree
Hide file tree
Showing 2 changed files with 25 additions and 100 deletions.
70 changes: 22 additions & 48 deletions heimdall/src/cfg/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,8 +30,7 @@ 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_selector;
use crate::cfg::util::resolve_entry_point;
use crate::cfg::util::map_contract;

#[derive(Debug, Clone, Parser)]
#[clap(about = "Generate a visual control flow graph for EVM bytecode",
Expand Down Expand Up @@ -238,12 +237,12 @@ pub fn cfg(args: CFGArgs) {
}

// add the creation to the trace
let vm_trace = trace.add_creation(cfg_call, line!(), "contract".to_string(), shortened_target, (contract_bytecode.len()/2usize).try_into().unwrap());
let vm_trace = trace.add_creation(cfg_call, line!(), "contract".to_string(), shortened_target.clone(), (contract_bytecode.len()/2usize).try_into().unwrap());

// find all selectors in the bytecode
let selectors = find_function_selectors(disassembled_bytecode);
logger.info(&format!("found {} possible function selectors.", selectors.len()));
logger.info(&format!("performing symbolic execution on '{}' .", &args.target));
logger.info(&format!("performing symbolic execution on '{}' .", &shortened_target));

// create a new progress bar
let progress = ProgressBar::new_spinner();
Expand All @@ -253,53 +252,28 @@ pub fn cfg(args: CFGArgs) {
// create a new petgraph StableGraph
let mut contract_cfg = Graph::<String, String>::new();

// perform EVM symbolic execution
for selector in selectors {
progress.set_message(format!("executing '0x{selector}'"));

// get the function's entry point
let function_entry_point = resolve_entry_point(&evm.clone(), selector.clone());
// add the call to the trace
let map_trace = trace.add_call(
vm_trace,
line!(),
"heimdall".to_string(),
"cfg".to_string(),
vec![format!("{} bytes", contract_bytecode.len()/2usize)],
"()".to_string()
);

// if the entry point is 0, then the function is not reachable
if function_entry_point == 0 {
continue;
}
// get a map of possible jump destinations
let (map, jumpdest_count) = map_contract(&evm.clone());

// add the call to the trace
let func_analysis_trace = trace.add_call(
vm_trace,
line!(),
"heimdall".to_string(),
"analyze".to_string(),
vec![format!("0x{selector}")],
"()".to_string()
);

// add the entry point to the trace
trace.add_info(
func_analysis_trace,
function_entry_point.try_into().unwrap(),
format!("discovered entry point: {function_entry_point}").to_string()
);

// get a map of possible jump destinations
let (map, jumpdest_count) = map_selector(&evm.clone(), selector.clone(), function_entry_point);

map.build_cfg(&mut contract_cfg, None);

// add the jumpdest count* to the trace
trace.add_debug(
func_analysis_trace,
function_entry_point.try_into().unwrap(),
format!("execution tree {}",

match jumpdest_count {
0 => "appears to be linear".to_string(),
_ => format!("has {jumpdest_count} unique branches")
}
).to_string()
);
}
// add jumpdests to the trace
trace.add_info(
map_trace,
line!(),
format!("traced and executed {jumpdest_count} possible paths.")
);

map.build_cfg(&mut contract_cfg, None);

progress.finish_and_clear();
logger.info("symbolic execution completed.");
Expand Down
55 changes: 3 additions & 52 deletions heimdall/src/cfg/util.rs
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
use std::{str::FromStr, collections::{HashMap, VecDeque}};
use std::{collections::{HashMap, VecDeque}};

use ethers::{
prelude::{
Expand Down Expand Up @@ -138,60 +138,11 @@ pub fn find_function_selectors(assembly: String) -> Vec<String> {
function_selectors
}

// resolve a selector's function entry point from the EVM bytecode
pub fn resolve_entry_point(evm: &VM, selector: String) -> u64 {
let mut vm = evm.clone();
let mut flag_next_jumpi = false;
let mut function_entry_point = 0;

// 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 EQ and it matched the selector, the next jumpi is the entry point
if call.last_instruction.opcode == "14"
&& call.last_instruction.inputs[0].eq(&U256::from_str(&selector.clone()).unwrap())
&& call.last_instruction.outputs[0].eq(&U256::from_str("1").unwrap())
{
flag_next_jumpi = true;
}

// if we are flagging the next jumpi, and the opcode is a JUMPI, we have found the entry point
if flag_next_jumpi && call.last_instruction.opcode == "57" {
// it's safe to convert here because we know max bytecode length is ~25kb, way less than 2^64
function_entry_point = call.last_instruction.inputs[0].as_u64();
break;
}

if vm.exitcode != 255 || !vm.returndata.is_empty() {
break;
}
}

function_entry_point
}

// build a map of function jump possibilities from the EVM bytecode
pub fn map_selector(
pub fn map_contract(
evm: &VM,
selector: String,
entry_point: u64,
) -> (VMTrace, u32) {
let mut vm = evm.clone();
vm.calldata = selector;

// step through the bytecode until we reach the entry point
while (vm.bytecode.len() >= (vm.instruction * 2 + 2) as usize)
&& (vm.instruction <= entry_point.into())
{
vm.step();

// this shouldn't be necessary, but it's safer to have it
if vm.exitcode != 255 || !vm.returndata.is_empty() {
break;
}
}
let vm = evm.clone();

// the VM is at the function entry point, begin tracing
let mut branch_count = 0;
Expand Down

0 comments on commit b0c41c6

Please sign in to comment.