Skip to content

Commit

Permalink
wip
Browse files Browse the repository at this point in the history
  • Loading branch information
@Jon-Becker
Jon-Becker committed Dec 22, 2024
1 parent 6214f67 commit 2c1a4ea
Show file tree
Hide file tree
Showing 3 changed files with 103 additions and 143 deletions.
148 changes: 96 additions & 52 deletions crates/common/src/ether/signatures.rs
View file
Original file line number Diff line number Diff line change
@@ -1,8 +1,13 @@
use std::path::PathBuf;
use std::{
path::PathBuf,
sync::{Arc, Mutex},
time::Instant,
};

use alloy_dyn_abi::{DynSolType, DynSolValue};
use alloy_json_abi::JsonAbi;
use async_trait::async_trait;
use hashbrown::HashMap;

use crate::{
ether::types::parse_function_parameters,
Expand All @@ -15,7 +20,7 @@ use crate::{
use eyre::{OptionExt, Result};
use heimdall_cache::{store_cache, with_cache};
use serde::{Deserialize, Serialize};
use tracing::{debug, trace};
use tracing::{debug, error, info, trace, warn};

use super::types::DynSolValueExt;

Expand Down Expand Up @@ -306,6 +311,56 @@ impl ResolveSelector for ResolvedFunction {
}
}

impl TryFrom<&ResolvedFunction> for TraceFactory {
// eyre
type Error = eyre::Report;

fn try_from(function: &ResolvedFunction) -> Result<Self, Self::Error> {
let mut trace = TraceFactory::default();
let decode_call = trace.add_call(
0,
line!(),
"heimdall".to_string(),
"decode".to_string(),
vec![],
"()".to_string(),
);
trace.br(decode_call);
trace.add_message(decode_call, line!(), vec![format!("signature: {}", function.signature)]);
trace.br(decode_call);

// build inputs
for (i, input) in function.decoded_inputs.as_ref().unwrap_or(&Vec::new()).iter().enumerate()
{
let mut decoded_inputs_as_message = display(vec![input.to_owned()], " ");
if decoded_inputs_as_message.is_empty() {
break;
}

if i == 0 {
decoded_inputs_as_message[0] = format!(
"input {}:{}{}",
i,
" ".repeat(4 - i.to_string().len()),
decoded_inputs_as_message[0].replacen(" ", "", 1)
)
} else {
decoded_inputs_as_message[0] = format!(
" {}:{}{}",
i,
" ".repeat(4 - i.to_string().len()),
decoded_inputs_as_message[0].replacen(" ", "", 1)
)
}

// add to trace and decoded string
trace.add_message(decode_call, 1, decoded_inputs_as_message);
}

Ok(trace)
}
}

/// Given the path to an ABI file, parses all [`ResolvedFunction`]s, [`ResolvedError`]s, and
/// [`ResolvedLog`]s from the ABI and saves them to the cache.
pub fn cache_signatures_from_abi(path: PathBuf) -> Result<()> {
Expand Down Expand Up @@ -356,6 +411,7 @@ pub fn cache_signatures_from_abi(path: PathBuf) -> Result<()> {
Ok(())
}

/// A heuristic function to score a function signature based on its spamminess.
pub fn score_signature(signature: &str, num_words: Option<usize>) -> u32 {
// the score starts at 1000
let mut score = 1000;
Expand Down Expand Up @@ -390,64 +446,52 @@ pub fn score_signature(signature: &str, num_words: Option<usize>) -> u32 {
score
}

/// trait impls
/// trait impls
/// trait impls
impl TryFrom<&ResolvedFunction> for TraceFactory {
// eyre
type Error = eyre::Report;

fn try_from(function: &ResolvedFunction) -> Result<Self, Self::Error> {
let mut trace = TraceFactory::default();
let decode_call = trace.add_call(
0,
line!(),
"heimdall".to_string(),
"decode".to_string(),
vec![],
"()".to_string(),
);
trace.br(decode_call);
trace.add_message(decode_call, line!(), vec![format!("signature: {}", function.signature)]);
trace.br(decode_call);

// build inputs
for (i, input) in function.decoded_inputs.as_ref().unwrap_or(&Vec::new()).iter().enumerate()
{
let mut decoded_inputs_as_message = display(vec![input.to_owned()], " ");
if decoded_inputs_as_message.is_empty() {
break;
}
/// Resolve a list of selectors to their function signatures.
pub async fn resolve_selectors<T>(selectors: Vec<String>) -> HashMap<String, Vec<T>>
where
T: ResolveSelector + Send + Clone + 'static, {
// short-circuit if there are no selectors
if selectors.is_empty() {
return HashMap::new();
}

if i == 0 {
decoded_inputs_as_message[0] = format!(
"input {}:{}{}",
i,
" ".repeat(4 - i.to_string().len()),
decoded_inputs_as_message[0].replacen(" ", "", 1)
)
} else {
decoded_inputs_as_message[0] = format!(
" {}:{}{}",
i,
" ".repeat(4 - i.to_string().len()),
decoded_inputs_as_message[0].replacen(" ", "", 1)
)
let resolved_functions: Arc<Mutex<HashMap<String, Vec<T>>>> =
Arc::new(Mutex::new(HashMap::new()));
let mut threads = Vec::new();
let start_time = Instant::now();
let selector_count = selectors.len();

for selector in selectors {
let function_clone = resolved_functions.clone();

// create a new thread for each selector
threads.push(tokio::task::spawn(async move {
if let Ok(Some(function)) = T::resolve(&selector).await {
let mut _resolved_functions =
function_clone.lock().expect("Could not obtain lock on function_clone.");
_resolved_functions.insert(selector, function);
}
}));
}

// add to trace and decoded string
trace.add_message(decode_call, 1, decoded_inputs_as_message);
// wait for all threads to finish
for thread in threads {
if let Err(e) = thread.await {
// Handle error
error!("failed to resolve selector: {:?}", e);
}
}

Ok(trace)
let signatures =
resolved_functions.lock().expect("failed to obtain lock on resolved_functions.").clone();
if signatures.is_empty() {
warn!("failed to resolve any signatures from {} selectors", selector_count);
}
info!("resolved {} signatures from {} selectors", signatures.len(), selector_count);
debug!("signature resolution took {:?}", start_time.elapsed());
signatures
}

/// tests
/// tests
/// tests
#[cfg(test)]
mod tests {
use heimdall_cache::delete_cache;
Expand Down
9 changes: 3 additions & 6 deletions crates/decompile/src/core/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,18 +12,15 @@ use heimdall_common::{
ether::{
compiler::detect_compiler,
signatures::{
cache_signatures_from_abi, score_signature, ResolvedError, ResolvedFunction,
ResolvedLog,
cache_signatures_from_abi, resolve_selectors, score_signature, ResolvedError,
ResolvedFunction, ResolvedLog,
},
types::to_type,
},
utils::strings::{decode_hex, encode_hex, encode_hex_reduced, StringExt},
};
use heimdall_disassembler::{disassemble, DisassemblerArgsBuilder};
use heimdall_vm::{
core::vm::VM,
ext::selectors::{find_function_selectors, resolve_selectors},
};
use heimdall_vm::{core::vm::VM, ext::selectors::find_function_selectors};
use std::time::{Duration, Instant};

use crate::{
Expand Down
89 changes: 4 additions & 85 deletions crates/vm/src/ext/selectors.rs
View file
Original file line number Diff line number Diff line change
@@ -1,44 +1,9 @@
use hashbrown::{HashMap, HashSet};
use std::{
sync::{Arc, Mutex},
time::Instant,
};

use eyre::Result;
use heimdall_common::{
ether::signatures::{ResolveSelector, ResolvedFunction},
utils::strings::decode_hex,
};
use tokio::task;
use tracing::{debug, error, info, trace, warn};

use crate::core::vm::VM;

// Find all function selectors and all the data associated to this function, represented by
// [`ResolvedFunction`]
pub async fn get_resolved_selectors(
disassembled_bytecode: &str,
skip_resolving: &bool,
evm: &VM,
) -> Result<(HashMap<String, u128>, HashMap<String, Vec<ResolvedFunction>>)> {
let selectors = find_function_selectors(evm, disassembled_bytecode);

let mut resolved_selectors = HashMap::new();
if !skip_resolving {
resolved_selectors =
resolve_selectors::<ResolvedFunction>(selectors.keys().cloned().collect()).await;

trace!(
"resolved {} possible functions from {} detected selectors.",
resolved_selectors.len(),
selectors.len()
);
} else {
trace!("found {} possible function selectors.", selectors.len());
}
use heimdall_common::utils::strings::decode_hex;
use tracing::{info, trace};

Ok((selectors, resolved_selectors))
}
use crate::core::vm::VM;

/// find all function selectors in the given EVM assembly.
// TODO: update get_resolved_selectors logic to support vyper, huff
Expand Down Expand Up @@ -97,7 +62,7 @@ pub fn find_function_selectors(evm: &VM, assembly: &str) -> HashMap<String, u128

/// resolve a selector's function entry point from the EVM bytecode
// TODO: update resolve_entry_point logic to support vyper
pub fn resolve_entry_point(vm: &mut VM, selector: &str) -> u128 {
fn resolve_entry_point(vm: &mut VM, selector: &str) -> u128 {
let mut handled_jumps = HashSet::new();

// execute the EVM call to find the entry point for the given selector
Expand Down Expand Up @@ -138,49 +103,3 @@ pub fn resolve_entry_point(vm: &mut VM, selector: &str) -> u128 {

0
}

/// Resolve a list of selectors to their function signatures.
pub async fn resolve_selectors<T>(selectors: Vec<String>) -> HashMap<String, Vec<T>>
where
T: ResolveSelector + Send + Clone + 'static, {
// short-circuit if there are no selectors
if selectors.is_empty() {
return HashMap::new();
}

let resolved_functions: Arc<Mutex<HashMap<String, Vec<T>>>> =
Arc::new(Mutex::new(HashMap::new()));
let mut threads = Vec::new();
let start_time = Instant::now();
let selector_count = selectors.len();

for selector in selectors {
let function_clone = resolved_functions.clone();

// create a new thread for each selector
threads.push(task::spawn(async move {
if let Ok(Some(function)) = T::resolve(&selector).await {
let mut _resolved_functions =
function_clone.lock().expect("Could not obtain lock on function_clone.");
_resolved_functions.insert(selector, function);
}
}));
}

// wait for all threads to finish
for thread in threads {
if let Err(e) = thread.await {
// Handle error
error!("failed to resolve selector: {:?}", e);
}
}

let signatures =
resolved_functions.lock().expect("failed to obtain lock on resolved_functions.").clone();
if signatures.is_empty() {
warn!("failed to resolve any signatures from {} selectors", selector_count);
}
info!("resolved {} signatures from {} selectors", signatures.len(), selector_count);
debug!("signature resolution took {:?}", start_time.elapsed());
signatures
}

0 comments on commit 2c1a4ea

Please sign in to comment.