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

foundry_cheatcodes/
env.rs
