mcp_utils/client/
variables.rs

1use regex::Regex;
2use std::env;
3
4/// Placeholder used to escape `$$` sequences during expansion.
5const ESCAPE_PLACEHOLDER: &str = "\x00ESCAPED_DOLLAR\x00";
6
7/// Expands environment variables in a string template.
8///
9/// Supports two formats:
10/// - `$VAR` - Simple variable reference
11/// - `${VAR}` - Bracketed variable reference
12/// - `$$` - Escape sequence for literal `$`
13pub fn expand_env_vars(template: &str) -> Result<String, VarError> {
14    let escape_re = Regex::new(r"\$\$").unwrap();
15    let bracketed_re = Regex::new(r"\$\{([^}]+)\}").unwrap();
16    let simple_re = Regex::new(r"\$([a-zA-Z_][a-zA-Z0-9_]*)").unwrap();
17
18    // Replace $$ with placeholder
19    let result = escape_re.replace_all(template, ESCAPE_PLACEHOLDER);
20
21    // Replace ${VAR} with env var, tracking any missing vars
22    let mut missing_var = None;
23    let result = bracketed_re.replace_all(&result, |caps: &regex::Captures| {
24        let var_name = &caps[1];
25        if let Ok(value) = env::var(var_name) {
26            value
27        } else {
28            missing_var = Some(var_name.to_string());
29            caps[0].to_string() // Keep original if not found
30        }
31    });
32    if let Some(var) = missing_var {
33        return Err(VarError::NotFound(var));
34    }
35
36    // Replace $VAR with env var
37    let result = simple_re.replace_all(&result, |caps: &regex::Captures| {
38        let var_name = &caps[1];
39        if let Ok(value) = env::var(var_name) {
40            value
41        } else {
42            missing_var = Some(var_name.to_string());
43            caps[0].to_string()
44        }
45    });
46    if let Some(var) = missing_var {
47        return Err(VarError::NotFound(var));
48    }
49
50    // Replace placeholder with $
51    let result = result.replace(ESCAPE_PLACEHOLDER, "$");
52
53    Ok(result)
54}
55
56#[derive(Debug)]
57pub enum VarError {
58    NotFound(String),
59}
60
61impl std::fmt::Display for VarError {
62    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
63        match self {
64            VarError::NotFound(name) => write!(f, "Environment variable '{name}' not found"),
65        }
66    }
67}
68
69impl std::error::Error for VarError {}
70
71#[cfg(test)]
72mod tests {
73    use super::*;
74
75    #[test]
76    fn test_simple_var() {
77        unsafe { env::set_var("TEST_VAR_SIMPLE", "hello") };
78        let result = expand_env_vars("$TEST_VAR_SIMPLE world").unwrap();
79        assert_eq!(result, "hello world");
80        unsafe { env::remove_var("TEST_VAR_SIMPLE") };
81    }
82
83    #[test]
84    fn test_bracketed_var() {
85        unsafe { env::set_var("TEST_VAR_BRACKET", "hello") };
86        let result = expand_env_vars("${TEST_VAR_BRACKET} world").unwrap();
87        assert_eq!(result, "hello world");
88        unsafe { env::remove_var("TEST_VAR_BRACKET") };
89    }
90
91    #[test]
92    fn test_escape_sequence() {
93        let result = expand_env_vars("$$VAR").unwrap();
94        assert_eq!(result, "$VAR");
95    }
96
97    #[test]
98    fn test_multiple_vars() {
99        unsafe {
100            env::set_var("VAR1", "hello");
101            env::set_var("VAR2", "world");
102        }
103        let result = expand_env_vars("$VAR1 ${VAR2}!").unwrap();
104        assert_eq!(result, "hello world!");
105        unsafe {
106            env::remove_var("VAR1");
107            env::remove_var("VAR2");
108        }
109    }
110
111    #[test]
112    fn test_missing_var() {
113        let result = expand_env_vars("$MISSING_VAR");
114        assert!(result.is_err());
115        match result {
116            Err(VarError::NotFound(name)) => assert_eq!(name, "MISSING_VAR"),
117            _ => panic!("Expected NotFound error"),
118        }
119    }
120
121    #[test]
122    fn test_unclosed_brace_left_as_is() {
123        // Unclosed braces are left as literal text since regex won't match
124        let result = expand_env_vars("${VAR").unwrap();
125        assert_eq!(result, "${VAR");
126    }
127
128    #[test]
129    fn test_empty_string() {
130        let result = expand_env_vars("").unwrap();
131        assert_eq!(result, "");
132    }
133
134    #[test]
135    fn test_no_vars() {
136        let result = expand_env_vars("plain text").unwrap();
137        assert_eq!(result, "plain text");
138    }
139
140    #[test]
141    fn test_dollar_at_end() {
142        let result = expand_env_vars("text$").unwrap();
143        assert_eq!(result, "text$");
144    }
145
146    #[test]
147    fn test_var_with_underscore() {
148        unsafe { env::set_var("MY_TEST_VAR", "value") };
149        let result = expand_env_vars("$MY_TEST_VAR").unwrap();
150        assert_eq!(result, "value");
151        unsafe { env::remove_var("MY_TEST_VAR") };
152    }
153
154    #[test]
155    fn test_var_with_numbers() {
156        unsafe { env::set_var("VAR123", "value") };
157        let result = expand_env_vars("$VAR123").unwrap();
158        assert_eq!(result, "value");
159        unsafe { env::remove_var("VAR123") };
160    }
161
162    #[test]
163    fn test_special_char_stops_var_name() {
164        unsafe { env::set_var("VAR", "value") };
165        let result = expand_env_vars("$VAR-suffix").unwrap();
166        assert_eq!(result, "value-suffix");
167        unsafe { env::remove_var("VAR") };
168    }
169}
mcp_utils/client/variables.rs

mcp_utils/client/
variables.rs
