Add promptSecretUint

[ZeroEkkusu]

Component

Forge

Describe the feature you would like

Problem

The combination broadcast + promptSecret leaks the private key because it requires parseUint, which doesn't obfuscate the value.

Example

Mnemonic:

vm.broadcast(vm.deriveKey(vm.promptSecret("Mnemonic"), 0));
revert();

├─ [0] VM::promptSecret("Mnemonic")
│   └─ ← [Return] <secret>
├─ [0] VM::deriveKey(<pk>) [staticcall]
│   └─ ← [Return] <pk>
├─ [0] VM::broadcast(<pk>)
│   └─ ← [Return] 
└─ ← [Revert] EvmError: Revert

Private key:

vm.broadcast(vm.parseUint(vm.promptSecret("Private key")));
revert();

├─ [0] VM::promptSecret("Private key")
│   └─ ← [Return] <secret>
├─ [0] VM::parseUint("0x0002a1b4984c17435bdd23c3e4d172af9217c4f63cac2edfd8044fd673459042") [staticcall]
│   └─ ← [Return] 4649744120560439147230634839216140256104967990126594206675684809117110338 [4.649e72]
├─ [0] VM::broadcast(<pk>)
│   └─ ← [Return] 
└─ ← [Revert] EvmError: Revert

Solution

Add promptSecretUint.

Additional context

No response

[kamuik16]

Hey @ZeroEkkusu! There exists a promptUint cheatcode, maybe that can help you?

/// Prompts the user for uint256 in the terminal.
#[cheatcode(group = Filesystem)]
function promptUint(string calldata promptText) external returns (uint256);

Then you wouldn't have to use parseUint.

[ZeroEkkusu]

Exactly, but it doesn't hide what's being typed in the terminal, like promptSecret does:

Private key: [hidden]

I don't think it would be too difficult to implement promptSecretUint.

[zerosnacks]

Relevant:

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

Lines 385 to 709 in 5494c33

	/// Custom decoding for cheatcode inputs.
	fn decode_cheatcode_inputs(&self, func: &Function, data: &[u8]) -> Option<Vec<String>> {
	match func.name.as_str() {
	"expectRevert" => Some(vec![self.revert_decoder.decode(data, None)]),
	"addr" | "createWallet" | "deriveKey" | "rememberKey" => {
	// Redact private key in all cases
	Some(vec!["<pk>".to_string()])
	}
	"broadcast" | "startBroadcast" => {
	// Redact private key if defined
	// broadcast(uint256) / startBroadcast(uint256)
	if !func.inputs.is_empty() && func.inputs[0].ty == "uint256" {
	Some(vec!["<pk>".to_string()])
	} else {
	None
	}
	}
	"getNonce" => {
	// Redact private key if defined
	// getNonce(Wallet)
	if !func.inputs.is_empty() && func.inputs[0].ty == "tuple" {
	Some(vec!["<pk>".to_string()])
	} else {
	None
	}
	}
	"sign" | "signP256" => {
	let mut decoded = func.abi_decode_input(&data[SELECTOR_LEN..], false).ok()?;
	// Redact private key and replace in trace
	// sign(uint256,bytes32) / signP256(uint256,bytes32) / sign(Wallet,bytes32)
	if !decoded.is_empty() &&
	(func.inputs[0].ty == "uint256" || func.inputs[0].ty == "tuple")
	{
	decoded[0] = DynSolValue::String("<pk>".to_string());
	}
	Some(decoded.iter().map(format_token).collect())
	}
	"parseJson" |
	"parseJsonUint" |
	"parseJsonUintArray" |
	"parseJsonInt" |
	"parseJsonIntArray" |
	"parseJsonString" |
	"parseJsonStringArray" |
	"parseJsonAddress" |
	"parseJsonAddressArray" |
	"parseJsonBool" |
	"parseJsonBoolArray" |
	"parseJsonBytes" |
	"parseJsonBytesArray" |
	"parseJsonBytes32" |
	"parseJsonBytes32Array" |
	"writeJson" |
	// `keyExists` is being deprecated in favor of `keyExistsJson`. It will be removed in future versions.
	"keyExists" |
	"keyExistsJson" |
	"serializeBool" |
	"serializeUint" |
	"serializeUintToHex" |
	"serializeInt" |
	"serializeAddress" |
	"serializeBytes32" |
	"serializeString" |
	"serializeBytes" => {
	if self.verbosity >= 5 {
	None
	} else {
	let mut decoded = func.abi_decode_input(&data[SELECTOR_LEN..], false).ok()?;
	let token = if func.name.as_str() == "parseJson" ||
	// `keyExists` is being deprecated in favor of `keyExistsJson`. It will be removed in future versions.
	func.name.as_str() == "keyExists" ||
	func.name.as_str() == "keyExistsJson"
	{
	"<JSON file>"
	} else {
	"<stringified JSON>"
	};
	decoded[0] = DynSolValue::String(token.to_string());
	Some(decoded.iter().map(format_token).collect())
	}
	}
	s if s.contains("Toml") => {
	if self.verbosity >= 5 {
	None
	} else {
	let mut decoded = func.abi_decode_input(&data[SELECTOR_LEN..], false).ok()?;
	let token = if func.name.as_str() == "parseToml" ||
	func.name.as_str() == "keyExistsToml"
	{
	"<TOML file>"
	} else {
	"<stringified TOML>"
	};
	decoded[0] = DynSolValue::String(token.to_string());
	Some(decoded.iter().map(format_token).collect())
	}
	}
	_ => None,
	}
	}
	/// Decodes a function's output into the given trace.
	fn decode_function_output(&self, trace: &CallTrace, funcs: &[Function]) -> Option<String> {
	let data = &trace.output;
	if trace.success {
	if trace.address == CHEATCODE_ADDRESS {
	if let Some(decoded) =
	funcs.iter().find_map(|func| self.decode_cheatcode_outputs(func))
	{
	return Some(decoded);
	}
	}
	if let Some(values) =
	funcs.iter().find_map(|func| func.abi_decode_output(data, false).ok())
	{
	// Functions coming from an external database do not have any outputs specified,
	// and will lead to returning an empty list of values.
	if values.is_empty() {
	return None;
	}
	return Some(
	values.iter().map(|value| self.apply_label(value)).format(", ").to_string(),
	);
	}
	None
	} else {
	Some(self.revert_decoder.decode(data, Some(trace.status)))
	}
	}
	/// Custom decoding for cheatcode outputs.
	fn decode_cheatcode_outputs(&self, func: &Function) -> Option<String> {
	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,
	}
	.map(Into::into)
	}
	/// Decodes an event.
	pub async fn decode_event<'a>(&self, log: &'a LogData) -> DecodedCallLog<'a> {
	let &[t0, ..] = log.topics() else { return DecodedCallLog::Raw(log) };
	let mut events = Vec::new();
	let events = match self.events.get(&(t0, log.topics().len() - 1)) {
	Some(es) => es,
	None => {
	if let Some(identifier) = &self.signature_identifier {
	if let Some(event) = identifier.write().await.identify_event(&t0[..]).await {
	events.push(get_indexed_event(event, log));
	}
	}
	&events
	}
	};
	for event in events {
	if let Ok(decoded) = event.decode_log(log, false) {
	let params = reconstruct_params(event, &decoded);
	return DecodedCallLog::Decoded(
	event.name.clone(),
	params
	.into_iter()
	.zip(event.inputs.iter())
	.map(|(param, input)| {
	// undo patched names
	let name = input.name.clone();
	(name, self.apply_label(&param))
	})
	.collect(),
	);
	}
	}
	DecodedCallLog::Raw(log)
	}
	/// Prefetches function and event signatures into the identifier cache
	pub async fn prefetch_signatures(&self, nodes: &[CallTraceNode]) {
	let Some(identifier) = &self.signature_identifier else { return };
	let events_it = nodes
	.iter()
	.flat_map(|node| node.logs.iter().filter_map(|log| log.topics().first()))
	.unique();
	identifier.write().await.identify_events(events_it).await;
	const DEFAULT_CREATE2_DEPLOYER_BYTES: [u8; 20] = DEFAULT_CREATE2_DEPLOYER.0 .0;
	let funcs_it = nodes
	.iter()
	.filter_map(|n| match n.trace.address.0 .0 {
	DEFAULT_CREATE2_DEPLOYER_BYTES => None,
	[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0x01..=0x0a] => None,
	_ => n.trace.data.get(..SELECTOR_LEN),
	})
	.filter(|v| !self.functions.contains_key(*v))
	.unique();
	identifier.write().await.identify_functions(funcs_it).await;
	}
	fn apply_label(&self, value: &DynSolValue) -> String {
	if let DynSolValue::Address(addr) = value {
	if let Some(label) = self.labels.get(addr) {
	return format!("{label}: [{addr}]");
	}
	}
	format_token(value)
	}
	}
	/// Restore the order of the params of a decoded event,
	/// as Alloy returns the indexed and unindexed params separately.
	fn reconstruct_params(event: &Event, decoded: &DecodedEvent) -> Vec<DynSolValue> {
	let mut indexed = 0;
	let mut unindexed = 0;
	let mut inputs = vec![];
	for input in event.inputs.iter() {
	if input.indexed {
	inputs.push(decoded.indexed[indexed].clone());
	indexed += 1;
	} else {
	inputs.push(decoded.body[unindexed].clone());
	unindexed += 1;
	}
	}
	inputs
	}
	fn indexed_inputs(event: &Event) -> usize {
	event.inputs.iter().filter(|param| param.indexed).count()
	}
	#[cfg(test)]
	mod tests {
	use super::*;
	use alloy_primitives::hex;
	#[test]
	fn test_should_redact_pk() {
	let decoder = CallTraceDecoder::new();
	// [function_signature, data, expected]
	let cheatcode_input_test_cases = vec![
	// Should redact private key from traces in all cases:
	("addr(uint256)", vec![], Some(vec!["<pk>".to_string()])),
	("createWallet(string)", vec![], Some(vec!["<pk>".to_string()])),
	("createWallet(uint256)", vec![], Some(vec!["<pk>".to_string()])),
	("deriveKey(string,uint32)", vec![], Some(vec!["<pk>".to_string()])),
	("deriveKey(string,string,uint32)", vec![], Some(vec!["<pk>".to_string()])),
	("deriveKey(string,uint32,string)", vec![], Some(vec!["<pk>".to_string()])),
	("deriveKey(string,string,uint32,string)", vec![], Some(vec!["<pk>".to_string()])),
	("rememberKey(uint256)", vec![], Some(vec!["<pk>".to_string()])),
	//
	// Should redact private key from traces in specific cases with exceptions:
	("broadcast(uint256)", vec![], Some(vec!["<pk>".to_string()])),
	("broadcast()", vec![], None), // Ignore: `private key` is not passed.
	("startBroadcast(uint256)", vec![], Some(vec!["<pk>".to_string()])),
	("startBroadcast()", vec![], None), // Ignore: `private key` is not passed.
	("getNonce((address,uint256,uint256,uint256))", vec![], Some(vec!["<pk>".to_string()])),
	("getNonce(address)", vec![], None), // Ignore: `address` is public.
	//
	// Should redact private key and replace in trace in cases:
	(
	"sign(uint256,bytes32)",
	hex!(
	"
	e341eaa4
	7c852118294e51e653712a81e05800f419141751be58f605c371e15141b007a6
	0000000000000000000000000000000000000000000000000000000000000000
	"
	)
	.to_vec(),
	Some(vec![
	"\"<pk>\"".to_string(),
	"0x0000000000000000000000000000000000000000000000000000000000000000"
	.to_string(),
	]),
	),
	(
	"signP256(uint256,bytes32)",
	hex!(
	"
	83211b40
	7c852118294e51e653712a81e05800f419141751be58f605c371e15141b007a6
	0000000000000000000000000000000000000000000000000000000000000000
	"
	)
	.to_vec(),
	Some(vec![
	"\"<pk>\"".to_string(),
	"0x0000000000000000000000000000000000000000000000000000000000000000"
	.to_string(),
	]),
	),
	];
	// [function_signature, expected]
	let cheatcode_output_test_cases = vec![
	// Should redact private key on output in all cases:
	("createWallet(string)", Some("<pk>".to_string())),
	("deriveKey(string,uint32)", Some("<pk>".to_string())),
	];
	for (function_signature, data, expected) in cheatcode_input_test_cases {
	let function = Function::parse(function_signature).unwrap();
	let result = decoder.decode_cheatcode_inputs(&function, &data);
	assert_eq!(result, expected, "Input case failed for: {}", function_signature);
	}
	for (function_signature, expected) in cheatcode_output_test_cases {
	let function = Function::parse(function_signature).unwrap();
	let result = Some(decoder.decode_cheatcode_outputs(&function).unwrap_or_default());
	assert_eq!(result, expected, "Output case failed for: {}", function_signature);
	}
	}
	}