foundry_cheatcodes/
error.rs

1use crate::Vm;
2use alloy_primitives::{Bytes, hex};
3use alloy_signer::Error as SignerError;
4use alloy_signer_local::LocalSignerError;
5use alloy_sol_types::SolError;
6use foundry_common::errors::FsPathError;
7use foundry_config::UnresolvedEnvVarError;
8use foundry_evm_core::backend::{BackendError, DatabaseError};
9use foundry_wallets::error::WalletSignerError;
10use k256::ecdsa::signature::Error as SignatureError;
11use revm::context_interface::result::EVMError;
12use std::{borrow::Cow, fmt};
13
14/// Cheatcode result type.
15///
16/// Type alias with a default Ok type of [`Vec<u8>`], and default Err type of [`Error`].
17pub type Result<T = Vec<u8>, E = Error> = std::result::Result<T, E>;
18
19macro_rules! fmt_err {
20    ($msg:literal $(,)?) => {
21        $crate::Error::fmt(::std::format_args!($msg))
22    };
23    ($err:expr $(,)?) => {
24        <$crate::Error as ::std::convert::From<_>>::from($err)
25    };
26    ($fmt:expr, $($arg:tt)*) => {
27        $crate::Error::fmt(::std::format_args!($fmt, $($arg)*))
28    };
29}
30
31macro_rules! bail {
32    ($msg:literal $(,)?) => {
33        return ::std::result::Result::Err(fmt_err!($msg))
34    };
35    ($err:expr $(,)?) => {
36        return ::std::result::Result::Err(fmt_err!($err))
37    };
38    ($fmt:expr, $($arg:tt)*) => {
39        return ::std::result::Result::Err(fmt_err!($fmt, $($arg)*))
40    };
41}
42
43macro_rules! ensure {
44    ($cond:expr $(,)?) => {
45        if !$cond {
46            return ::std::result::Result::Err($crate::Error::custom(
47                ::std::concat!("Condition failed: `", ::std::stringify!($cond), "`")
48            ));
49        }
50    };
51    ($cond:expr, $msg:literal $(,)?) => {
52        if !$cond {
53            return ::std::result::Result::Err(fmt_err!($msg));
54        }
55    };
56    ($cond:expr, $err:expr $(,)?) => {
57        if !$cond {
58            return ::std::result::Result::Err(fmt_err!($err));
59        }
60    };
61    ($cond:expr, $fmt:expr, $($arg:tt)*) => {
62        if !$cond {
63            return ::std::result::Result::Err(fmt_err!($fmt, $($arg)*));
64        }
65    };
66}
67
68/// Error thrown by cheatcodes.
69// This uses a custom repr to minimize the size of the error.
70// The repr is basically `enum { Cow<'static, str>, Cow<'static, [u8]> }`
71pub struct Error {
72    /// If true, encode `data` as `Error(string)`, otherwise encode it directly as `bytes`.
73    is_str: bool,
74    /// Whether this was constructed from an owned byte vec, which means we have to drop the data
75    /// in `impl Drop`.
76    drop: bool,
77    /// The error data. Always a valid pointer, and never modified.
78    data: *const [u8],
79}
80
81impl std::error::Error for Error {}
82
83impl fmt::Debug for Error {
84    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
85        f.write_str("Error::")?;
86        self.kind().fmt(f)
87    }
88}
89
90impl fmt::Display for Error {
91    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
92        self.kind().fmt(f)
93    }
94}
95
96/// Kind of cheatcode errors.
97///
98/// Constructed by [`Error::kind`].
99#[derive(Debug)]
100pub enum ErrorKind<'a> {
101    /// A string error, ABI-encoded as `CheatcodeError(string)`.
102    String(&'a str),
103    /// A raw bytes error. Does not get encoded.
104    Bytes(&'a [u8]),
105}
106
107impl fmt::Display for ErrorKind<'_> {
108    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
109        match *self {
110            Self::String(ss) => f.write_str(ss),
111            Self::Bytes(b) => f.write_str(&hex::encode_prefixed(b)),
112        }
113    }
114}
115
116impl Error {
117    /// Creates a new error and ABI encodes it as `CheatcodeError(string)`.
118    pub fn encode(error: impl Into<Self>) -> Bytes {
119        error.into().abi_encode().into()
120    }
121
122    /// Creates a new error with a custom message.
123    pub fn display(msg: impl fmt::Display) -> Self {
124        Self::fmt(format_args!("{msg}"))
125    }
126
127    /// Creates a new error with a custom [`fmt::Arguments`] message.
128    pub fn fmt(args: fmt::Arguments<'_>) -> Self {
129        match args.as_str() {
130            Some(s) => Self::new_str(s),
131            None => Self::new_string(std::fmt::format(args)),
132        }
133    }
134
135    /// ABI-encodes this error as `CheatcodeError(string)` if the inner message is a string,
136    /// otherwise returns the raw bytes.
137    pub fn abi_encode(&self) -> Vec<u8> {
138        match self.kind() {
139            ErrorKind::String(string) => Vm::CheatcodeError { message: string.into() }.abi_encode(),
140            ErrorKind::Bytes(bytes) => bytes.into(),
141        }
142    }
143
144    /// Returns the kind of this error.
145    pub fn kind(&self) -> ErrorKind<'_> {
146        let data = self.data();
147        if self.is_str {
148            debug_assert!(std::str::from_utf8(data).is_ok());
149            ErrorKind::String(unsafe { std::str::from_utf8_unchecked(data) })
150        } else {
151            ErrorKind::Bytes(data)
152        }
153    }
154
155    /// Returns the raw data of this error.
156    pub const fn data(&self) -> &[u8] {
157        unsafe { &*self.data }
158    }
159
160    /// Returns `true` if this error is a human-readable string.
161    pub const fn is_str(&self) -> bool {
162        self.is_str
163    }
164
165    /// Prefixes a string error with the name of the cheatcode that produced it.
166    pub(crate) fn with_cheatcode_name(self, cheat: &spec::Cheatcode<'static>) -> Self {
167        let name = cheat.func.signature.split('(').next().unwrap();
168        // Skip showing the cheatcode name for:
169        // - assertions: too verbose, and can already be inferred from the error message
170        // - `rpcUrl`: forge-std relies on it in `getChainWithUpdatedRpcUrl`
171        if self.is_str && !name.contains("assert") && name != "rpcUrl" {
172            fmt_err!("vm.{name}: {self}")
173        } else {
174            self
175        }
176    }
177
178    fn new_str(data: &'static str) -> Self {
179        Self::_new(true, false, data.as_bytes())
180    }
181
182    fn new_string(data: String) -> Self {
183        Self::_new(true, true, Box::into_raw(data.into_boxed_str().into_boxed_bytes()))
184    }
185
186    fn new_bytes(data: &'static [u8]) -> Self {
187        Self::_new(false, false, data)
188    }
189
190    fn new_vec(data: Vec<u8>) -> Self {
191        Self::_new(false, true, Box::into_raw(data.into_boxed_slice()))
192    }
193
194    fn _new(is_str: bool, drop: bool, data: *const [u8]) -> Self {
195        debug_assert!(!data.is_null());
196        Self { is_str, drop, data }
197    }
198}
199
200impl Drop for Error {
201    fn drop(&mut self) {
202        if self.drop {
203            drop(unsafe { Box::<[u8]>::from_raw(self.data.cast_mut()) });
204        }
205    }
206}
207
208impl From<Cow<'static, str>> for Error {
209    fn from(value: Cow<'static, str>) -> Self {
210        match value {
211            Cow::Borrowed(str) => Self::new_str(str),
212            Cow::Owned(string) => Self::new_string(string),
213        }
214    }
215}
216
217impl From<String> for Error {
218    fn from(value: String) -> Self {
219        Self::new_string(value)
220    }
221}
222
223impl From<&'static str> for Error {
224    fn from(value: &'static str) -> Self {
225        Self::new_str(value)
226    }
227}
228
229impl From<Cow<'static, [u8]>> for Error {
230    fn from(value: Cow<'static, [u8]>) -> Self {
231        match value {
232            Cow::Borrowed(bytes) => Self::new_bytes(bytes),
233            Cow::Owned(vec) => Self::new_vec(vec),
234        }
235    }
236}
237
238impl From<&'static [u8]> for Error {
239    fn from(value: &'static [u8]) -> Self {
240        Self::new_bytes(value)
241    }
242}
243
244impl<const N: usize> From<&'static [u8; N]> for Error {
245    fn from(value: &'static [u8; N]) -> Self {
246        Self::new_bytes(value)
247    }
248}
249
250impl From<Vec<u8>> for Error {
251    fn from(value: Vec<u8>) -> Self {
252        Self::new_vec(value)
253    }
254}
255
256impl From<Bytes> for Error {
257    fn from(value: Bytes) -> Self {
258        Self::new_vec(value.into())
259    }
260}
261
262// So we can use `?` on `Result<_, Error>`.
263macro_rules! impl_from {
264    ($($t:ty),* $(,)?) => {$(
265        impl From<$t> for Error {
266            fn from(value: $t) -> Self {
267                Self::display(value)
268            }
269        }
270    )*};
271}
272
273impl_from!(
274    alloy_sol_types::Error,
275    alloy_dyn_abi::Error,
276    alloy_primitives::SignatureError,
277    alloy_consensus::crypto::RecoveryError,
278    FsPathError,
279    hex::FromHexError,
280    BackendError,
281    DatabaseError,
282    jsonpath_lib::JsonPathError,
283    serde_json::Error,
284    SignatureError,
285    std::io::Error,
286    std::num::TryFromIntError,
287    std::str::Utf8Error,
288    std::string::FromUtf8Error,
289    UnresolvedEnvVarError,
290    LocalSignerError,
291    SignerError,
292    WalletSignerError,
293);
294
295impl<T: Into<BackendError>> From<EVMError<T>> for Error {
296    fn from(err: EVMError<T>) -> Self {
297        Self::display(BackendError::from(err))
298    }
299}
300
301impl From<eyre::Report> for Error {
302    fn from(err: eyre::Report) -> Self {
303        Self::new_string(foundry_common::errors::display_chain(&err))
304    }
305}
306
307#[cfg(test)]
308mod tests {
309    use super::*;
310
311    #[test]
312    fn encode() {
313        let error = Vm::CheatcodeError { message: "hello".into() }.abi_encode();
314        assert_eq!(Error::from("hello").abi_encode(), error);
315        assert_eq!(Error::encode("hello"), error);
316
317        assert_eq!(Error::from(b"hello").abi_encode(), b"hello");
318        assert_eq!(Error::encode(b"hello"), b"hello"[..]);
319    }
320}
foundry_cheatcodes/error.rs

foundry_cheatcodes/
error.rs
