Skip to content

Commit

Permalink
Implement vm.prompt cheatcode
Browse files Browse the repository at this point in the history
  • Loading branch information
Tudmotu committed Mar 21, 2024
1 parent 63ea108 commit 8da53cf
Show file tree
Hide file tree
Showing 12 changed files with 134 additions and 0 deletions.
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 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
17 changes: 17 additions & 0 deletions testdata/cheats/Prompt.t.sol
Original file line number Diff line number Diff line change
@@ -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");
}
}
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.

0 comments on commit 8da53cf

Please sign in to comment.