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

foundry_cheatcodes/
env.rs
