Skip to main content

utils/
variables.rs

1use regex::Regex;
2use std::borrow::Cow;
3use std::collections::BTreeMap;
4use std::env;
5use std::fmt;
6use std::sync::Arc;
7use thiserror::Error;
8
9const ESCAPE_PLACEHOLDER: &str = "\x00ESCAPED_DOLLAR\x00";
10
11type GetEnvVarFn = Arc<dyn Fn(&str) -> Option<String> + Send + Sync>;
12
13/// A small key→value lookup used when expanding `$VAR` / `${VAR}` references.
14///
15/// `Vars` is checked before falling back to an env lookup (the process
16/// environment by default), so callers can inject locally-scoped values
17/// (e.g. `WORKSPACE`) without leaking them into global state.
18#[derive(Clone)]
19pub struct Vars {
20    map: BTreeMap<String, String>,
21    get_env_var: GetEnvVarFn,
22    escape_re: Regex,
23    bracketed_re: Regex,
24    simple_re: Regex,
25}
26
27impl fmt::Debug for Vars {
28    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
29        f.debug_struct("Vars").field("map", &self.map).finish_non_exhaustive()
30    }
31}
32
33impl Default for Vars {
34    fn default() -> Self {
35        Self::new()
36    }
37}
38
39impl Vars {
40    pub fn new() -> Self {
41        Self {
42            map: BTreeMap::new(),
43            get_env_var: Arc::new(|k| env::var(k).ok()),
44            escape_re: Regex::new(r"\$\$").unwrap(),
45            bracketed_re: Regex::new(r"\$\{([^}]+)\}").unwrap(),
46            simple_re: Regex::new(r"\$([a-zA-Z_][a-zA-Z0-9_]*)").unwrap(),
47        }
48    }
49
50    /// Override the env-var fallback. Useful in tests to avoid mutating the
51    /// process environment.
52    pub fn with_env_lookup(mut self, f: impl Fn(&str) -> Option<String> + Send + Sync + 'static) -> Self {
53        self.get_env_var = Arc::new(f);
54        self
55    }
56
57    pub fn with(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
58        self.insert(key, value);
59        self
60    }
61
62    pub fn insert(&mut self, key: impl Into<String>, value: impl Into<String>) -> &mut Self {
63        self.map.insert(key.into(), value.into());
64        self
65    }
66
67    pub fn get(&self, key: &str) -> Option<Cow<'_, str>> {
68        if let Some(value) = self.map.get(key) {
69            Some(Cow::Borrowed(value.as_str()))
70        } else {
71            (self.get_env_var)(key).map(Cow::Owned)
72        }
73    }
74
75    /// Expand `$VAR` / `${VAR}` references in `template`, consulting `self` before
76    /// the env lookup. `$$` escapes to a literal `$`.
77    pub fn expand(&self, template: &str) -> Result<String, VarError> {
78        let escaped = self.escape_re.replace_all(template, ESCAPE_PLACEHOLDER);
79        let bracketed = self.substitute(&self.bracketed_re, &escaped)?;
80        let simple = self.substitute(&self.simple_re, &bracketed)?;
81        Ok(simple.replace(ESCAPE_PLACEHOLDER, "$"))
82    }
83
84    fn substitute(&self, re: &Regex, text: &str) -> Result<String, VarError> {
85        let mut missing = None;
86        let result = re.replace_all(text, |caps: &regex::Captures| {
87            let var_name = &caps[1];
88            if let Some(value) = self.get(var_name) {
89                value.into_owned()
90            } else {
91                missing = Some(var_name.to_string());
92                caps[0].to_string()
93            }
94        });
95        match missing {
96            Some(var) => Err(VarError::NotFound(var)),
97            None => Ok(result.into_owned()),
98        }
99    }
100
101    /// Returns `true` if `s` contains a `$VAR` or `${VAR}` style reference.
102    pub fn has_reference(&self, s: &str) -> bool {
103        self.bracketed_re.is_match(s) || self.simple_re.is_match(s)
104    }
105}
106
107#[derive(Debug, Error)]
108pub enum VarError {
109    #[error("Environment variable '{0}' not found")]
110    NotFound(String),
111}
112
113#[cfg(test)]
114mod tests {
115    use super::*;
116
117    fn no_env() -> Vars {
118        Vars::new().with_env_lookup(|_| None)
119    }
120
121    #[test]
122    fn test_simple_var() {
123        let vars = no_env().with("TEST_VAR_SIMPLE", "hello");
124        let result = vars.expand("$TEST_VAR_SIMPLE world").unwrap();
125        assert_eq!(result, "hello world");
126    }
127
128    #[test]
129    fn test_bracketed_var() {
130        let vars = no_env().with("TEST_VAR_BRACKET", "hello");
131        let result = vars.expand("${TEST_VAR_BRACKET} world").unwrap();
132        assert_eq!(result, "hello world");
133    }
134
135    #[test]
136    fn test_escape_sequence() {
137        let result = no_env().expand("$$VAR").unwrap();
138        assert_eq!(result, "$VAR");
139    }
140
141    #[test]
142    fn test_multiple_vars() {
143        let vars = no_env().with("VAR1", "hello").with("VAR2", "world");
144        let result = vars.expand("$VAR1 ${VAR2}!").unwrap();
145        assert_eq!(result, "hello world!");
146    }
147
148    #[test]
149    fn test_missing_var() {
150        let result = no_env().expand("$MISSING_VAR");
151        assert!(matches!(result, Err(VarError::NotFound(ref name)) if name == "MISSING_VAR"));
152    }
153
154    #[test]
155    fn test_unclosed_brace_left_as_is() {
156        let result = no_env().expand("${VAR").unwrap();
157        assert_eq!(result, "${VAR");
158    }
159
160    #[test]
161    fn test_empty_string() {
162        let result = no_env().expand("").unwrap();
163        assert_eq!(result, "");
164    }
165
166    #[test]
167    fn test_no_vars() {
168        let result = no_env().expand("plain text").unwrap();
169        assert_eq!(result, "plain text");
170    }
171
172    #[test]
173    fn test_dollar_at_end() {
174        let result = no_env().expand("text$").unwrap();
175        assert_eq!(result, "text$");
176    }
177
178    #[test]
179    fn test_var_with_underscore() {
180        let vars = no_env().with("MY_TEST_VAR", "value");
181        let result = vars.expand("$MY_TEST_VAR").unwrap();
182        assert_eq!(result, "value");
183    }
184
185    #[test]
186    fn test_var_with_numbers() {
187        let vars = no_env().with("VAR123", "value");
188        let result = vars.expand("$VAR123").unwrap();
189        assert_eq!(result, "value");
190    }
191
192    #[test]
193    fn test_special_char_stops_var_name() {
194        let vars = no_env().with("VAR", "value");
195        let result = vars.expand("$VAR-suffix").unwrap();
196        assert_eq!(result, "value-suffix");
197    }
198
199    #[test]
200    fn vars_lookup_takes_precedence_over_env() {
201        let vars = Vars::new()
202            .with_env_lookup(|k| (k == "WORKSPACE").then(|| "/from-env".into()))
203            .with("WORKSPACE", "/from-vars");
204        let result = vars.expand("${WORKSPACE}/foo").unwrap();
205        assert_eq!(result, "/from-vars/foo");
206    }
207
208    #[test]
209    fn vars_falls_through_to_env_when_missing() {
210        let vars = Vars::new().with_env_lookup(|k| (k == "ONLY_IN_ENV").then(|| "from-env".into()));
211        let result = vars.expand("$ONLY_IN_ENV").unwrap();
212        assert_eq!(result, "from-env");
213    }
214
215    #[test]
216    fn has_reference_detects_both_forms() {
217        let vars = no_env();
218        assert!(vars.has_reference("${WORKSPACE}/foo"));
219        assert!(vars.has_reference("$HOME/bar"));
220        assert!(!vars.has_reference("plain/path"));
221        assert!(!vars.has_reference("$"));
222        assert!(!vars.has_reference("$$"));
223    }
224}