nextest_runner/config/scripts/
env_map.rs

1// Copyright (c) The nextest Contributors
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4use crate::errors::EnvVarError;
5use serde::{Deserialize, de::Error};
6use std::{collections::BTreeMap, fmt, process::Command};
7
8/// A map of environment variables associated with a [`super::ScriptCommand`].
9///
10/// Map keys are validated at construction time.
11#[derive(Clone, Debug, Default, PartialEq)]
12pub struct ScriptCommandEnvMap(BTreeMap<String, String>);
13
14impl ScriptCommandEnvMap {
15    /// Returns the value for the given key, if present.
16    #[cfg(test)]
17    pub(crate) fn get(&self, key: &str) -> Option<&str> {
18        self.0.get(key).map(String::as_str)
19    }
20
21    /// Applies all environment variables to the command.
22    ///
23    /// All keys are validated at construction time (via [`Deserialize`] or
24    /// [`ScriptCommandEnvMap::new`]), so this method is infallible.
25    pub(crate) fn apply_env(&self, cmd: &mut Command) {
26        for (key, value) in &self.0 {
27            cmd.env(key, value);
28        }
29    }
30
31    /// Creates a new `ScriptCommandEnvMap` from a `BTreeMap`, validating that
32    /// all keys are valid environment variable names.
33    #[cfg(test)]
34    pub(crate) fn new(
35        value: BTreeMap<String, String>,
36    ) -> Result<Self, crate::errors::ErrorList<crate::errors::EnvVarError>> {
37        use crate::errors::ErrorList;
38
39        let mut errors = Vec::new();
40        for key in value.keys() {
41            if let Err(err) = validate_env_var_key(key) {
42                errors.push(err);
43            }
44        }
45        if let Some(err) = ErrorList::new("unsupported environment variables", errors) {
46            return Err(err);
47        }
48        Ok(Self(value))
49    }
50}
51
52impl<'de> Deserialize<'de> for ScriptCommandEnvMap {
53    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
54    where
55        D: serde::Deserializer<'de>,
56    {
57        struct EnvMapVisitor;
58
59        impl<'de> serde::de::Visitor<'de> for EnvMapVisitor {
60            type Value = ScriptCommandEnvMap;
61
62            fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
63                formatter.write_str("a map of environment variable names to values")
64            }
65
66            fn visit_map<A>(self, mut map: A) -> Result<Self::Value, A::Error>
67            where
68                A: serde::de::MapAccess<'de>,
69            {
70                let mut env = BTreeMap::new();
71                while let Some((key, value)) = map.next_entry::<String, String>()? {
72                    if let Err(err) = validate_env_var_key(&key) {
73                        return Err(A::Error::invalid_value(
74                            serde::de::Unexpected::Str(&key),
75                            &err,
76                        ));
77                    }
78                    env.insert(key, value);
79                }
80                Ok(ScriptCommandEnvMap(env))
81            }
82        }
83
84        deserializer.deserialize_map(EnvMapVisitor)
85    }
86}
87
88// Validates against the most conservative definition of a valid environment
89// variable key. The definition of "Name" is taken from POSIX.1-2024 [1]:
90//
91// > In the shell command language, a word consisting solely of underscores,
92// > digits, and alphabetics from the portable character set. The first
93// > character of a name is not a digit.
94//
95// This is more conservative than strictly necessary: chapter 8 of
96// POSIX.1-2024 [2] and Microsoft's documentation [3] only prohibit '=' in
97// keys, and POSIX notes that implementations may permit other characters.
98// However, restricting to the portable character set is not wrong and avoids
99// cross-platform surprises (e.g. shells that reject non-POSIX names, or NUL
100// bytes that are not portable).
101//
102// [1]: https://pubs.opengroup.org/onlinepubs/9799919799/basedefs/V1_chap03.html#tag_03_216
103// [2]: https://pubs.opengroup.org/onlinepubs/9699919799/basedefs/V1_chap08.html
104// [3]: https://learn.microsoft.com/en-us/windows/win32/procthread/environment-variables
105
106/// Validates a str to see if it is suitable to be a key of an environment
107/// variable.
108pub(crate) fn validate_env_var_key(key: &str) -> Result<(), EnvVarError> {
109    let mut chars = key.chars();
110    match chars.next() {
111        Some(c) if c.is_ascii_alphabetic() || c == '_' => {}
112        _ => {
113            return Err(EnvVarError::InvalidKeyStartChar {
114                key: key.to_owned(),
115            });
116        }
117    }
118    if !chars.all(|c| c.is_ascii_alphanumeric() || c == '_') {
119        return Err(EnvVarError::InvalidKey {
120            key: key.to_owned(),
121        });
122    }
123    if key.starts_with("NEXTEST") {
124        return Err(EnvVarError::ReservedKey {
125            key: key.to_owned(),
126        });
127    }
128    Ok(())
129}
130
131#[cfg(test)]
132mod tests {
133    use super::*;
134    use std::collections::BTreeMap;
135
136    #[test]
137    fn apply_env() {
138        let mut cmd = std::process::Command::new("demo");
139        cmd.env_clear();
140        let env = ScriptCommandEnvMap::new(BTreeMap::from([
141            (String::from("KEY_A"), String::from("value_a")),
142            (String::from("KEY_B"), String::from("value_b")),
143        ]))
144        .expect("valid env var keys");
145
146        env.apply_env(&mut cmd);
147        let applied: BTreeMap<_, _> = cmd.get_envs().collect();
148        assert_eq!(applied.len(), 2, "all env vars applied");
149    }
150
151    #[test]
152    fn new_rejects_invalid_key() {
153        let err =
154            ScriptCommandEnvMap::new(BTreeMap::from([(String::from("INVALID "), String::new())]))
155                .unwrap_err();
156        assert_eq!(
157            err.to_string(),
158            "key `INVALID ` does not consist solely of letters, digits, and underscores",
159        );
160    }
161
162    #[test]
163    fn validate_env_var_key_valid() {
164        validate_env_var_key("MY_ENV_VAR").unwrap();
165        validate_env_var_key("MY_ENV_VAR_1").unwrap();
166        validate_env_var_key("__NEXTEST_TEST").unwrap();
167    }
168
169    #[test]
170    fn validate_env_var_key_invalid() {
171        let cases = [
172            ("", "key `` does not start with a letter or underscore"),
173            (" ", "key ` ` does not start with a letter or underscore"),
174            ("=", "key `=` does not start with a letter or underscore"),
175            ("0", "key `0` does not start with a letter or underscore"),
176            (
177                "0TEST ",
178                "key `0TEST ` does not start with a letter or underscore",
179            ),
180            (
181                "=TEST=",
182                "key `=TEST=` does not start with a letter or underscore",
183            ),
184            (
185                "TEST TEST",
186                "key `TEST TEST` does not consist solely of letters, digits, and underscores",
187            ),
188            (
189                "TESTTEST\n",
190                "key `TESTTEST\n` does not consist solely of letters, digits, and underscores",
191            ),
192            (
193                "TEST=TEST",
194                "key `TEST=TEST` does not consist solely of letters, digits, and underscores",
195            ),
196            (
197                "TEST=",
198                "key `TEST=` does not consist solely of letters, digits, and underscores",
199            ),
200            (
201                "NEXTEST",
202                "key `NEXTEST` begins with `NEXTEST`, which is reserved for internal use",
203            ),
204            (
205                "NEXTEST_NEXTEST",
206                "key `NEXTEST_NEXTEST` begins with `NEXTEST`, which is reserved for internal use",
207            ),
208        ];
209
210        for (key, message) in cases {
211            let err = validate_env_var_key(key).unwrap_err();
212            let actual_message = err.to_string();
213
214            assert_eq!(
215                actual_message, *message,
216                "key validation error message equals expected"
217            );
218        }
219    }
220}