Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 4 commits into from
Mar 21, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions crates/cheatcodes/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -43,3 +43,4 @@ walkdir = "2"
p256 = "0.13.2"
thiserror = "1"
rustc-hash.workspace = true
dialoguer = "0.11.0"
40 changes: 40 additions & 0 deletions crates/cheatcodes/assets/cheatcodes.json

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

10 changes: 10 additions & 0 deletions crates/cheatcodes/spec/src/vm.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
5 changes: 5 additions & 0 deletions crates/cheatcodes/src/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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(),
Expand Down Expand Up @@ -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("./"),
Expand Down
50 changes: 50 additions & 0 deletions crates/cheatcodes/src/fs.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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::*;
Expand Down
1 change: 1 addition & 0 deletions crates/config/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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'
Expand Down
3 changes: 3 additions & 0 deletions crates/config/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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),
Expand Down
1 change: 1 addition & 0 deletions crates/evm/traces/src/decoder/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
1 change: 1 addition & 0 deletions crates/forge/tests/cli/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand Down
3 changes: 3 additions & 0 deletions crates/forge/tests/it/test_helpers.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down
2 changes: 2 additions & 0 deletions testdata/cheats/Vm.sol

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

18 changes: 18 additions & 0 deletions testdata/default/cheats/Prompt.t.sol
Original file line number Diff line number Diff line change
@@ -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");
}
}
Loading