From d3d66ae9f029682136497d7f47e388a47719866f Mon Sep 17 00:00:00 2001 From: Jonathan Becker Date: Sat, 30 Dec 2023 10:23:42 -0600 Subject: [PATCH] feat(symbolic-execution): implement & utilize `run_with_timeout` (#257) * fix(symbolic-exec): break out of infinite `JUMP` loops where stack grows infinitely * chore(symbolic-exec): consolidate `JUMP`/`JUMPI` logic * feat(threading): implement `run_with_timeout` * feat(cfg): use `run_with_timeout` for symbolic execution * feat(decompile): use `run_with_timeout` for symbolic execution * feat(snapshot): use `run_with_timeout` for symbolic execution --- common/src/ether/evm/ext/exec/mod.rs | 2 +- common/src/utils/threading.rs | 64 ++++++++++++++++++++++++++++ core/src/cfg/mod.rs | 17 +++++++- core/src/decompile/mod.rs | 39 ++++++++++++++--- core/src/snapshot/mod.rs | 25 +++++++++-- core/tests/test_cfg.rs | 4 ++ core/tests/test_decompile.rs | 10 +++++ core/tests/test_snapshot.rs | 5 +++ 8 files changed, 155 insertions(+), 11 deletions(-) diff --git a/common/src/ether/evm/ext/exec/mod.rs b/common/src/ether/evm/ext/exec/mod.rs index 779cb88c..af84c06c 100644 --- a/common/src/ether/evm/ext/exec/mod.rs +++ b/common/src/ether/evm/ext/exec/mod.rs @@ -24,7 +24,7 @@ use crate::{ }; use std::collections::HashMap; -#[derive(Clone, Debug)] +#[derive(Clone, Debug, Default)] pub struct VMTrace { pub instruction: u128, pub gas_used: u128, diff --git a/common/src/utils/threading.rs b/common/src/utils/threading.rs index 0c7fe456..acacb410 100644 --- a/common/src/utils/threading.rs +++ b/common/src/utils/threading.rs @@ -63,6 +63,28 @@ pub fn task_pool< results } +/// Takes a function and some arguments, and runs the function in a separate thread. If the function +/// doesnt finish within the given timeout, the thread is killed, and the function returns None. +pub fn run_with_timeout(f: F, timeout: std::time::Duration) -> Option +where + T: Send + 'static, + F: FnOnce() -> T + Send + 'static, { + let (tx, rx) = unbounded(); + let handle = thread::spawn(move || { + let result = f(); + tx.send(result).unwrap(); + }); + + let result = rx.recv_timeout(timeout); + if result.is_err() { + handle.thread().unpark(); + return None + } + + handle.join().unwrap(); + Some(result.unwrap()) +} + #[cfg(test)] mod tests { use crate::utils::threading::*; @@ -109,4 +131,46 @@ mod tests { let results = task_pool(items, num_threads, f); assert!(results.is_empty()); } + + #[test] + fn test_run_with_timeout() { + // Test case with a function that finishes within the timeout + let timeout = std::time::Duration::from_secs(1); + let f = || 1; + let result = run_with_timeout(f, timeout); + assert_eq!(result, Some(1)); + + // Test case with a function that doesnt finish within the timeout + let timeout = std::time::Duration::from_millis(1); + let f = || std::thread::sleep(std::time::Duration::from_secs(1)); + let result = run_with_timeout(f, timeout); + assert_eq!(result, None); + } + + #[test] + fn test_run_with_timeout_with_panic() { + // Test case with a function that panics + let timeout = std::time::Duration::from_secs(1); + let f = || panic!("test"); + let result = run_with_timeout(f, timeout); + assert_eq!(result, None); + } + + #[test] + fn test_run_with_timeout_with_args() { + // Test case with a function that takes arguments + let timeout = std::time::Duration::from_secs(1); + let f = |x: i32| x * 2; + let result = run_with_timeout(move || f(2), timeout); + assert_eq!(result, Some(4)); + } + + #[test] + fn test_run_with_timeout_infinite_loop() { + // Test case with a function that runs an infinite loop + let timeout = std::time::Duration::from_secs(1); + let f = || loop {}; + let result = run_with_timeout(f, timeout); + assert_eq!(result, None); + } } diff --git a/core/src/cfg/mod.rs b/core/src/cfg/mod.rs index 06a10303..af2d85d6 100644 --- a/core/src/cfg/mod.rs +++ b/core/src/cfg/mod.rs @@ -7,6 +7,7 @@ use heimdall_common::{ bytecode::get_bytecode_from_target, compiler::detect_compiler, selectors::find_function_selectors, }, + utils::threading::run_with_timeout, }; use indicatif::ProgressBar; use std::time::Duration; @@ -56,6 +57,10 @@ pub struct CFGArgs { /// The name for the output file #[clap(long, short, default_value = "", hide_default_value = true)] pub name: String, + + /// Timeout for symbolic execution + #[clap(long, short, default_value = "10000", hide_default_value = true)] + pub timeout: u64, } impl CFGArgsBuilder { @@ -68,6 +73,7 @@ impl CFGArgsBuilder { color_edges: Some(false), output: Some(String::new()), name: Some(String::new()), + timeout: Some(10000), } } } @@ -194,7 +200,14 @@ pub async fn cfg(args: CFGArgs) -> Result, Box map, + None => { + logger.error("symbolic execution timed out."); + return Err("symbolic execution timed out.".into()) + } + }; // add jumpdests to the trace trace.add_info( @@ -204,7 +217,7 @@ pub async fn cfg(args: CFGArgs) -> Result, Box map, + None => { + trace.add_error( + func_analysis_trace, + line!(), + &format!("symbolic execution timed out!"), + ); + (VMTrace::default(), 0) + } + }; trace.add_debug( func_analysis_trace, @@ -276,7 +297,7 @@ pub async fn decompile( if args.include_yul { debug_max!("analyzing symbolic execution trace '0x{}' with yul analyzer", selector); analyzed_function = analyze_yul( - map, + &map, Function { selector: selector.clone(), entry_point: function_entry_point, @@ -301,7 +322,7 @@ pub async fn decompile( } else { debug_max!("analyzing symbolic execution trace '0x{}' with sol analyzer", selector); analyzed_function = analyze_sol( - map, + &map, Function { selector: selector.clone(), entry_point: function_entry_point, @@ -326,6 +347,14 @@ pub async fn decompile( ); } + // add notice to analyzed_function if jumpdest_count == 0, indicating that + // symbolic execution timed out + if jumpdest_count == 0 { + analyzed_function + .notices + .push("symbolic execution timed out. please report this!".to_string()); + } + let argument_count = analyzed_function.arguments.len(); if argument_count != 0 { diff --git a/core/src/snapshot/mod.rs b/core/src/snapshot/mod.rs index 14b838a9..69c5e2c0 100644 --- a/core/src/snapshot/mod.rs +++ b/core/src/snapshot/mod.rs @@ -4,7 +4,7 @@ pub mod menus; pub mod resolve; pub mod structures; pub mod util; -use heimdall_common::debug_max; +use heimdall_common::{debug_max, utils::threading::run_with_timeout}; use std::{ collections::{HashMap, HashSet}, @@ -76,6 +76,10 @@ pub struct SnapshotArgs { /// The output directory to write the output to, or 'print' to print to the console. #[clap(long = "output", short = 'o', default_value = "output", hide_default_value = true)] pub output: String, + + /// The timeout for each function's symbolic execution in milliseconds. + #[clap(long, short, default_value = "10000", hide_default_value = true)] + pub timeout: u64, } impl SnapshotArgsBuilder { @@ -89,6 +93,7 @@ impl SnapshotArgsBuilder { no_tui: Some(true), name: Some(String::new()), output: Some(String::new()), + timeout: Some(10000), } } } @@ -257,8 +262,22 @@ async fn get_snapshots( ); // get a map of possible jump destinations - let (map, jumpdest_count) = - evm.clone().symbolic_exec_selector(&selector, function_entry_point); + let mut evm_clone = evm.clone(); + let selector_clone = selector.clone(); + let (map, jumpdest_count) = match run_with_timeout( + move || evm_clone.symbolic_exec_selector(&selector_clone, function_entry_point), + Duration::from_millis(args.timeout), + ) { + Some(map) => map, + None => { + trace.add_error( + func_analysis_trace, + line!(), + &format!("symbolic execution timed out, skipping snapshotting."), + ); + continue + } + }; trace.add_debug( func_analysis_trace, diff --git a/core/tests/test_cfg.rs b/core/tests/test_cfg.rs index d1d767e4..92d2bc78 100644 --- a/core/tests/test_cfg.rs +++ b/core/tests/test_cfg.rs @@ -16,6 +16,7 @@ mod benchmark { color_edges: false, output: String::from(""), name: String::from(""), + timeout: 10000, }; let _ = heimdall_core::cfg::cfg(args).await; } @@ -34,6 +35,7 @@ mod benchmark { color_edges: false, output: String::from(""), name: String::from(""), + timeout: 10000, }; let _ = heimdall_core::cfg::cfg(args).await; } @@ -58,6 +60,7 @@ mod integration_tests { color_edges: false, output: String::from(""), name: String::from(""), + timeout: 10000, }) .await .unwrap(); @@ -82,6 +85,7 @@ mod integration_tests { color_edges: false, output: String::from(""), name: String::from(""), + timeout: 10000, }) .await .unwrap(); diff --git a/core/tests/test_decompile.rs b/core/tests/test_decompile.rs index cef193e8..93f9f2ac 100644 --- a/core/tests/test_decompile.rs +++ b/core/tests/test_decompile.rs @@ -18,6 +18,7 @@ mod benchmark { include_yul: false, output: String::from(""), name: String::from(""), + timeout: 10000, }; let _ = heimdall_core::decompile::decompile(args).await; } @@ -38,6 +39,7 @@ mod benchmark { include_yul: false, output: String::from(""), name: String::from(""), + timeout: 10000, }; let _ = heimdall_core::decompile::decompile(args).await; } @@ -58,6 +60,7 @@ mod benchmark { include_yul: true, output: String::from(""), name: String::from(""), + timeout: 10000, }; let _ = heimdall_core::decompile::decompile(args).await; } @@ -78,6 +81,7 @@ mod benchmark { include_yul: true, output: String::from(""), name: String::from(""), + timeout: 10000, }; let _ = heimdall_core::decompile::decompile(args).await; } @@ -98,6 +102,7 @@ mod benchmark { include_yul: false, output: String::from(""), name: String::from(""), + timeout: 10000, }; let _ = heimdall_core::decompile::decompile(args).await; } @@ -118,6 +123,7 @@ mod benchmark { include_yul: false, output: String::from(""), name: String::from(""), + timeout: 10000, }; let _ = heimdall_core::decompile::decompile(args).await; } @@ -144,6 +150,7 @@ mod integration_tests { include_yul: false, output: String::from(""), name: String::from(""), + timeout: 10000, }) .await .unwrap(); @@ -173,6 +180,7 @@ mod integration_tests { include_yul: false, output: String::from(""), name: String::from(""), + timeout: 10000, }) .await .unwrap(); @@ -209,6 +217,7 @@ mod integration_tests { include_yul: false, output: String::from(""), name: String::from(""), + timeout: 10000, }) .await .unwrap(); @@ -322,6 +331,7 @@ mod integration_tests { include_yul: false, output: String::from(""), name: String::from(""), + timeout: 10000, }) .await .unwrap(); diff --git a/core/tests/test_snapshot.rs b/core/tests/test_snapshot.rs index a43b595c..33e2dfc8 100644 --- a/core/tests/test_snapshot.rs +++ b/core/tests/test_snapshot.rs @@ -17,6 +17,7 @@ mod benchmark { no_tui: true, name: String::from(""), output: String::from(""), + timeout: 10000, }; let _ = heimdall_core::snapshot::snapshot(args).await.unwrap(); } @@ -36,6 +37,7 @@ mod benchmark { no_tui: true, name: String::from(""), output: String::from(""), + timeout: 10000, }; let _ = heimdall_core::snapshot::snapshot(args).await.unwrap(); } @@ -61,6 +63,7 @@ mod integration_tests { no_tui: true, name: String::from(""), output: String::from(""), + timeout: 10000, }; let _ = heimdall_core::snapshot::snapshot(args).await.unwrap(); @@ -77,6 +80,7 @@ mod integration_tests { no_tui: true, name: String::from(""), output: String::from(""), + timeout: 10000, }; let _ = heimdall_core::snapshot::snapshot(args).await.unwrap(); @@ -167,6 +171,7 @@ mod integration_tests { no_tui: true, name: String::from(""), output: String::from(""), + timeout: 10000, }; let _ = heimdall_core::snapshot::snapshot(args).await.unwrap(); }