From cb6cdeb09ad9a15e96172579308813e1451a96f7 Mon Sep 17 00:00:00 2001
From: Jon-Becker <jonathan@jbecker.dev>
Date: Sat, 30 Dec 2023 10:19:24 -0600
Subject: [PATCH] feat(snapshot): use `run_with_timeout` for symbolic execution

---
 core/src/snapshot/mod.rs    | 25 ++++++++++++++++++++++---
 core/tests/test_snapshot.rs |  5 +++++
 2 files changed, 27 insertions(+), 3 deletions(-)

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