|
| 1 | +use super::{fmt_err, Cheatcodes, Result}; |
| 2 | +use crate::abi::HEVMCalls; |
| 3 | +use ethers::{ |
| 4 | + abi::{ParamType, Token}, |
| 5 | + prelude::*, |
| 6 | +}; |
| 7 | +use foundry_macros::UIfmt; |
| 8 | +use revm::{Database, EVMData}; |
| 9 | + |
| 10 | +pub fn parse(s: &str, ty: &ParamType) -> Result { |
| 11 | + parse_token(s, ty) |
| 12 | + .map(|token| abi::encode(&[token]).into()) |
| 13 | + .map_err(|e| fmt_err!("Failed to parse `{s}` as type `{ty}`: {e}")) |
| 14 | +} |
| 15 | + |
| 16 | +pub fn parse_array<I, T>(values: I, ty: &ParamType) -> Result |
| 17 | +where |
| 18 | + I: IntoIterator<Item = T>, |
| 19 | + T: AsRef<str>, |
| 20 | +{ |
| 21 | + let mut values = values.into_iter(); |
| 22 | + match values.next() { |
| 23 | + Some(first) if !first.as_ref().is_empty() => { |
| 24 | + let tokens = std::iter::once(first) |
| 25 | + .chain(values) |
| 26 | + .map(|v| parse_token(v.as_ref(), ty)) |
| 27 | + .collect::<Result<Vec<_>, _>>()?; |
| 28 | + Ok(abi::encode(&[Token::Array(tokens)]).into()) |
| 29 | + } |
| 30 | + // return the empty encoded Bytes when values is empty or the first element is empty |
| 31 | + _ => Ok(abi::encode(&[Token::String(String::new())]).into()), |
| 32 | + } |
| 33 | +} |
| 34 | + |
| 35 | +fn parse_token(s: &str, ty: &ParamType) -> Result<Token, String> { |
| 36 | + match ty { |
| 37 | + ParamType::Bool => { |
| 38 | + s.to_ascii_lowercase().parse().map(Token::Bool).map_err(|e| e.to_string()) |
| 39 | + } |
| 40 | + ParamType::Uint(256) => parse_uint(s).map(Token::Uint), |
| 41 | + ParamType::Int(256) => parse_int(s).map(Token::Int), |
| 42 | + ParamType::Address => s.parse().map(Token::Address).map_err(|e| e.to_string()), |
| 43 | + ParamType::FixedBytes(32) => parse_bytes(s).map(Token::FixedBytes), |
| 44 | + ParamType::Bytes => parse_bytes(s).map(Token::Bytes), |
| 45 | + ParamType::String => Ok(Token::String(s.to_string())), |
| 46 | + _ => Err("unsupported type".into()), |
| 47 | + } |
| 48 | +} |
| 49 | + |
| 50 | +fn parse_int(s: &str) -> Result<U256, String> { |
| 51 | + // hex string may start with "0x", "+0x", or "-0x" which needs to be stripped for |
| 52 | + // `I256::from_hex_str` |
| 53 | + if s.starts_with("0x") || s.starts_with("+0x") || s.starts_with("-0x") { |
| 54 | + s.replacen("0x", "", 1).parse::<I256>().map_err(|err| err.to_string()) |
| 55 | + } else { |
| 56 | + match I256::from_dec_str(s) { |
| 57 | + Ok(val) => Ok(val), |
| 58 | + Err(dec_err) => s.parse::<I256>().map_err(|hex_err| { |
| 59 | + format!("could not parse value as decimal or hex: {dec_err}, {hex_err}") |
| 60 | + }), |
| 61 | + } |
| 62 | + } |
| 63 | + .map(|v| v.into_raw()) |
| 64 | +} |
| 65 | + |
| 66 | +fn parse_uint(s: &str) -> Result<U256, String> { |
| 67 | + if s.starts_with("0x") { |
| 68 | + s.parse::<U256>().map_err(|err| err.to_string()) |
| 69 | + } else { |
| 70 | + match U256::from_dec_str(s) { |
| 71 | + Ok(val) => Ok(val), |
| 72 | + Err(dec_err) => s.parse::<U256>().map_err(|hex_err| { |
| 73 | + format!("could not parse value as decimal or hex: {dec_err}, {hex_err}") |
| 74 | + }), |
| 75 | + } |
| 76 | + } |
| 77 | +} |
| 78 | + |
| 79 | +fn parse_bytes(s: &str) -> Result<Vec<u8>, String> { |
| 80 | + hex::decode(s).map_err(|e| e.to_string()) |
| 81 | +} |
| 82 | + |
| 83 | +#[instrument(level = "error", name = "util", target = "evm::cheatcodes", skip_all)] |
| 84 | +pub fn apply<DB: Database>( |
| 85 | + _state: &mut Cheatcodes, |
| 86 | + _data: &mut EVMData<'_, DB>, |
| 87 | + call: &HEVMCalls, |
| 88 | +) -> Option<Result> { |
| 89 | + Some(match call { |
| 90 | + HEVMCalls::ToString0(inner) => { |
| 91 | + Ok(ethers::abi::encode(&[Token::String(inner.0.pretty())]).into()) |
| 92 | + } |
| 93 | + HEVMCalls::ToString1(inner) => { |
| 94 | + Ok(ethers::abi::encode(&[Token::String(inner.0.pretty())]).into()) |
| 95 | + } |
| 96 | + HEVMCalls::ToString2(inner) => { |
| 97 | + Ok(ethers::abi::encode(&[Token::String(inner.0.pretty())]).into()) |
| 98 | + } |
| 99 | + HEVMCalls::ToString3(inner) => { |
| 100 | + Ok(ethers::abi::encode(&[Token::String(inner.0.pretty())]).into()) |
| 101 | + } |
| 102 | + HEVMCalls::ToString4(inner) => { |
| 103 | + Ok(ethers::abi::encode(&[Token::String(inner.0.pretty())]).into()) |
| 104 | + } |
| 105 | + HEVMCalls::ToString5(inner) => { |
| 106 | + Ok(ethers::abi::encode(&[Token::String(inner.0.pretty())]).into()) |
| 107 | + } |
| 108 | + HEVMCalls::ParseBytes(inner) => parse(&inner.0, &ParamType::Bytes), |
| 109 | + HEVMCalls::ParseAddress(inner) => parse(&inner.0, &ParamType::Address), |
| 110 | + HEVMCalls::ParseUint(inner) => parse(&inner.0, &ParamType::Uint(256)), |
| 111 | + HEVMCalls::ParseInt(inner) => parse(&inner.0, &ParamType::Int(256)), |
| 112 | + HEVMCalls::ParseBytes32(inner) => parse(&inner.0, &ParamType::FixedBytes(32)), |
| 113 | + HEVMCalls::ParseBool(inner) => parse(&inner.0, &ParamType::Bool), |
| 114 | + _ => return None, |
| 115 | + }) |
| 116 | +} |
| 117 | + |
| 118 | +#[cfg(test)] |
| 119 | +mod tests { |
| 120 | + use super::*; |
| 121 | + use ethers::abi::AbiDecode; |
| 122 | + |
| 123 | + #[test] |
| 124 | + fn test_uint_env() { |
| 125 | + let pk = "0x10532cc9d0d992825c3f709c62c969748e317a549634fb2a9fa949326022e81f"; |
| 126 | + let val: U256 = pk.parse().unwrap(); |
| 127 | + let parsed = parse(pk, &ParamType::Uint(256)).unwrap(); |
| 128 | + let decoded = U256::decode(&parsed).unwrap(); |
| 129 | + assert_eq!(val, decoded); |
| 130 | + |
| 131 | + let parsed = parse(pk.strip_prefix("0x").unwrap(), &ParamType::Uint(256)).unwrap(); |
| 132 | + let decoded = U256::decode(&parsed).unwrap(); |
| 133 | + assert_eq!(val, decoded); |
| 134 | + |
| 135 | + let parsed = parse("1337", &ParamType::Uint(256)).unwrap(); |
| 136 | + let decoded = U256::decode(&parsed).unwrap(); |
| 137 | + assert_eq!(U256::from(1337u64), decoded); |
| 138 | + } |
| 139 | + |
| 140 | + #[test] |
| 141 | + fn test_int_env() { |
| 142 | + let val = U256::from(100u64); |
| 143 | + let parsed = parse(&val.to_string(), &ParamType::Int(256)).unwrap(); |
| 144 | + let decoded = I256::decode(parsed).unwrap(); |
| 145 | + assert_eq!(val, decoded.try_into().unwrap()); |
| 146 | + |
| 147 | + let parsed = parse("100", &ParamType::Int(256)).unwrap(); |
| 148 | + let decoded = I256::decode(parsed).unwrap(); |
| 149 | + assert_eq!(U256::from(100u64), decoded.try_into().unwrap()); |
| 150 | + } |
| 151 | +} |
0 commit comments