From 8da53cfb53ff743a3d9ac1c1bb75e7320e5be17c Mon Sep 17 00:00:00 2001 From: Yotam Bar-On Date: Mon, 18 Mar 2024 23:43:07 +0200 Subject: [PATCH] Implement vm.prompt cheatcode --- Cargo.lock | 3 ++ crates/cheatcodes/Cargo.toml | 1 + crates/cheatcodes/assets/cheatcodes.json | 40 +++++++++++++++++++ crates/cheatcodes/spec/src/vm.rs | 10 +++++ crates/cheatcodes/src/config.rs | 5 +++ crates/cheatcodes/src/fs.rs | 50 ++++++++++++++++++++++++ crates/config/README.md | 1 + crates/config/src/lib.rs | 3 ++ crates/evm/traces/src/decoder/mod.rs | 1 + crates/forge/tests/cli/config.rs | 1 + testdata/cheats/Prompt.t.sol | 17 ++++++++ testdata/cheats/Vm.sol | 2 + 12 files changed, 134 insertions(+) create mode 100644 testdata/cheats/Prompt.t.sol diff --git a/Cargo.lock b/Cargo.lock index ab4caf50ef2d..1f40fa2bb3f4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2091,7 +2091,9 @@ checksum = "658bce805d770f407bc62102fca7c2c64ceef2fbcb2b8bd19d2765ce093980de" dependencies = [ "console", "shell-words", + "tempfile", "thiserror", + "zeroize", ] [[package]] @@ -3167,6 +3169,7 @@ dependencies = [ "alloy-sol-types", "base64 0.22.0", "const-hex", + "dialoguer", "eyre", "foundry-cheatcodes-spec", "foundry-common", diff --git a/crates/cheatcodes/Cargo.toml b/crates/cheatcodes/Cargo.toml index 57b39966e034..47315f5ce996 100644 --- a/crates/cheatcodes/Cargo.toml +++ b/crates/cheatcodes/Cargo.toml @@ -43,3 +43,4 @@ walkdir = "2" p256 = "0.13.2" thiserror = "1" rustc-hash.workspace = true +dialoguer = "0.11.0" diff --git a/crates/cheatcodes/assets/cheatcodes.json b/crates/cheatcodes/assets/cheatcodes.json index b086072263d1..9c69772b69f9 100644 --- a/crates/cheatcodes/assets/cheatcodes.json +++ b/crates/cheatcodes/assets/cheatcodes.json @@ -6038,6 +6038,46 @@ "status": "stable", "safety": "safe" }, + { + "func": { + "id": "prompt", + "description": "Prompts the user for a string value in the terminal.", + "declaration": "function prompt(string calldata promptText) external returns (string memory input);", + "visibility": "external", + "mutability": "", + "signature": "prompt(string)", + "selector": "0x47eaf474", + "selectorBytes": [ + 71, + 234, + 244, + 116 + ] + }, + "group": "filesystem", + "status": "stable", + "safety": "safe" + }, + { + "func": { + "id": "promptSecret", + "description": "Prompts the user for a hidden string value in the terminal.", + "declaration": "function promptSecret(string calldata promptText) external returns (string memory input);", + "visibility": "external", + "mutability": "", + "signature": "promptSecret(string)", + "selector": "0x1e279d41", + "selectorBytes": [ + 30, + 39, + 157, + 65 + ] + }, + "group": "filesystem", + "status": "stable", + "safety": "safe" + }, { "func": { "id": "readCallers", diff --git a/crates/cheatcodes/spec/src/vm.rs b/crates/cheatcodes/spec/src/vm.rs index bd6500fb7c81..cdcca3331785 100644 --- a/crates/cheatcodes/spec/src/vm.rs +++ b/crates/cheatcodes/spec/src/vm.rs @@ -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. diff --git a/crates/cheatcodes/src/config.rs b/crates/cheatcodes/src/config.rs index eb8f56ab7f6e..1048f6b6937d 100644 --- a/crates/cheatcodes/src/config.rs +++ b/crates/cheatcodes/src/config.rs @@ -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("./"), diff --git a/crates/cheatcodes/src/fs.rs b/crates/cheatcodes/src/fs.rs index 3e345db94a97..7789915b5657 100644 --- a/crates/cheatcodes/src/fs.rs +++ b/crates/cheatcodes/src/fs.rs @@ -4,6 +4,7 @@ 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::{ @@ -11,6 +12,8 @@ use std::{ 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 { }) } +fn prompt_input(prompt_text: &str) -> Result { + Input::new().allow_empty(true).with_prompt(prompt_text).interact_text() +} + +fn prompt_password(prompt_text: &str) -> Result { + Password::new().with_prompt(prompt_text).interact() +} + +fn prompt( + state: &Cheatcodes, + prompt_text: &str, + input: fn(&str) -> Result, +) -> Result { + 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::*; diff --git a/crates/config/README.md b/crates/config/README.md index 46fa83e5a302..ffa4cdc75025 100644 --- a/crates/config/README.md +++ b/crates/config/README.md @@ -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' diff --git a/crates/config/src/lib.rs b/crates/config/src/lib.rs index 4f27b58b8edd..dd5834dee158 100644 --- a/crates/config/src/lib.rs +++ b/crates/config/src/lib.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 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), diff --git a/crates/evm/traces/src/decoder/mod.rs b/crates/evm/traces/src/decoder/mod.rs index 2eac34d003a8..f53d8b516258 100644 --- a/crates/evm/traces/src/decoder/mod.rs +++ b/crates/evm/traces/src/decoder/mod.rs @@ -520,6 +520,7 @@ impl CallTraceDecoder { match func.name.as_str() { s if s.starts_with("env") => Some(""), "createWallet" | "deriveKey" => Some(""), + "promptSecret" => Some(""), "parseJson" if self.verbosity < 5 => Some(""), "readFile" if self.verbosity < 5 => Some(""), _ => None, diff --git a/crates/forge/tests/cli/config.rs b/crates/forge/tests/cli/config.rs index 6714df59dd83..81884f3bcd9f 100644 --- a/crates/forge/tests/cli/config.rs +++ b/crates/forge/tests/cli/config.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), diff --git a/testdata/cheats/Prompt.t.sol b/testdata/cheats/Prompt.t.sol new file mode 100644 index 000000000000..6d84ed8553a0 --- /dev/null +++ b/testdata/cheats/Prompt.t.sol @@ -0,0 +1,17 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 +pragma solidity 0.8.18; + +import "ds-test/test.sol"; +import "./Vm.sol"; + +contract PromptTest is DSTest { + Vm constant vm = Vm(HEVM_ADDRESS); + + function testPrompt_revertNotATerminal() public { + vm._expectCheatcodeRevert("IO error: not a terminal"); + vm.prompt("test"); + + vm._expectCheatcodeRevert("IO error: not a terminal"); + vm.promptSecret("test"); + } +} diff --git a/testdata/cheats/Vm.sol b/testdata/cheats/Vm.sol index 623ef254bf06..b5a604e14d96 100644 --- a/testdata/cheats/Vm.sol +++ b/testdata/cheats/Vm.sol @@ -299,6 +299,8 @@ interface Vm { function prank(address msgSender, address txOrigin) external; function prevrandao(bytes32 newPrevrandao) external; function projectRoot() external view returns (string memory path); + function prompt(string calldata promptText) external returns (string memory input); + function promptSecret(string calldata promptText) external returns (string memory input); function readCallers() external returns (CallerMode callerMode, address msgSender, address txOrigin); function readDir(string calldata path) external view returns (DirEntry[] memory entries); function readDir(string calldata path, uint64 maxDepth) external view returns (DirEntry[] memory entries);