foundry_cheatcodes/
env.rs

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

foundry_cheatcodes/
env.rs
