foundry_cheatcodes/
env.rs

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

foundry_cheatcodes/
env.rs
