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

@@ -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,

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<T, F>(f: F, timeout: std::time::Duration) -> Option<T>
+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);
+    }
 }

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<Graph<String, String>, Box<dyn std::er
     );
 
     // get a map of possible jump destinations
-    let (map, jumpdest_count) = &evm.symbolic_exec();
+    let (map, jumpdest_count) =
+        match run_with_timeout(move || evm.symbolic_exec(), Duration::from_millis(args.timeout)) {
+            Some(map) => 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<Graph<String, String>, Box<dyn std::er
     );
 
     debug_max!("building control flow graph from symbolic execution trace");
-    build_cfg(map, &mut contract_cfg, None, false);
+    build_cfg(&map, &mut contract_cfg, None, false);
 
     progress.finish_and_clear();
     logger.info("symbolic execution completed.");

core/src/decompile/mod.rs

@@ -5,7 +5,9 @@ pub mod precompile;
 pub mod resolve;
 pub mod util;
 use heimdall_common::{
-    debug_max, ether::bytecode::get_bytecode_from_target, utils::strings::get_shortned_target,
+    debug_max,
+    ether::{bytecode::get_bytecode_from_target, evm::ext::exec::VMTrace},
+    utils::{strings::get_shortned_target, threading::run_with_timeout},
 };
 
 use crate::{
@@ -80,6 +82,10 @@ pub struct DecompilerArgs {
     /// The name for the output file
     #[clap(long, short, default_value = "", hide_default_value = true)]
     pub name: 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 DecompilerArgsBuilder {
@@ -94,6 +100,7 @@ impl DecompilerArgsBuilder {
             include_yul: Some(false),
             output: Some(String::new()),
             name: Some(String::new()),
+            timeout: Some(10000),
         }
     }
 }
@@ -252,8 +259,22 @@ pub async fn decompile(
         );
 
         // 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!"),
+                );
+                (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 {

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,

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();

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();