From afe45981470d74f1cca0064e5252e5309a54bb60 Mon Sep 17 00:00:00 2001 From: Yotam Bar-On Date: Sat, 3 Feb 2024 15:38:47 +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/fs.rs | 28 +++++++++++++++++ testdata/cheats/Vm.sol | 2 ++ 6 files changed, 84 insertions(+) diff --git a/Cargo.lock b/Cargo.lock index c65fd0e82b4d..b35a4f7b9a16 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2014,7 +2014,9 @@ checksum = "658bce805d770f407bc62102fca7c2c64ceef2fbcb2b8bd19d2765ce093980de" dependencies = [ "console", "shell-words", + "tempfile", "thiserror", + "zeroize", ] [[package]] @@ -2981,6 +2983,7 @@ dependencies = [ "alloy-sol-types", "base64 0.21.7", "const-hex", + "dialoguer", "eyre", "foundry-cheatcodes-spec", "foundry-common", diff --git a/crates/cheatcodes/Cargo.toml b/crates/cheatcodes/Cargo.toml index 4c2e69a36483..423b2d5d6e60 100644 --- a/crates/cheatcodes/Cargo.toml +++ b/crates/cheatcodes/Cargo.toml @@ -40,3 +40,4 @@ k256.workspace = true walkdir = "2" p256 = "0.13.2" thiserror = "1" +dialoguer = "0.11.0" diff --git a/crates/cheatcodes/assets/cheatcodes.json b/crates/cheatcodes/assets/cheatcodes.json index 7198f615288a..172c64f9321a 100644 --- a/crates/cheatcodes/assets/cheatcodes.json +++ b/crates/cheatcodes/assets/cheatcodes.json @@ -5653,6 +5653,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 6ab18d10935b..15b03b7d6efa 100644 --- a/crates/cheatcodes/spec/src/vm.rs +++ b/crates/cheatcodes/spec/src/vm.rs @@ -1393,6 +1393,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/fs.rs b/crates/cheatcodes/src/fs.rs index 3e345db94a97..0c289ec7862b 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::{ @@ -296,6 +297,20 @@ impl Cheatcode for tryFfiCall { } } +impl Cheatcode for promptCall { + fn apply(&self, _state: &mut Cheatcodes) -> Result { + let Self { promptText: text } = self; + prompt(text).map(|res| res.abi_encode()) + } +} + +impl Cheatcode for promptSecretCall { + fn apply(&self, _state: &mut Cheatcodes) -> Result { + let Self { promptText: text } = self; + prompt_secret(text).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 +385,19 @@ fn ffi(state: &Cheatcodes, input: &[String]) -> Result { }) } +fn prompt(prompt_text: &String) -> Result { + let input: String = + Input::new().allow_empty(true).with_prompt(prompt_text).interact_text().unwrap(); + + Ok(input) +} + +fn prompt_secret(prompt_text: &String) -> Result { + let input: String = Password::new().with_prompt(prompt_text).interact().unwrap(); + + Ok(input) +} + #[cfg(test)] mod tests { use super::*; diff --git a/testdata/cheats/Vm.sol b/testdata/cheats/Vm.sol index bac97febb2f9..27f96acd0fd8 100644 --- a/testdata/cheats/Vm.sol +++ b/testdata/cheats/Vm.sol @@ -280,6 +280,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);