foundry_cheatcodes/
env.rs

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

foundry_cheatcodes/
env.rs
