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 fn data(&self) -> &[u8] {
157        unsafe { &*self.data }
158    }
159
160    /// Returns `true` if this error is a human-readable string.
161    pub fn is_str(&self) -> bool {
162        self.is_str
163    }
164
165    fn new_str(data: &'static str) -> Self {
166        Self::_new(true, false, data.as_bytes())
167    }
168
169    fn new_string(data: String) -> Self {
170        Self::_new(true, true, Box::into_raw(data.into_boxed_str().into_boxed_bytes()))
171    }
172
173    fn new_bytes(data: &'static [u8]) -> Self {
174        Self::_new(false, false, data)
175    }
176
177    fn new_vec(data: Vec<u8>) -> Self {
178        Self::_new(false, true, Box::into_raw(data.into_boxed_slice()))
179    }
180
181    fn _new(is_str: bool, drop: bool, data: *const [u8]) -> Self {
182        debug_assert!(!data.is_null());
183        Self { is_str, drop, data }
184    }
185}
186
187impl Drop for Error {
188    fn drop(&mut self) {
189        if self.drop {
190            drop(unsafe { Box::<[u8]>::from_raw(self.data.cast_mut()) });
191        }
192    }
193}
194
195impl From<Cow<'static, str>> for Error {
196    fn from(value: Cow<'static, str>) -> Self {
197        match value {
198            Cow::Borrowed(str) => Self::new_str(str),
199            Cow::Owned(string) => Self::new_string(string),
200        }
201    }
202}
203
204impl From<String> for Error {
205    fn from(value: String) -> Self {
206        Self::new_string(value)
207    }
208}
209
210impl From<&'static str> for Error {
211    fn from(value: &'static str) -> Self {
212        Self::new_str(value)
213    }
214}
215
216impl From<Cow<'static, [u8]>> for Error {
217    fn from(value: Cow<'static, [u8]>) -> Self {
218        match value {
219            Cow::Borrowed(bytes) => Self::new_bytes(bytes),
220            Cow::Owned(vec) => Self::new_vec(vec),
221        }
222    }
223}
224
225impl From<&'static [u8]> for Error {
226    fn from(value: &'static [u8]) -> Self {
227        Self::new_bytes(value)
228    }
229}
230
231impl<const N: usize> From<&'static [u8; N]> for Error {
232    fn from(value: &'static [u8; N]) -> Self {
233        Self::new_bytes(value)
234    }
235}
236
237impl From<Vec<u8>> for Error {
238    fn from(value: Vec<u8>) -> Self {
239        Self::new_vec(value)
240    }
241}
242
243impl From<Bytes> for Error {
244    fn from(value: Bytes) -> Self {
245        Self::new_vec(value.into())
246    }
247}
248
249// So we can use `?` on `Result<_, Error>`.
250macro_rules! impl_from {
251    ($($t:ty),* $(,)?) => {$(
252        impl From<$t> for Error {
253            fn from(value: $t) -> Self {
254                Self::display(value)
255            }
256        }
257    )*};
258}
259
260impl_from!(
261    alloy_sol_types::Error,
262    alloy_dyn_abi::Error,
263    alloy_primitives::SignatureError,
264    alloy_consensus::crypto::RecoveryError,
265    eyre::Report,
266    FsPathError,
267    hex::FromHexError,
268    BackendError,
269    DatabaseError,
270    jsonpath_lib::JsonPathError,
271    serde_json::Error,
272    SignatureError,
273    std::io::Error,
274    std::num::TryFromIntError,
275    std::str::Utf8Error,
276    std::string::FromUtf8Error,
277    UnresolvedEnvVarError,
278    LocalSignerError,
279    SignerError,
280    WalletSignerError,
281);
282
283impl<T: Into<BackendError>> From<EVMError<T>> for Error {
284    fn from(err: EVMError<T>) -> Self {
285        Self::display(BackendError::from(err))
286    }
287}
288
289#[cfg(test)]
290mod tests {
291    use super::*;
292
293    #[test]
294    fn encode() {
295        let error = Vm::CheatcodeError { message: "hello".into() }.abi_encode();
296        assert_eq!(Error::from("hello").abi_encode(), error);
297        assert_eq!(Error::encode("hello"), error);
298
299        assert_eq!(Error::from(b"hello").abi_encode(), b"hello");
300        assert_eq!(Error::encode(b"hello"), b"hello"[..]);
301    }
302}
foundry_cheatcodes/error.rs

foundry_cheatcodes/
error.rs
