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#[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 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 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: ®ex::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 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}