Add configurable timeout for vm.prompt

crates/cheatcodes/src/config.rs

@@ -1,6 +1,7 @@
 use super::Result;
 use crate::Vm::Rpc;
 use alloy_primitives::Address;
+use core::time::Duration;
 use foundry_common::fs::normalize_path;
 use foundry_compilers::{utils::canonicalize, ProjectPathsConfig};
 use foundry_config::{
@@ -20,6 +21,8 @@ use std::{
 pub struct CheatsConfig {
     /// Whether the FFI cheatcode is enabled.
     pub ffi: bool,
+    /// Sets a timeout for vm.prompt cheatcodes
+    pub prompt_timeout: Duration,
     /// RPC storage caching settings determines what chains and endpoints to cache
     pub rpc_storage_caching: StorageCachingConfig,
     /// All known endpoints and their aliases
@@ -50,6 +53,7 @@ impl CheatsConfig {
 
         Self {
             ffi: evm_opts.ffi,
+            prompt_timeout: Duration::from_secs(evm_opts.prompt_timeout),
             rpc_storage_caching: config.rpc_storage_caching.clone(),
             rpc_endpoints,
             paths: config.project_paths(),
@@ -164,6 +168,7 @@ impl Default for CheatsConfig {
     fn default() -> Self {
         Self {
             ffi: false,
+            prompt_timeout: Duration::from_secs(120),
             rpc_storage_caching: Default::default(),
             rpc_endpoints: Default::default(),
             paths: ProjectPathsConfig::builder().build_with_root("./"),

crates/cheatcodes/src/fs.rs

@@ -12,6 +12,8 @@ use std::{
     io::{BufRead, BufReader, Write},
     path::Path,
     process::Command,
+    sync::mpsc,
+    thread,
     time::{SystemTime, UNIX_EPOCH},
 };
 use walkdir::WalkDir;
@@ -298,16 +300,16 @@ impl Cheatcode for tryFfiCall {
 }
 
 impl Cheatcode for promptCall {
-    fn apply(&self, _state: &mut Cheatcodes) -> Result {
+    fn apply(&self, state: &mut Cheatcodes) -> Result {
         let Self { promptText: text } = self;
-        prompt(text).map(|res| res.abi_encode())
+        prompt(state, text, prompt_input).map(|res| res.abi_encode())
     }
 }
 
 impl Cheatcode for promptSecretCall {
-    fn apply(&self, _state: &mut Cheatcodes) -> Result {
+    fn apply(&self, state: &mut Cheatcodes) -> Result {
         let Self { promptText: text } = self;
-        prompt_secret(text).map(|res| res.abi_encode())
+        prompt(state, text, prompt_password).map(|res| res.abi_encode())
     }
 }
 
@@ -385,17 +387,37 @@ fn ffi(state: &Cheatcodes, input: &[String]) -> Result<FfiResult> {
     })
 }
 
-fn prompt(prompt_text: &String) -> Result<String> {
-    let input: String =
-        Input::new().allow_empty(true).with_prompt(prompt_text).interact_text().unwrap();
+fn prompt_input(prompt_text: &String) -> Result<String, dialoguer::Error> {
+    Input::new().allow_empty(true).with_prompt(prompt_text).interact_text()
+}
 
-    Ok(input)
+fn prompt_password(prompt_text: &String) -> Result<String, dialoguer::Error> {
+    Password::new().with_prompt(prompt_text).interact()
 }
 
-fn prompt_secret(prompt_text: &String) -> Result<String> {
-    let input: String = Password::new().with_prompt(prompt_text).interact().unwrap();
+fn prompt(
+    state: &Cheatcodes,
+    prompt_text: &str,
+    input: fn(&String) -> Result<String, dialoguer::Error>,
+) -> Result<String> {
+    let text_clone = prompt_text.to_string();
+    let timeout = state.config.prompt_timeout;
+    let (send, recv) = mpsc::channel();
+
+    thread::spawn(move || {
+        send.send(input(&text_clone)).unwrap();
+    });
 
-    Ok(input)
+    match recv.recv_timeout(timeout) {
+        Ok(res) => res.map_err(|_| {
+            println!();
+            "I/O error occured".into()
+        }),
+        Err(_) => {
+            println!();
+            Err("Prompt timed out".into())
+        }
+    }
 }
 
 #[cfg(test)]

crates/config/README.md

@@ -114,6 +114,7 @@ no_match_contract = "Bar"
 match_path = "*/Foo*"
 no_match_path = "*/Bar*"
 ffi = false
+prompt_timeout = 120
 # These are the default callers, generated using `address(uint160(uint256(keccak256("foundry default caller"))))`
 sender = '0x1804c8AB1F12E6bbf3894d4083f33e07309d1f38'
 tx_origin = '0x1804c8AB1F12E6bbf3894d4083f33e07309d1f38'

crates/config/src/lib.rs

@@ -240,6 +240,8 @@ pub struct Config {
     pub invariant: InvariantConfig,
     /// Whether to allow ffi cheatcodes in test
     pub ffi: bool,
+    /// Sets a timeout for vm.prompt cheatcodes
+    pub prompt_timeout: u64,
     /// The address which will be executing all tests
     pub sender: Address,
     /// The tx.origin value during EVM execution
@@ -1794,6 +1796,7 @@ impl Default for Config {
             fuzz: Default::default(),
             invariant: Default::default(),
             ffi: false,
+            prompt_timeout: 120,
             sender: Config::DEFAULT_SENDER,
             tx_origin: Config::DEFAULT_SENDER,
             initial_balance: U256::from(0xffffffffffffffffffffffffu128),

crates/evm/core/src/opts.rs

@@ -53,6 +53,9 @@ pub struct EvmOpts {
     /// Enables the FFI cheatcode.
     pub ffi: bool,
 
+    /// Sets a timeout for vm.prompt cheatcodes
+    pub prompt_timeout: u64,
+
     /// Verbosity mode of EVM output as number of occurrences.
     pub verbosity: u8,
 

crates/forge/tests/cli/config.rs

@@ -64,6 +64,7 @@ forgetest!(can_extract_config_values, |prj, cmd| {
         },
         invariant: InvariantConfig { runs: 256, ..Default::default() },
         ffi: true,
+        prompt_timeout: 0,
         sender: "00a329c0648769A73afAc7F9381D08FB43dBEA72".parse().unwrap(),
         tx_origin: "00a329c0648769A73afAc7F9F81E08FB43dBEA72".parse().unwrap(),
         initial_balance: U256::from(0xffffffffffffffffffffffffu128),