feat(cheatcodes) vm.prompt: Prompt user for interactive input (foundry-rs#7012)

* Implement vm.prompt cheatcode

* chore: speedup prompt test locally

* move prompt.sol

---------

Co-authored-by: Matthias Seitz <[email protected]>

Author: Tudmotu

Cargo.lock

@@ -43,3 +43,4 @@ walkdir = "2"
 p256 = "0.13.2"
 thiserror = "1"
 rustc-hash.workspace = true
+dialoguer = "0.11.0"

crates/cheatcodes/Cargo.toml

@@ -1399,6 +1399,16 @@ interface Vm {
     #[cheatcode(group = Filesystem)]
     function tryFfi(string[] calldata commandInput) external returns (FfiResult memory result);
 
+    // -------- User Interaction --------
+
+    /// Prompts the user for a string value in the terminal.
+    #[cheatcode(group = Filesystem)]
+    function prompt(string calldata promptText) external returns (string memory input);
+
+    /// Prompts the user for a hidden string value in the terminal.
+    #[cheatcode(group = Filesystem)]
+    function promptSecret(string calldata promptText) external returns (string memory input);
+
     // ======== Environment Variables ========
 
     /// Sets environment variables.

crates/cheatcodes/assets/cheatcodes.json

@@ -11,6 +11,7 @@ use foundry_evm_core::opts::EvmOpts;
 use std::{
     collections::HashMap,
     path::{Path, PathBuf},
+    time::Duration,
 };
 
 /// Additional, configurable context the `Cheatcodes` inspector has access to
@@ -22,6 +23,8 @@ pub struct CheatsConfig {
     pub ffi: bool,
     /// Use the create 2 factory in all cases including tests and non-broadcasting scripts.
     pub always_use_create_2_factory: 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
@@ -55,6 +58,7 @@ impl CheatsConfig {
         Self {
             ffi: evm_opts.ffi,
             always_use_create_2_factory: evm_opts.always_use_create_2_factory,
+            prompt_timeout: Duration::from_secs(config.prompt_timeout),
             rpc_storage_caching: config.rpc_storage_caching.clone(),
             rpc_endpoints,
             paths: config.project_paths(),
@@ -171,6 +175,7 @@ impl Default for CheatsConfig {
         Self {
             ffi: false,
             always_use_create_2_factory: false,
+            prompt_timeout: Duration::from_secs(120),
             rpc_storage_caching: Default::default(),
             rpc_endpoints: Default::default(),
             paths: ProjectPathsConfig::builder().build_with_root("./"),

crates/cheatcodes/spec/src/vm.rs

@@ -4,13 +4,16 @@ use crate::{Cheatcode, Cheatcodes, Result, Vm::*};
 use alloy_json_abi::ContractObject;
 use alloy_primitives::U256;
 use alloy_sol_types::SolValue;
+use dialoguer::{Input, Password};
 use foundry_common::{fs, get_artifact_path};
 use foundry_config::fs_permissions::FsAccessKind;
 use std::{
     collections::hash_map::Entry,
     io::{BufRead, BufReader, Write},
     path::Path,
     process::Command,
+    sync::mpsc,
+    thread,
     time::{SystemTime, UNIX_EPOCH},
 };
 use walkdir::WalkDir;
@@ -296,6 +299,20 @@ impl Cheatcode for tryFfiCall {
     }
 }
 
+impl Cheatcode for promptCall {
+    fn apply(&self, state: &mut Cheatcodes) -> Result {
+        let Self { promptText: text } = self;
+        prompt(state, text, prompt_input).map(|res| res.abi_encode())
+    }
+}
+
+impl Cheatcode for promptSecretCall {
+    fn apply(&self, state: &mut Cheatcodes) -> Result {
+        let Self { promptText: text } = self;
+        prompt(state, text, prompt_password).map(|res| res.abi_encode())
+    }
+}
+
 pub(super) fn write_file(state: &Cheatcodes, path: &Path, contents: &[u8]) -> Result {
     let path = state.config.ensure_path_allowed(path, FsAccessKind::Write)?;
     // write access to foundry.toml is not allowed
@@ -370,6 +387,39 @@ fn ffi(state: &Cheatcodes, input: &[String]) -> Result<FfiResult> {
     })
 }
 
+fn prompt_input(prompt_text: &str) -> Result<String, dialoguer::Error> {
+    Input::new().allow_empty(true).with_prompt(prompt_text).interact_text()
+}
+
+fn prompt_password(prompt_text: &str) -> Result<String, dialoguer::Error> {
+    Password::new().with_prompt(prompt_text).interact()
+}
+
+fn prompt(
+    state: &Cheatcodes,
+    prompt_text: &str,
+    input: fn(&str) -> 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();
+    });
+
+    match recv.recv_timeout(timeout) {
+        Ok(res) => res.map_err(|err| {
+            println!();
+            err.to_string().into()
+        }),
+        Err(_) => {
+            println!();
+            Err("Prompt timed out".into())
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

crates/cheatcodes/src/config.rs

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

crates/cheatcodes/src/fs.rs

@@ -244,6 +244,8 @@ pub struct Config {
     pub ffi: bool,
     /// Use the create 2 factory in all cases including tests and non-broadcasting scripts.
     pub always_use_create_2_factory: bool,
+    /// Sets a timeout in seconds 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
@@ -1873,6 +1875,7 @@ impl Default for Config {
             invariant: Default::default(),
             always_use_create_2_factory: false,
             ffi: false,
+            prompt_timeout: 120,
             sender: Config::DEFAULT_SENDER,
             tx_origin: Config::DEFAULT_SENDER,
             initial_balance: U256::from(0xffffffffffffffffffffffffu128),

crates/config/README.md

@@ -520,6 +520,7 @@ impl CallTraceDecoder {
         match func.name.as_str() {
             s if s.starts_with("env") => Some("<env var value>"),
             "createWallet" | "deriveKey" => Some("<pk>"),
+            "promptSecret" => Some("<secret>"),
             "parseJson" if self.verbosity < 5 => Some("<encoded JSON value>"),
             "readFile" if self.verbosity < 5 => Some("<file>"),
             _ => None,

crates/config/src/lib.rs

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

crates/evm/traces/src/decoder/mod.rs

@@ -197,6 +197,9 @@ impl ForgeTestData {
         config.rpc_endpoints = rpc_endpoints();
         config.allow_paths.push(manifest_root().to_path_buf());
 
+        // no prompt testing
+        config.prompt_timeout = 0;
+
         let root = self.project.root();
         let opts = self.evm_opts.clone();
         let env = opts.local_evm_env();

crates/forge/tests/cli/config.rs

@@ -0,0 +1,18 @@
+// SPDX-License-Identifier: MIT OR Apache-2.0
+pragma solidity 0.8.18;
+
+import "ds-test/test.sol";
+import "cheats/Vm.sol";
+
+contract PromptTest is DSTest {
+    Vm constant vm = Vm(HEVM_ADDRESS);
+
+    function testPrompt_revertNotATerminal() public {
+        // should revert in CI and testing environments either with timout or because no terminal is available
+        vm._expectCheatcodeRevert();
+        vm.prompt("test");
+
+        vm._expectCheatcodeRevert();
+        vm.promptSecret("test");
+    }
+}