foundry_cheatcodes/
env.rs

1//! Implementations of [`Environment`](spec::Group::Environment) cheatcodes.
2
3use crate::{Cheatcode, Cheatcodes, Error, Result, Vm::*, string};
4use alloy_dyn_abi::DynSolType;
5use alloy_sol_types::SolValue;
6use foundry_evm_core::evm::FoundryEvmNetwork;
7use std::{env, sync::OnceLock};
8
9/// Stores the forge execution context for the duration of the program.
10pub static FORGE_CONTEXT: OnceLock<ForgeContext> = OnceLock::new();
11
12/// Returns the current forge execution context, if it has been set.
13pub fn current_execution_context() -> Option<ForgeContext> {
14    FORGE_CONTEXT.get().copied()
15}
16
17impl Cheatcode for setEnvCall {
18    fn apply<FEN: FoundryEvmNetwork>(&self, _state: &mut Cheatcodes<FEN>) -> Result {
19        let Self { name: key, value } = self;
20        if key.is_empty() {
21            Err(fmt_err!("environment variable key can't be empty"))
22        } else if key.contains('=') {
23            Err(fmt_err!("environment variable key can't contain equal sign `=`"))
24        } else if key.contains('\0') {
25            Err(fmt_err!("environment variable key can't contain NUL character `\\0`"))
26        } else if value.contains('\0') {
27            Err(fmt_err!("environment variable value can't contain NUL character `\\0`"))
28        } else {
29            unsafe {
30                env::set_var(key, value);
31            }
32            Ok(Default::default())
33        }
34    }
35}
36
37impl Cheatcode for resolveEnvCall {
38    fn apply<FEN: FoundryEvmNetwork>(&self, _state: &mut Cheatcodes<FEN>) -> Result {
39        let Self { input } = self;
40        let resolved = foundry_config::resolve::interpolate(input)
41            .map_err(|e| fmt_err!("failed to resolve env var: {e}"))?;
42        Ok(resolved.abi_encode())
43    }
44}
45
46impl Cheatcode for envExistsCall {
47    fn apply<FEN: FoundryEvmNetwork>(&self, _state: &mut Cheatcodes<FEN>) -> Result {
48        let Self { name } = self;
49        Ok(env::var(name).is_ok().abi_encode())
50    }
51}
52
53impl Cheatcode for envBool_0Call {
54    fn apply<FEN: FoundryEvmNetwork>(&self, _state: &mut Cheatcodes<FEN>) -> Result {
55        let Self { name } = self;
56        env(name, &DynSolType::Bool)
57    }
58}
59
60impl Cheatcode for envUint_0Call {
61    fn apply<FEN: FoundryEvmNetwork>(&self, _state: &mut Cheatcodes<FEN>) -> Result {
62        let Self { name } = self;
63        env(name, &DynSolType::Uint(256))
64    }
65}
66
67impl Cheatcode for envInt_0Call {
68    fn apply<FEN: FoundryEvmNetwork>(&self, _state: &mut Cheatcodes<FEN>) -> Result {
69        let Self { name } = self;
70        env(name, &DynSolType::Int(256))
71    }
72}
73
74impl Cheatcode for envAddress_0Call {
75    fn apply<FEN: FoundryEvmNetwork>(&self, _state: &mut Cheatcodes<FEN>) -> Result {
76        let Self { name } = self;
77        env(name, &DynSolType::Address)
78    }
79}
80
81impl Cheatcode for envBytes32_0Call {
82    fn apply<FEN: FoundryEvmNetwork>(&self, _state: &mut Cheatcodes<FEN>) -> Result {
83        let Self { name } = self;
84        env(name, &DynSolType::FixedBytes(32))
85    }
86}
87
88impl Cheatcode for envString_0Call {
89    fn apply<FEN: FoundryEvmNetwork>(&self, _state: &mut Cheatcodes<FEN>) -> Result {
90        let Self { name } = self;
91        env(name, &DynSolType::String)
92    }
93}
94
95impl Cheatcode for envBytes_0Call {
96    fn apply<FEN: FoundryEvmNetwork>(&self, _state: &mut Cheatcodes<FEN>) -> Result {
97        let Self { name } = self;
98        env(name, &DynSolType::Bytes)
99    }
100}
101
102impl Cheatcode for envBool_1Call {
103    fn apply<FEN: FoundryEvmNetwork>(&self, _state: &mut Cheatcodes<FEN>) -> Result {
104        let Self { name, delim } = self;
105        env_array(name, delim, &DynSolType::Bool)
106    }
107}
108
109impl Cheatcode for envUint_1Call {
110    fn apply<FEN: FoundryEvmNetwork>(&self, _state: &mut Cheatcodes<FEN>) -> Result {
111        let Self { name, delim } = self;
112        env_array(name, delim, &DynSolType::Uint(256))
113    }
114}
115
116impl Cheatcode for envInt_1Call {
117    fn apply<FEN: FoundryEvmNetwork>(&self, _state: &mut Cheatcodes<FEN>) -> Result {
118        let Self { name, delim } = self;
119        env_array(name, delim, &DynSolType::Int(256))
120    }
121}
122
123impl Cheatcode for envAddress_1Call {
124    fn apply<FEN: FoundryEvmNetwork>(&self, _state: &mut Cheatcodes<FEN>) -> Result {
125        let Self { name, delim } = self;
126        env_array(name, delim, &DynSolType::Address)
127    }
128}
129
130impl Cheatcode for envBytes32_1Call {
131    fn apply<FEN: FoundryEvmNetwork>(&self, _state: &mut Cheatcodes<FEN>) -> Result {
132        let Self { name, delim } = self;
133        env_array(name, delim, &DynSolType::FixedBytes(32))
134    }
135}
136
137impl Cheatcode for envString_1Call {
138    fn apply<FEN: FoundryEvmNetwork>(&self, _state: &mut Cheatcodes<FEN>) -> Result {
139        let Self { name, delim } = self;
140        env_array(name, delim, &DynSolType::String)
141    }
142}
143
144impl Cheatcode for envBytes_1Call {
145    fn apply<FEN: FoundryEvmNetwork>(&self, _state: &mut Cheatcodes<FEN>) -> Result {
146        let Self { name, delim } = self;
147        env_array(name, delim, &DynSolType::Bytes)
148    }
149}
150
151// bool
152impl Cheatcode for envOr_0Call {
153    fn apply<FEN: FoundryEvmNetwork>(&self, _state: &mut Cheatcodes<FEN>) -> Result {
154        let Self { name, defaultValue } = self;
155        env_default(name, defaultValue, &DynSolType::Bool)
156    }
157}
158
159// uint256
160impl Cheatcode for envOr_1Call {
161    fn apply<FEN: FoundryEvmNetwork>(&self, _state: &mut Cheatcodes<FEN>) -> Result {
162        let Self { name, defaultValue } = self;
163        env_default(name, defaultValue, &DynSolType::Uint(256))
164    }
165}
166
167// int256
168impl Cheatcode for envOr_2Call {
169    fn apply<FEN: FoundryEvmNetwork>(&self, _state: &mut Cheatcodes<FEN>) -> Result {
170        let Self { name, defaultValue } = self;
171        env_default(name, defaultValue, &DynSolType::Int(256))
172    }
173}
174
175// address
176impl Cheatcode for envOr_3Call {
177    fn apply<FEN: FoundryEvmNetwork>(&self, _state: &mut Cheatcodes<FEN>) -> Result {
178        let Self { name, defaultValue } = self;
179        env_default(name, defaultValue, &DynSolType::Address)
180    }
181}
182
183// bytes32
184impl Cheatcode for envOr_4Call {
185    fn apply<FEN: FoundryEvmNetwork>(&self, _state: &mut Cheatcodes<FEN>) -> Result {
186        let Self { name, defaultValue } = self;
187        env_default(name, defaultValue, &DynSolType::FixedBytes(32))
188    }
189}
190
191// string
192impl Cheatcode for envOr_5Call {
193    fn apply<FEN: FoundryEvmNetwork>(&self, _state: &mut Cheatcodes<FEN>) -> Result {
194        let Self { name, defaultValue } = self;
195        env_default(name, defaultValue, &DynSolType::String)
196    }
197}
198
199// bytes
200impl Cheatcode for envOr_6Call {
201    fn apply<FEN: FoundryEvmNetwork>(&self, _state: &mut Cheatcodes<FEN>) -> Result {
202        let Self { name, defaultValue } = self;
203        env_default(name, defaultValue, &DynSolType::Bytes)
204    }
205}
206
207// bool[]
208impl Cheatcode for envOr_7Call {
209    fn apply<FEN: FoundryEvmNetwork>(&self, _state: &mut Cheatcodes<FEN>) -> Result {
210        let Self { name, delim, defaultValue } = self;
211        env_array_default(name, delim, defaultValue, &DynSolType::Bool)
212    }
213}
214
215// uint256[]
216impl Cheatcode for envOr_8Call {
217    fn apply<FEN: FoundryEvmNetwork>(&self, _state: &mut Cheatcodes<FEN>) -> Result {
218        let Self { name, delim, defaultValue } = self;
219        env_array_default(name, delim, defaultValue, &DynSolType::Uint(256))
220    }
221}
222
223// int256[]
224impl Cheatcode for envOr_9Call {
225    fn apply<FEN: FoundryEvmNetwork>(&self, _state: &mut Cheatcodes<FEN>) -> Result {
226        let Self { name, delim, defaultValue } = self;
227        env_array_default(name, delim, defaultValue, &DynSolType::Int(256))
228    }
229}
230
231// address[]
232impl Cheatcode for envOr_10Call {
233    fn apply<FEN: FoundryEvmNetwork>(&self, _state: &mut Cheatcodes<FEN>) -> Result {
234        let Self { name, delim, defaultValue } = self;
235        env_array_default(name, delim, defaultValue, &DynSolType::Address)
236    }
237}
238
239// bytes32[]
240impl Cheatcode for envOr_11Call {
241    fn apply<FEN: FoundryEvmNetwork>(&self, _state: &mut Cheatcodes<FEN>) -> Result {
242        let Self { name, delim, defaultValue } = self;
243        env_array_default(name, delim, defaultValue, &DynSolType::FixedBytes(32))
244    }
245}
246
247// string[]
248impl Cheatcode for envOr_12Call {
249    fn apply<FEN: FoundryEvmNetwork>(&self, _state: &mut Cheatcodes<FEN>) -> Result {
250        let Self { name, delim, defaultValue } = self;
251        env_array_default(name, delim, defaultValue, &DynSolType::String)
252    }
253}
254
255// bytes[]
256impl Cheatcode for envOr_13Call {
257    fn apply<FEN: FoundryEvmNetwork>(&self, _state: &mut Cheatcodes<FEN>) -> Result {
258        let Self { name, delim, defaultValue } = self;
259        let default = defaultValue.clone();
260        env_array_default(name, delim, &default, &DynSolType::Bytes)
261    }
262}
263
264impl Cheatcode for isContextCall {
265    fn apply<FEN: FoundryEvmNetwork>(&self, _state: &mut Cheatcodes<FEN>) -> Result {
266        let Self { context } = self;
267        Ok((FORGE_CONTEXT.get() == Some(context)).abi_encode())
268    }
269}
270
271/// Set `forge` command current execution context for the duration of the program.
272/// Execution context is immutable, subsequent calls of this function won't change the context.
273pub fn set_execution_context(context: ForgeContext) {
274    let _ = FORGE_CONTEXT.set(context);
275}
276
277fn env(key: &str, ty: &DynSolType) -> Result {
278    get_env(key).and_then(|val| string::parse(&val, ty).map_err(map_env_err(key, &val)))
279}
280
281fn env_default<T: SolValue>(key: &str, default: &T, ty: &DynSolType) -> Result {
282    Ok(env(key, ty).unwrap_or_else(|_| default.abi_encode()))
283}
284
285fn env_array(key: &str, delim: &str, ty: &DynSolType) -> Result {
286    get_env(key).and_then(|val| {
287        string::parse_array(val.split(delim).map(str::trim), ty).map_err(map_env_err(key, &val))
288    })
289}
290
291fn env_array_default<T: SolValue>(key: &str, delim: &str, default: &T, ty: &DynSolType) -> Result {
292    Ok(env_array(key, delim, ty).unwrap_or_else(|_| default.abi_encode()))
293}
294
295fn get_env(key: &str) -> Result<String> {
296    match env::var(key) {
297        Ok(val) => Ok(val),
298        Err(env::VarError::NotPresent) => Err(fmt_err!("environment variable {key:?} not found")),
299        Err(env::VarError::NotUnicode(s)) => {
300            Err(fmt_err!("environment variable {key:?} was not valid unicode: {s:?}"))
301        }
302    }
303}
304
305/// Converts the error message of a failed parsing attempt to a more user-friendly message that
306/// doesn't leak the value.
307fn map_env_err<'a>(key: &'a str, value: &'a str) -> impl FnOnce(Error) -> Error + 'a {
308    move |e| {
309        // failed parsing <value> as type `uint256`: parser error:
310        // <value>
311        //   ^
312        //   expected at least one digit
313        let mut e = e.to_string();
314        e = e.replacen(&format!("\"{value}\""), &format!("${key}"), 1);
315        e = e.replacen(&format!("\n{value}\n"), &format!("\n${key}\n"), 1);
316        Error::from(e)
317    }
318}
319
320#[cfg(test)]
321mod tests {
322    use super::*;
323
324    #[test]
325    fn parse_env_uint() {
326        let key = "parse_env_uint";
327        let value = "t";
328        unsafe {
329            env::set_var(key, value);
330        }
331
332        let err = env(key, &DynSolType::Uint(256)).unwrap_err().to_string();
333        assert_eq!(err.matches("$parse_env_uint").count(), 2, "{err:?}");
334        unsafe {
335            env::remove_var(key);
336        }
337    }
338}
foundry_cheatcodes/env.rs

foundry_cheatcodes/
env.rs
