1use crate::{Cheatcode, Cheatcodes, Error, Result, Vm::*, string};
4use alloy_dyn_abi::DynSolType;
5use alloy_sol_types::SolValue;
6use std::{env, sync::OnceLock};
7
8pub 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
145impl 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
153impl 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
161impl 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
169impl 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
177impl 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
185impl 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
193impl 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
201impl 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
209impl 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
217impl 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
225impl 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
233impl 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
241impl 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
249impl 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
265pub 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
299fn map_env_err<'a>(key: &'a str, value: &'a str) -> impl FnOnce(Error) -> Error + 'a {
302 move |e| {
303 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}