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    if env_is_missing(key) {
283        return Ok(default.abi_encode());
284    }
285    let val = get_env(key)?;
286    // An empty string is a valid `string` value, but it can never parse as any other type, so a
287    // variable that is set but empty is treated as not found for the other typed overloads.
288    if val.is_empty() && !matches!(ty, DynSolType::String) {
289        return Ok(default.abi_encode());
290    }
291    string::parse(&val, ty).map_err(map_env_err(key, &val))
292}
293
294fn env_array(key: &str, delim: &str, ty: &DynSolType) -> Result {
295    get_env(key).and_then(|val| {
296        string::parse_array(val.split(delim).map(str::trim), ty).map_err(map_env_err(key, &val))
297    })
298}
299
300fn env_array_default<T: SolValue>(key: &str, delim: &str, default: &T, ty: &DynSolType) -> Result {
301    if env_is_missing(key) {
302        return Ok(default.abi_encode());
303    }
304    env_array(key, delim, ty)
305}
306
307/// Returns `true` if the variable is not set at all. The `envOr` cheatcodes fall back to their
308/// default only then (or when the variable is set but empty, for the scalar typed overloads); a
309/// variable that is set to a value that cannot be parsed is an error.
310fn env_is_missing(key: &str) -> bool {
311    matches!(env::var(key), Err(env::VarError::NotPresent))
312}
313
314fn get_env(key: &str) -> Result<String> {
315    match env::var(key) {
316        Ok(val) => Ok(val),
317        Err(env::VarError::NotPresent) => Err(fmt_err!("environment variable {key:?} not found")),
318        Err(env::VarError::NotUnicode(s)) => {
319            Err(fmt_err!("environment variable {key:?} was not valid unicode: {s:?}"))
320        }
321    }
322}
323
324/// Converts the error message of a failed parsing attempt to a more user-friendly message that
325/// doesn't leak the value.
326fn map_env_err<'a>(key: &'a str, value: &'a str) -> impl FnOnce(Error) -> Error + 'a {
327    move |e| {
328        // failed parsing <value> as type `uint256`: parser error:
329        // <value>
330        //   ^
331        //   expected at least one digit
332        let mut e = e.to_string();
333        e = e.replacen(&format!("\"{value}\""), &format!("${key}"), 1);
334        e = e.replacen(&format!("\n{value}\n"), &format!("\n${key}\n"), 1);
335        Error::from(e)
336    }
337}
338
339#[cfg(test)]
340mod tests {
341    use super::*;
342
343    #[test]
344    fn env_or_returns_default_when_missing() {
345        let key = "env_or_missing";
346        unsafe {
347            env::remove_var(key);
348        }
349        let default = 7u64;
350        let encoded = default.abi_encode();
351        assert_eq!(env_default(key, &default, &DynSolType::Uint(256)).unwrap(), encoded);
352        assert_eq!(env_array_default(key, ",", &default, &DynSolType::Uint(256)).unwrap(), encoded);
353    }
354
355    #[test]
356    fn env_or_returns_default_when_empty() {
357        let key = "env_or_empty";
358        unsafe {
359            env::set_var(key, "");
360        }
361        let default = 7u64;
362        assert_eq!(
363            env_default(key, &default, &DynSolType::Uint(256)).unwrap(),
364            default.abi_encode()
365        );
366        assert_eq!(
367            env_array_default(key, ",", &default, &DynSolType::Uint(256)).unwrap(),
368            Vec::<u64>::new().abi_encode()
369        );
370        // An empty string is a valid `string` value and must not be replaced by the default.
371        assert_eq!(
372            env_default(key, &"default", &DynSolType::String).unwrap(),
373            String::new().abi_encode()
374        );
375        unsafe {
376            env::remove_var(key);
377        }
378    }
379
380    #[test]
381    fn env_or_rejects_unparsable_value() {
382        let key = "env_or_unparsable";
383        unsafe {
384            env::set_var(key, "not_a_number");
385        }
386
387        let err = env_default(key, &7u64, &DynSolType::Uint(256)).unwrap_err().to_string();
388        assert!(err.contains("$env_or_unparsable"), "{err:?}");
389        let err =
390            env_array_default(key, ",", &7u64, &DynSolType::Uint(256)).unwrap_err().to_string();
391        assert!(err.contains("$env_or_unparsable"), "{err:?}");
392        unsafe {
393            env::remove_var(key);
394        }
395    }
396
397    #[test]
398    fn parse_env_uint() {
399        let key = "parse_env_uint";
400        let value = "t";
401        unsafe {
402            env::set_var(key, value);
403        }
404
405        let err = env(key, &DynSolType::Uint(256)).unwrap_err().to_string();
406        assert_eq!(err.matches("$parse_env_uint").count(), 2, "{err:?}");
407        unsafe {
408            env::remove_var(key);
409        }
410    }
411}
foundry_cheatcodes/env.rs

foundry_cheatcodes/
env.rs
