diff --git a/crates/evm/src/executor/inspector/cheatcodes/ext.rs b/crates/evm/src/executor/inspector/cheatcodes/ext.rs index 95f88e4cd89f0..e2a6a5cab11a7 100644 --- a/crates/evm/src/executor/inspector/cheatcodes/ext.rs +++ b/crates/evm/src/executor/inspector/cheatcodes/ext.rs @@ -1,5 +1,5 @@ -use super::{bail, ensure, fmt_err, Cheatcodes, Result}; -use crate::{abi::HEVMCalls, executor::inspector::cheatcodes::util}; +use super::{bail, ensure, fmt_err, util::MAGIC_SKIP_BYTES, Cheatcodes, Error, Result}; +use crate::{abi::HEVMCalls, executor::inspector::cheatcodes::parse}; use ethers::{ abi::{self, AbiEncode, JsonAbi, ParamType, Token}, prelude::artifacts::CompactContractBytecode, @@ -7,6 +7,7 @@ use ethers::{ }; use foundry_common::{fmt::*, fs, get_artifact_path}; use foundry_config::fs_permissions::FsAccessKind; +use revm::{Database, EVMData}; use serde::Deserialize; use serde_json::Value; use std::{collections::BTreeMap, env, path::Path, process::Command}; @@ -190,9 +191,9 @@ fn get_env(key: &str, ty: ParamType, delim: Option<&str>, default: Option) -> Result { }; trace!(target : "forge::evm", ?values, "parsign values"); return if let Some(array) = values[0].as_array() { - util::parse_array(array.iter().map(to_string), &coercion_type) + parse::parse_array(array.iter().map(to_string), &coercion_type) } else { - util::parse(&to_string(values[0]), &coercion_type) + parse::parse(&to_string(values[0]), &coercion_type) } } @@ -482,7 +483,7 @@ fn key_exists(json_str: &str, key: &str) -> Result { let json: Value = serde_json::from_str(json_str).map_err(|e| format!("Could not convert to JSON: {e}"))?; let values = jsonpath_lib::select(&json, &canonicalize_json_key(key))?; - let exists = util::parse(&(!values.is_empty()).to_string(), &ParamType::Bool)?; + let exists = parse::parse(&(!values.is_empty()).to_string(), &ParamType::Bool)?; Ok(exists) } @@ -494,8 +495,28 @@ fn sleep(milliseconds: &U256) -> Result { Ok(Default::default()) } +/// Skip the current test, by returning a magic value that will be checked by the test runner. +pub fn skip(state: &mut Cheatcodes, depth: u64, skip: bool) -> Result { + if !skip { + return Ok(b"".into()) + } + + // Skip should not work if called deeper than at test level. + // As we're not returning the magic skip bytes, this will cause a test failure. + if depth > 1 { + return Err(Error::custom("The skip cheatcode can only be used at test level")) + } + + state.skip = true; + Err(Error::custom_bytes(MAGIC_SKIP_BYTES)) +} + #[instrument(level = "error", name = "ext", target = "evm::cheatcodes", skip_all)] -pub fn apply(state: &mut Cheatcodes, call: &HEVMCalls) -> Option { +pub fn apply( + state: &mut Cheatcodes, + data: &mut EVMData<'_, DB>, + call: &HEVMCalls, +) -> Option { Some(match call { HEVMCalls::Ffi(inner) => { if state.config.ffi { @@ -680,6 +701,7 @@ pub fn apply(state: &mut Cheatcodes, call: &HEVMCalls) -> Option { HEVMCalls::WriteJson0(inner) => write_json(state, &inner.0, &inner.1, None), HEVMCalls::WriteJson1(inner) => write_json(state, &inner.0, &inner.1, Some(&inner.2)), HEVMCalls::KeyExists(inner) => key_exists(&inner.0, &inner.1), + HEVMCalls::Skip(inner) => skip(state, data.journaled_state.depth(), inner.0), _ => return None, }) } diff --git a/crates/evm/src/executor/inspector/cheatcodes/mod.rs b/crates/evm/src/executor/inspector/cheatcodes/mod.rs index ebd4d9c0b3107..0e048399d4201 100644 --- a/crates/evm/src/executor/inspector/cheatcodes/mod.rs +++ b/crates/evm/src/executor/inspector/cheatcodes/mod.rs @@ -57,10 +57,15 @@ mod fs; mod fuzz; /// Mapping related cheatcodes mod mapping; +/// Parsing related cheatcodes. +/// Does not include JSON-related cheatcodes to cut complexity. +mod parse; /// Snapshot related cheatcodes mod snapshot; -/// Utility cheatcodes (`sign` etc.) +/// Utility functions and constants. pub mod util; +/// Wallet / key management related cheatcodes +mod wallet; pub use util::{BroadcastableTransaction, DEFAULT_CREATE2_DEPLOYER}; mod config; @@ -219,10 +224,11 @@ impl Cheatcodes { let opt = env::apply(self, data, caller, &decoded) .transpose() - .or_else(|| util::apply(self, data, &decoded)) + .or_else(|| wallet::apply(self, data, &decoded)) + .or_else(|| parse::apply(self, data, &decoded)) .or_else(|| expect::apply(self, data, &decoded)) .or_else(|| fuzz::apply(&decoded)) - .or_else(|| ext::apply(self, &decoded)) + .or_else(|| ext::apply(self, data, &decoded)) .or_else(|| fs::apply(self, &decoded)) .or_else(|| snapshot::apply(data, &decoded)) .or_else(|| fork::apply(self, data, &decoded)); diff --git a/crates/evm/src/executor/inspector/cheatcodes/parse.rs b/crates/evm/src/executor/inspector/cheatcodes/parse.rs new file mode 100644 index 0000000000000..904a3d502574e --- /dev/null +++ b/crates/evm/src/executor/inspector/cheatcodes/parse.rs @@ -0,0 +1,151 @@ +use super::{fmt_err, Cheatcodes, Result}; +use crate::abi::HEVMCalls; +use ethers::{ + abi::{ParamType, Token}, + prelude::*, +}; +use foundry_macros::UIfmt; +use revm::{Database, EVMData}; + +pub fn parse(s: &str, ty: &ParamType) -> Result { + parse_token(s, ty) + .map(|token| abi::encode(&[token]).into()) + .map_err(|e| fmt_err!("Failed to parse `{s}` as type `{ty}`: {e}")) +} + +pub fn parse_array(values: I, ty: &ParamType) -> Result +where + I: IntoIterator, + T: AsRef, +{ + let mut values = values.into_iter(); + match values.next() { + Some(first) if !first.as_ref().is_empty() => { + let tokens = std::iter::once(first) + .chain(values) + .map(|v| parse_token(v.as_ref(), ty)) + .collect::, _>>()?; + Ok(abi::encode(&[Token::Array(tokens)]).into()) + } + // return the empty encoded Bytes when values is empty or the first element is empty + _ => Ok(abi::encode(&[Token::String(String::new())]).into()), + } +} + +fn parse_token(s: &str, ty: &ParamType) -> Result { + match ty { + ParamType::Bool => { + s.to_ascii_lowercase().parse().map(Token::Bool).map_err(|e| e.to_string()) + } + ParamType::Uint(256) => parse_uint(s).map(Token::Uint), + ParamType::Int(256) => parse_int(s).map(Token::Int), + ParamType::Address => s.parse().map(Token::Address).map_err(|e| e.to_string()), + ParamType::FixedBytes(32) => parse_bytes(s).map(Token::FixedBytes), + ParamType::Bytes => parse_bytes(s).map(Token::Bytes), + ParamType::String => Ok(Token::String(s.to_string())), + _ => Err("unsupported type".into()), + } +} + +fn parse_int(s: &str) -> Result { + // hex string may start with "0x", "+0x", or "-0x" which needs to be stripped for + // `I256::from_hex_str` + if s.starts_with("0x") || s.starts_with("+0x") || s.starts_with("-0x") { + s.replacen("0x", "", 1).parse::().map_err(|err| err.to_string()) + } else { + match I256::from_dec_str(s) { + Ok(val) => Ok(val), + Err(dec_err) => s.parse::().map_err(|hex_err| { + format!("could not parse value as decimal or hex: {dec_err}, {hex_err}") + }), + } + } + .map(|v| v.into_raw()) +} + +fn parse_uint(s: &str) -> Result { + if s.starts_with("0x") { + s.parse::().map_err(|err| err.to_string()) + } else { + match U256::from_dec_str(s) { + Ok(val) => Ok(val), + Err(dec_err) => s.parse::().map_err(|hex_err| { + format!("could not parse value as decimal or hex: {dec_err}, {hex_err}") + }), + } + } +} + +fn parse_bytes(s: &str) -> Result, String> { + hex::decode(s).map_err(|e| e.to_string()) +} + +#[instrument(level = "error", name = "util", target = "evm::cheatcodes", skip_all)] +pub fn apply( + _state: &mut Cheatcodes, + _data: &mut EVMData<'_, DB>, + call: &HEVMCalls, +) -> Option { + Some(match call { + HEVMCalls::ToString0(inner) => { + Ok(ethers::abi::encode(&[Token::String(inner.0.pretty())]).into()) + } + HEVMCalls::ToString1(inner) => { + Ok(ethers::abi::encode(&[Token::String(inner.0.pretty())]).into()) + } + HEVMCalls::ToString2(inner) => { + Ok(ethers::abi::encode(&[Token::String(inner.0.pretty())]).into()) + } + HEVMCalls::ToString3(inner) => { + Ok(ethers::abi::encode(&[Token::String(inner.0.pretty())]).into()) + } + HEVMCalls::ToString4(inner) => { + Ok(ethers::abi::encode(&[Token::String(inner.0.pretty())]).into()) + } + HEVMCalls::ToString5(inner) => { + Ok(ethers::abi::encode(&[Token::String(inner.0.pretty())]).into()) + } + HEVMCalls::ParseBytes(inner) => parse(&inner.0, &ParamType::Bytes), + HEVMCalls::ParseAddress(inner) => parse(&inner.0, &ParamType::Address), + HEVMCalls::ParseUint(inner) => parse(&inner.0, &ParamType::Uint(256)), + HEVMCalls::ParseInt(inner) => parse(&inner.0, &ParamType::Int(256)), + HEVMCalls::ParseBytes32(inner) => parse(&inner.0, &ParamType::FixedBytes(32)), + HEVMCalls::ParseBool(inner) => parse(&inner.0, &ParamType::Bool), + _ => return None, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use ethers::abi::AbiDecode; + + #[test] + fn test_uint_env() { + let pk = "0x10532cc9d0d992825c3f709c62c969748e317a549634fb2a9fa949326022e81f"; + let val: U256 = pk.parse().unwrap(); + let parsed = parse(pk, &ParamType::Uint(256)).unwrap(); + let decoded = U256::decode(&parsed).unwrap(); + assert_eq!(val, decoded); + + let parsed = parse(pk.strip_prefix("0x").unwrap(), &ParamType::Uint(256)).unwrap(); + let decoded = U256::decode(&parsed).unwrap(); + assert_eq!(val, decoded); + + let parsed = parse("1337", &ParamType::Uint(256)).unwrap(); + let decoded = U256::decode(&parsed).unwrap(); + assert_eq!(U256::from(1337u64), decoded); + } + + #[test] + fn test_int_env() { + let val = U256::from(100u64); + let parsed = parse(&val.to_string(), &ParamType::Int(256)).unwrap(); + let decoded = I256::decode(parsed).unwrap(); + assert_eq!(val, decoded.try_into().unwrap()); + + let parsed = parse("100", &ParamType::Int(256)).unwrap(); + let decoded = I256::decode(parsed).unwrap(); + assert_eq!(U256::from(100u64), decoded.try_into().unwrap()); + } +} diff --git a/crates/evm/src/executor/inspector/cheatcodes/util.rs b/crates/evm/src/executor/inspector/cheatcodes/util.rs index 3989421a70e09..886948e7e7890 100644 --- a/crates/evm/src/executor/inspector/cheatcodes/util.rs +++ b/crates/evm/src/executor/inspector/cheatcodes/util.rs @@ -1,6 +1,5 @@ -use super::{ensure, fmt_err, Cheatcodes, Error, Result}; +use super::{ensure, Result}; use crate::{ - abi::HEVMCalls, executor::backend::{ error::{DatabaseError, DatabaseResult}, DatabaseExt, @@ -9,43 +8,29 @@ use crate::{ }; use bytes::{BufMut, Bytes, BytesMut}; use ethers::{ - abi::{AbiEncode, Address, ParamType, Token}, + abi::Address, core::k256::elliptic_curve::Curve, prelude::{ - k256::{ - ecdsa::SigningKey, - elliptic_curve::{bigint::Encoding, sec1::ToEncodedPoint}, - Secp256k1, - }, - LocalWallet, Signer, H160, *, + k256::{ecdsa::SigningKey, elliptic_curve::bigint::Encoding, Secp256k1}, + H160, *, }, - signers::{ - coins_bip39::{ - ChineseSimplified, ChineseTraditional, Czech, English, French, Italian, Japanese, - Korean, Portuguese, Spanish, Wordlist, - }, - MnemonicBuilder, - }, - types::{transaction::eip2718::TypedTransaction, NameOrAddress, H256, U256}, - utils::{self, keccak256}, + types::{transaction::eip2718::TypedTransaction, NameOrAddress, U256}, }; -use foundry_common::{fmt::*, RpcUrl}; +use foundry_common::RpcUrl; use revm::{ interpreter::CreateInputs, primitives::{Account, TransactTo}, Database, EVMData, JournaledState, }; -use std::{collections::VecDeque, str::FromStr}; +use std::collections::VecDeque; -const DEFAULT_DERIVATION_PATH_PREFIX: &str = "m/44'/60'/0'/0/"; +pub const MAGIC_SKIP_BYTES: &[u8] = b"FOUNDRY::SKIP"; /// Address of the default CREATE2 deployer 0x4e59b44847b379578588920ca78fbf26c0b4956c pub const DEFAULT_CREATE2_DEPLOYER: H160 = H160([ 78, 89, 180, 72, 71, 179, 121, 87, 133, 136, 146, 12, 167, 143, 191, 38, 192, 180, 149, 108, ]); -pub const MAGIC_SKIP_BYTES: &[u8] = b"FOUNDRY::SKIP"; - /// Helps collecting transactions from different forks. #[derive(Debug, Clone, Default)] pub struct BroadcastableTransaction { @@ -100,223 +85,6 @@ where Ok(f(account)) } -fn addr(private_key: U256) -> Result { - let key = parse_private_key(private_key)?; - let addr = utils::secret_key_to_address(&key); - Ok(addr.encode().into()) -} - -fn sign(private_key: U256, digest: H256, chain_id: U256) -> Result { - let key = parse_private_key(private_key)?; - let wallet = LocalWallet::from(key).with_chain_id(chain_id.as_u64()); - - // The `ecrecover` precompile does not use EIP-155 - let sig = wallet.sign_hash(digest)?; - let recovered = sig.recover(digest)?; - - assert_eq!(recovered, wallet.address()); - - let mut r_bytes = [0u8; 32]; - let mut s_bytes = [0u8; 32]; - sig.r.to_big_endian(&mut r_bytes); - sig.s.to_big_endian(&mut s_bytes); - - Ok((sig.v, r_bytes, s_bytes).encode().into()) -} - -/// Using a given private key, return its public ETH address, its public key affine x and y -/// coodinates, and its private key (see the 'Wallet' struct) -/// -/// If 'label' is set to 'Some()', assign that label to the associated ETH address in state -fn create_wallet(private_key: U256, label: Option, state: &mut Cheatcodes) -> Result { - let key = parse_private_key(private_key)?; - let addr = utils::secret_key_to_address(&key); - - let pub_key = key.verifying_key().as_affine().to_encoded_point(false); - let pub_key_x = pub_key.x().ok_or("No x coordinate was found")?; - let pub_key_y = pub_key.y().ok_or("No y coordinate was found")?; - - let pub_key_x = U256::from(pub_key_x.as_slice()); - let pub_key_y = U256::from(pub_key_y.as_slice()); - - if let Some(label) = label { - state.labels.insert(addr, label); - } - - Ok((addr, pub_key_x, pub_key_y, private_key).encode().into()) -} - -enum WordlistLang { - ChineseSimplified, - ChineseTraditional, - Czech, - English, - French, - Italian, - Japanese, - Korean, - Portuguese, - Spanish, -} - -impl FromStr for WordlistLang { - type Err = String; - - fn from_str(input: &str) -> Result { - match input { - "chinese_simplified" => Ok(Self::ChineseSimplified), - "chinese_traditional" => Ok(Self::ChineseTraditional), - "czech" => Ok(Self::Czech), - "english" => Ok(Self::English), - "french" => Ok(Self::French), - "italian" => Ok(Self::Italian), - "japanese" => Ok(Self::Japanese), - "korean" => Ok(Self::Korean), - "portuguese" => Ok(Self::Portuguese), - "spanish" => Ok(Self::Spanish), - _ => Err(format!("the language `{}` has no wordlist", input)), - } - } -} - -fn derive_key(mnemonic: &str, path: &str, index: u32) -> Result { - let derivation_path = - if path.ends_with('/') { format!("{path}{index}") } else { format!("{path}/{index}") }; - - let wallet = MnemonicBuilder::::default() - .phrase(mnemonic) - .derivation_path(&derivation_path)? - .build()?; - - let private_key = U256::from_big_endian(wallet.signer().to_bytes().as_slice()); - - Ok(private_key.encode().into()) -} - -fn derive_key_with_wordlist(mnemonic: &str, path: &str, index: u32, lang: &str) -> Result { - let lang = WordlistLang::from_str(lang)?; - match lang { - WordlistLang::ChineseSimplified => derive_key::(mnemonic, path, index), - WordlistLang::ChineseTraditional => derive_key::(mnemonic, path, index), - WordlistLang::Czech => derive_key::(mnemonic, path, index), - WordlistLang::English => derive_key::(mnemonic, path, index), - WordlistLang::French => derive_key::(mnemonic, path, index), - WordlistLang::Italian => derive_key::(mnemonic, path, index), - WordlistLang::Japanese => derive_key::(mnemonic, path, index), - WordlistLang::Korean => derive_key::(mnemonic, path, index), - WordlistLang::Portuguese => derive_key::(mnemonic, path, index), - WordlistLang::Spanish => derive_key::(mnemonic, path, index), - } -} - -fn remember_key(state: &mut Cheatcodes, private_key: U256, chain_id: U256) -> Result { - let key = parse_private_key(private_key)?; - let wallet = LocalWallet::from(key).with_chain_id(chain_id.as_u64()); - let address = wallet.address(); - - state.script_wallets.push(wallet); - - Ok(address.encode().into()) -} - -pub fn parse(s: &str, ty: &ParamType) -> Result { - parse_token(s, ty) - .map(|token| abi::encode(&[token]).into()) - .map_err(|e| fmt_err!("Failed to parse `{s}` as type `{ty}`: {e}")) -} - -pub fn skip(state: &mut Cheatcodes, depth: u64, skip: bool) -> Result { - if !skip { - return Ok(b"".into()) - } - - // Skip should not work if called deeper than at test level. - // As we're not returning the magic skip bytes, this will cause a test failure. - if depth > 1 { - return Err(Error::custom("The skip cheatcode can only be used at test level")) - } - - state.skip = true; - Err(Error::custom_bytes(MAGIC_SKIP_BYTES)) -} - -#[instrument(level = "error", name = "util", target = "evm::cheatcodes", skip_all)] -pub fn apply( - state: &mut Cheatcodes, - data: &mut EVMData<'_, DB>, - call: &HEVMCalls, -) -> Option { - Some(match call { - HEVMCalls::Addr(inner) => addr(inner.0), - // [function sign(uint256,bytes32)] Used to sign bytes32 digests using the given private key - HEVMCalls::Sign0(inner) => sign(inner.0, inner.1.into(), data.env.cfg.chain_id.into()), - // [function createWallet(string)] Used to derive private key and label the wallet with the - // same string - HEVMCalls::CreateWallet0(inner) => { - create_wallet(U256::from(keccak256(&inner.0)), Some(inner.0.clone()), state) - } - // [function createWallet(uint256)] creates a new wallet with the given private key - HEVMCalls::CreateWallet1(inner) => create_wallet(inner.0, None, state), - // [function createWallet(uint256,string)] creates a new wallet with the given private key - // and labels it with the given string - HEVMCalls::CreateWallet2(inner) => create_wallet(inner.0, Some(inner.1.clone()), state), - // [function sign(uint256,bytes32)] Used to sign bytes32 digests using the given Wallet's - // private key - HEVMCalls::Sign1(inner) => { - sign(inner.0.private_key, inner.1.into(), data.env.cfg.chain_id.into()) - } - HEVMCalls::DeriveKey0(inner) => { - derive_key::(&inner.0, DEFAULT_DERIVATION_PATH_PREFIX, inner.1) - } - HEVMCalls::DeriveKey1(inner) => derive_key::(&inner.0, &inner.1, inner.2), - HEVMCalls::DeriveKey2(inner) => { - derive_key_with_wordlist(&inner.0, DEFAULT_DERIVATION_PATH_PREFIX, inner.1, &inner.2) - } - HEVMCalls::DeriveKey3(inner) => { - derive_key_with_wordlist(&inner.0, &inner.1, inner.2, &inner.3) - } - HEVMCalls::RememberKey(inner) => remember_key(state, inner.0, data.env.cfg.chain_id.into()), - HEVMCalls::Label(inner) => { - state.labels.insert(inner.0, inner.1.clone()); - Ok(Default::default()) - } - HEVMCalls::GetLabel(inner) => { - let label = state - .labels - .get(&inner.0) - .cloned() - .unwrap_or_else(|| format!("unlabeled:{:?}", inner.0)); - Ok(ethers::abi::encode(&[Token::String(label)]).into()) - } - HEVMCalls::ToString0(inner) => { - Ok(ethers::abi::encode(&[Token::String(inner.0.pretty())]).into()) - } - HEVMCalls::ToString1(inner) => { - Ok(ethers::abi::encode(&[Token::String(inner.0.pretty())]).into()) - } - HEVMCalls::ToString2(inner) => { - Ok(ethers::abi::encode(&[Token::String(inner.0.pretty())]).into()) - } - HEVMCalls::ToString3(inner) => { - Ok(ethers::abi::encode(&[Token::String(inner.0.pretty())]).into()) - } - HEVMCalls::ToString4(inner) => { - Ok(ethers::abi::encode(&[Token::String(inner.0.pretty())]).into()) - } - HEVMCalls::ToString5(inner) => { - Ok(ethers::abi::encode(&[Token::String(inner.0.pretty())]).into()) - } - HEVMCalls::ParseBytes(inner) => parse(&inner.0, &ParamType::Bytes), - HEVMCalls::ParseAddress(inner) => parse(&inner.0, &ParamType::Address), - HEVMCalls::ParseUint(inner) => parse(&inner.0, &ParamType::Uint(256)), - HEVMCalls::ParseInt(inner) => parse(&inner.0, &ParamType::Int(256)), - HEVMCalls::ParseBytes32(inner) => parse(&inner.0, &ParamType::FixedBytes(32)), - HEVMCalls::ParseBool(inner) => parse(&inner.0, &ParamType::Bool), - HEVMCalls::Skip(inner) => skip(state, data.journaled_state.depth(), inner.0), - _ => return None, - }) -} - pub fn process_create( broadcast_sender: Address, bytecode: Bytes, @@ -375,73 +143,6 @@ where } } -pub fn parse_array(values: I, ty: &ParamType) -> Result -where - I: IntoIterator, - T: AsRef, -{ - let mut values = values.into_iter(); - match values.next() { - Some(first) if !first.as_ref().is_empty() => { - let tokens = std::iter::once(first) - .chain(values) - .map(|v| parse_token(v.as_ref(), ty)) - .collect::, _>>()?; - Ok(abi::encode(&[Token::Array(tokens)]).into()) - } - // return the empty encoded Bytes when values is empty or the first element is empty - _ => Ok(abi::encode(&[Token::String(String::new())]).into()), - } -} - -fn parse_token(s: &str, ty: &ParamType) -> Result { - match ty { - ParamType::Bool => { - s.to_ascii_lowercase().parse().map(Token::Bool).map_err(|e| e.to_string()) - } - ParamType::Uint(256) => parse_uint(s).map(Token::Uint), - ParamType::Int(256) => parse_int(s).map(Token::Int), - ParamType::Address => s.parse().map(Token::Address).map_err(|e| e.to_string()), - ParamType::FixedBytes(32) => parse_bytes(s).map(Token::FixedBytes), - ParamType::Bytes => parse_bytes(s).map(Token::Bytes), - ParamType::String => Ok(Token::String(s.to_string())), - _ => Err("unsupported type".into()), - } -} - -fn parse_int(s: &str) -> Result { - // hex string may start with "0x", "+0x", or "-0x" which needs to be stripped for - // `I256::from_hex_str` - if s.starts_with("0x") || s.starts_with("+0x") || s.starts_with("-0x") { - s.replacen("0x", "", 1).parse::().map_err(|err| err.to_string()) - } else { - match I256::from_dec_str(s) { - Ok(val) => Ok(val), - Err(dec_err) => s.parse::().map_err(|hex_err| { - format!("could not parse value as decimal or hex: {dec_err}, {hex_err}") - }), - } - } - .map(|v| v.into_raw()) -} - -fn parse_uint(s: &str) -> Result { - if s.starts_with("0x") { - s.parse::().map_err(|err| err.to_string()) - } else { - match U256::from_dec_str(s) { - Ok(val) => Ok(val), - Err(dec_err) => s.parse::().map_err(|hex_err| { - format!("could not parse value as decimal or hex: {dec_err}, {hex_err}") - }), - } - } -} - -fn parse_bytes(s: &str) -> Result, String> { - hex::decode(s).map_err(|e| e.to_string()) -} - pub fn parse_private_key(private_key: U256) -> Result { ensure!(!private_key.is_zero(), "Private key cannot be 0."); ensure!( @@ -475,38 +176,3 @@ pub fn check_if_fixed_gas_limit( pub fn is_potential_precompile(address: H160) -> bool { address < H160::from_low_u64_be(10) && address != H160::zero() } - -#[cfg(test)] -mod tests { - use super::*; - use ethers::abi::AbiDecode; - - #[test] - fn test_uint_env() { - let pk = "0x10532cc9d0d992825c3f709c62c969748e317a549634fb2a9fa949326022e81f"; - let val: U256 = pk.parse().unwrap(); - let parsed = parse(pk, &ParamType::Uint(256)).unwrap(); - let decoded = U256::decode(&parsed).unwrap(); - assert_eq!(val, decoded); - - let parsed = parse(pk.strip_prefix("0x").unwrap(), &ParamType::Uint(256)).unwrap(); - let decoded = U256::decode(&parsed).unwrap(); - assert_eq!(val, decoded); - - let parsed = parse("1337", &ParamType::Uint(256)).unwrap(); - let decoded = U256::decode(&parsed).unwrap(); - assert_eq!(U256::from(1337u64), decoded); - } - - #[test] - fn test_int_env() { - let val = U256::from(100u64); - let parsed = parse(&val.to_string(), &ParamType::Int(256)).unwrap(); - let decoded = I256::decode(parsed).unwrap(); - assert_eq!(val, decoded.try_into().unwrap()); - - let parsed = parse("100", &ParamType::Int(256)).unwrap(); - let decoded = I256::decode(parsed).unwrap(); - assert_eq!(U256::from(100u64), decoded.try_into().unwrap()); - } -} diff --git a/crates/evm/src/executor/inspector/cheatcodes/wallet.rs b/crates/evm/src/executor/inspector/cheatcodes/wallet.rs new file mode 100644 index 0000000000000..3c1d55048aa90 --- /dev/null +++ b/crates/evm/src/executor/inspector/cheatcodes/wallet.rs @@ -0,0 +1,211 @@ +use super::{ensure, Cheatcodes, Result}; +use crate::abi::HEVMCalls; +use ethers::{ + abi::{AbiEncode, Token}, + core::k256::elliptic_curve::Curve, + prelude::{ + k256::{ + ecdsa::SigningKey, + elliptic_curve::{bigint::Encoding, sec1::ToEncodedPoint}, + Secp256k1, + }, + LocalWallet, Signer, + }, + signers::{ + coins_bip39::{ + ChineseSimplified, ChineseTraditional, Czech, English, French, Italian, Japanese, + Korean, Portuguese, Spanish, Wordlist, + }, + MnemonicBuilder, + }, + types::{H256, U256}, + utils::{self, keccak256}, +}; +use revm::{Database, EVMData}; +use std::str::FromStr; + +/// The BIP32 default derivation path prefix. +const DEFAULT_DERIVATION_PATH_PREFIX: &str = "m/44'/60'/0'/0/"; + +pub fn parse_private_key(private_key: U256) -> Result { + ensure!(!private_key.is_zero(), "Private key cannot be 0."); + ensure!( + private_key < U256::from_big_endian(&Secp256k1::ORDER.to_be_bytes()), + "Private key must be less than the secp256k1 curve order \ + (115792089237316195423570985008687907852837564279074904382605163141518161494337).", + ); + let mut bytes: [u8; 32] = [0; 32]; + private_key.to_big_endian(&mut bytes); + SigningKey::from_bytes((&bytes).into()).map_err(Into::into) +} + +fn addr(private_key: U256) -> Result { + let key = parse_private_key(private_key)?; + let addr = utils::secret_key_to_address(&key); + Ok(addr.encode().into()) +} + +fn sign(private_key: U256, digest: H256, chain_id: U256) -> Result { + let key = parse_private_key(private_key)?; + let wallet = LocalWallet::from(key).with_chain_id(chain_id.as_u64()); + + // The `ecrecover` precompile does not use EIP-155 + let sig = wallet.sign_hash(digest)?; + let recovered = sig.recover(digest)?; + + assert_eq!(recovered, wallet.address()); + + let mut r_bytes = [0u8; 32]; + let mut s_bytes = [0u8; 32]; + sig.r.to_big_endian(&mut r_bytes); + sig.s.to_big_endian(&mut s_bytes); + + Ok((sig.v, r_bytes, s_bytes).encode().into()) +} + +/// Using a given private key, return its public ETH address, its public key affine x and y +/// coodinates, and its private key (see the 'Wallet' struct) +/// +/// If 'label' is set to 'Some()', assign that label to the associated ETH address in state +fn create_wallet(private_key: U256, label: Option, state: &mut Cheatcodes) -> Result { + let key = parse_private_key(private_key)?; + let addr = utils::secret_key_to_address(&key); + + let pub_key = key.verifying_key().as_affine().to_encoded_point(false); + let pub_key_x = pub_key.x().ok_or("No x coordinate was found")?; + let pub_key_y = pub_key.y().ok_or("No y coordinate was found")?; + + let pub_key_x = U256::from(pub_key_x.as_slice()); + let pub_key_y = U256::from(pub_key_y.as_slice()); + + if let Some(label) = label { + state.labels.insert(addr, label); + } + + Ok((addr, pub_key_x, pub_key_y, private_key).encode().into()) +} + +enum WordlistLang { + ChineseSimplified, + ChineseTraditional, + Czech, + English, + French, + Italian, + Japanese, + Korean, + Portuguese, + Spanish, +} + +impl FromStr for WordlistLang { + type Err = String; + + fn from_str(input: &str) -> Result { + match input { + "chinese_simplified" => Ok(Self::ChineseSimplified), + "chinese_traditional" => Ok(Self::ChineseTraditional), + "czech" => Ok(Self::Czech), + "english" => Ok(Self::English), + "french" => Ok(Self::French), + "italian" => Ok(Self::Italian), + "japanese" => Ok(Self::Japanese), + "korean" => Ok(Self::Korean), + "portuguese" => Ok(Self::Portuguese), + "spanish" => Ok(Self::Spanish), + _ => Err(format!("the language `{}` has no wordlist", input)), + } + } +} + +fn derive_key(mnemonic: &str, path: &str, index: u32) -> Result { + let derivation_path = + if path.ends_with('/') { format!("{path}{index}") } else { format!("{path}/{index}") }; + + let wallet = MnemonicBuilder::::default() + .phrase(mnemonic) + .derivation_path(&derivation_path)? + .build()?; + + let private_key = U256::from_big_endian(wallet.signer().to_bytes().as_slice()); + + Ok(private_key.encode().into()) +} + +fn derive_key_with_wordlist(mnemonic: &str, path: &str, index: u32, lang: &str) -> Result { + let lang = WordlistLang::from_str(lang)?; + match lang { + WordlistLang::ChineseSimplified => derive_key::(mnemonic, path, index), + WordlistLang::ChineseTraditional => derive_key::(mnemonic, path, index), + WordlistLang::Czech => derive_key::(mnemonic, path, index), + WordlistLang::English => derive_key::(mnemonic, path, index), + WordlistLang::French => derive_key::(mnemonic, path, index), + WordlistLang::Italian => derive_key::(mnemonic, path, index), + WordlistLang::Japanese => derive_key::(mnemonic, path, index), + WordlistLang::Korean => derive_key::(mnemonic, path, index), + WordlistLang::Portuguese => derive_key::(mnemonic, path, index), + WordlistLang::Spanish => derive_key::(mnemonic, path, index), + } +} + +fn remember_key(state: &mut Cheatcodes, private_key: U256, chain_id: U256) -> Result { + let key = parse_private_key(private_key)?; + let wallet = LocalWallet::from(key).with_chain_id(chain_id.as_u64()); + let address = wallet.address(); + + state.script_wallets.push(wallet); + + Ok(address.encode().into()) +} + +#[instrument(level = "error", name = "util", target = "evm::cheatcodes", skip_all)] +pub fn apply( + state: &mut Cheatcodes, + data: &mut EVMData<'_, DB>, + call: &HEVMCalls, +) -> Option { + Some(match call { + HEVMCalls::Addr(inner) => addr(inner.0), + // [function sign(uint256,bytes32)] Used to sign bytes32 digests using the given private key + HEVMCalls::Sign0(inner) => sign(inner.0, inner.1.into(), data.env.cfg.chain_id.into()), + // [function createWallet(string)] Used to derive private key and label the wallet with the + // same string + HEVMCalls::CreateWallet0(inner) => { + create_wallet(U256::from(keccak256(&inner.0)), Some(inner.0.clone()), state) + } + // [function createWallet(uint256)] creates a new wallet with the given private key + HEVMCalls::CreateWallet1(inner) => create_wallet(inner.0, None, state), + // [function createWallet(uint256,string)] creates a new wallet with the given private key + // and labels it with the given string + HEVMCalls::CreateWallet2(inner) => create_wallet(inner.0, Some(inner.1.clone()), state), + // [function sign(uint256,bytes32)] Used to sign bytes32 digests using the given Wallet's + // private key + HEVMCalls::Sign1(inner) => { + sign(inner.0.private_key, inner.1.into(), data.env.cfg.chain_id.into()) + } + HEVMCalls::DeriveKey0(inner) => { + derive_key::(&inner.0, DEFAULT_DERIVATION_PATH_PREFIX, inner.1) + } + HEVMCalls::DeriveKey1(inner) => derive_key::(&inner.0, &inner.1, inner.2), + HEVMCalls::DeriveKey2(inner) => { + derive_key_with_wordlist(&inner.0, DEFAULT_DERIVATION_PATH_PREFIX, inner.1, &inner.2) + } + HEVMCalls::DeriveKey3(inner) => { + derive_key_with_wordlist(&inner.0, &inner.1, inner.2, &inner.3) + } + HEVMCalls::RememberKey(inner) => remember_key(state, inner.0, data.env.cfg.chain_id.into()), + HEVMCalls::Label(inner) => { + state.labels.insert(inner.0, inner.1.clone()); + Ok(Default::default()) + } + HEVMCalls::GetLabel(inner) => { + let label = state + .labels + .get(&inner.0) + .cloned() + .unwrap_or_else(|| format!("unlabeled:{:?}", inner.0)); + Ok(ethers::abi::encode(&[Token::String(label)]).into()) + } + _ => return None, + }) +}