Skip to main content

zlayer_builder/dockerfile/
variable.rs

1//! Variable expansion for Dockerfile ARG and ENV values
2//!
3//! This module provides functionality to expand variables in Dockerfile strings,
4//! supporting the following syntax:
5//!
6//! - `$VAR` - Simple variable reference
7//! - `${VAR}` - Explicit variable reference
8//! - `${VAR:-default}` - Use default if VAR is unset or empty
9//! - `${VAR:+alternate}` - Use alternate if VAR is set and non-empty
10//! - `${VAR-default}` - Use default if VAR is unset (empty string is valid)
11//! - `${VAR+alternate}` - Use alternate if VAR is set (including empty)
12
13use std::collections::HashMap;
14
15/// Expand variables in a string using the provided ARG and ENV maps.
16///
17/// Variables are expanded in the following order of precedence:
18/// 1. Build-time ARGs (provided at build time)
19/// 2. ENV variables (from Dockerfile ENV instructions)
20/// 3. Default ARG values (from Dockerfile ARG instructions)
21///
22/// # Arguments
23///
24/// * `input` - The string containing variable references
25/// * `args` - ARG variable values (build-time and defaults)
26/// * `env` - ENV variable values
27///
28/// # Returns
29///
30/// The string with all variables expanded
31pub fn expand_variables(
32    input: &str,
33    args: &HashMap<String, String>,
34    env: &HashMap<String, String>,
35) -> String {
36    let mut result = String::with_capacity(input.len());
37    let mut chars = input.chars().peekable();
38
39    while let Some(c) = chars.next() {
40        if c == '$' {
41            if let Some(&next) = chars.peek() {
42                if next == '{' {
43                    // ${VAR} or ${VAR:-default} or ${VAR:+alternate} format
44                    chars.next(); // consume '{'
45                    let (expanded, _) = expand_braced_variable(&mut chars, args, env);
46                    result.push_str(&expanded);
47                } else if next == '$' {
48                    // Escaped dollar sign
49                    chars.next();
50                    result.push('$');
51                } else if next.is_ascii_alphabetic() || next == '_' {
52                    // $VAR format
53                    let var_name = consume_var_name(&mut chars);
54                    if let Some(value) = lookup_variable(&var_name, args, env) {
55                        result.push_str(&value);
56                    }
57                    // If variable not found, expand to empty string (Docker behavior)
58                } else {
59                    // Just a dollar sign
60                    result.push(c);
61                }
62            } else {
63                result.push(c);
64            }
65        } else if c == '\\' {
66            // Check for escaped dollar sign
67            if let Some(&next) = chars.peek() {
68                if next == '$' {
69                    chars.next();
70                    result.push('$');
71                } else {
72                    result.push(c);
73                }
74            } else {
75                result.push(c);
76            }
77        } else {
78            result.push(c);
79        }
80    }
81
82    result
83}
84
85/// Consume a variable name from the character stream (for $VAR format)
86fn consume_var_name(chars: &mut std::iter::Peekable<std::str::Chars>) -> String {
87    let mut name = String::new();
88
89    while let Some(&c) = chars.peek() {
90        if c.is_ascii_alphanumeric() || c == '_' {
91            name.push(c);
92            chars.next();
93        } else {
94            break;
95        }
96    }
97
98    name
99}
100
101/// Expand a braced variable expression ${...}
102fn expand_braced_variable(
103    chars: &mut std::iter::Peekable<std::str::Chars>,
104    args: &HashMap<String, String>,
105    env: &HashMap<String, String>,
106) -> (String, bool) {
107    let mut var_name = String::new();
108    let mut operator = None;
109    let mut default_value = String::new();
110    let mut in_default = false;
111    let mut brace_depth = 1;
112
113    while let Some(c) = chars.next() {
114        if c == '}' {
115            brace_depth -= 1;
116            if brace_depth == 0 {
117                break;
118            }
119            if in_default {
120                default_value.push(c);
121            }
122        } else if c == '{' {
123            brace_depth += 1;
124            if in_default {
125                default_value.push(c);
126            }
127        } else if !in_default && (c == ':' || c == '-' || c == '+') {
128            // Check for modifier operators
129            if c == ':' {
130                if let Some(&next) = chars.peek() {
131                    if next == '-' || next == '+' {
132                        chars.next();
133                        operator = Some(format!(":{}", next));
134                        in_default = true;
135                        continue;
136                    }
137                }
138                var_name.push(c);
139            } else if c == '-' || c == '+' {
140                operator = Some(c.to_string());
141                in_default = true;
142            } else {
143                var_name.push(c);
144            }
145        } else if in_default {
146            default_value.push(c);
147        } else {
148            var_name.push(c);
149        }
150    }
151
152    // Look up the variable
153    let value = lookup_variable(&var_name, args, env);
154
155    match operator.as_deref() {
156        Some(":-") => {
157            // ${VAR:-default} - use default if unset OR empty
158            match value {
159                Some(v) if !v.is_empty() => (v, true),
160                _ => {
161                    // Recursively expand the default value
162                    (expand_variables(&default_value, args, env), false)
163                }
164            }
165        }
166        Some("-") => {
167            // ${VAR-default} - use default only if unset (empty is valid)
168            match value {
169                Some(v) => (v, true),
170                None => (expand_variables(&default_value, args, env), false),
171            }
172        }
173        Some(":+") => {
174            // ${VAR:+alternate} - use alternate if set AND non-empty
175            match value {
176                Some(v) if !v.is_empty() => (expand_variables(&default_value, args, env), true),
177                _ => (String::new(), false),
178            }
179        }
180        Some("+") => {
181            // ${VAR+alternate} - use alternate if set (including empty)
182            match value {
183                Some(_) => (expand_variables(&default_value, args, env), true),
184                None => (String::new(), false),
185            }
186        }
187        None | Some(_) => {
188            // Simple ${VAR}
189            let is_set = value.is_some();
190            (value.unwrap_or_default(), is_set)
191        }
192    }
193}
194
195/// Look up a variable in the provided maps.
196/// Checks ENV first (higher precedence), then ARGs.
197fn lookup_variable(
198    name: &str,
199    args: &HashMap<String, String>,
200    env: &HashMap<String, String>,
201) -> Option<String> {
202    // ENV takes precedence over ARG
203    env.get(name).cloned().or_else(|| args.get(name).cloned())
204}
205
206/// Expands variables in a list of strings
207pub fn expand_variables_in_list(
208    inputs: &[String],
209    args: &HashMap<String, String>,
210    env: &HashMap<String, String>,
211) -> Vec<String> {
212    inputs
213        .iter()
214        .map(|s| expand_variables(s, args, env))
215        .collect()
216}
217
218/// Builder for tracking variable state during Dockerfile processing
219#[derive(Debug, Default, Clone)]
220pub struct VariableContext {
221    /// Build-time ARG values (from --build-arg)
222    pub build_args: HashMap<String, String>,
223
224    /// ARG default values from Dockerfile
225    pub arg_defaults: HashMap<String, String>,
226
227    /// ENV values accumulated during processing
228    pub env_vars: HashMap<String, String>,
229}
230
231impl VariableContext {
232    /// Create a new variable context
233    pub fn new() -> Self {
234        Self::default()
235    }
236
237    /// Create with build-time arguments
238    pub fn with_build_args(build_args: HashMap<String, String>) -> Self {
239        Self {
240            build_args,
241            ..Default::default()
242        }
243    }
244
245    /// Register an ARG with optional default
246    pub fn add_arg(&mut self, name: impl Into<String>, default: Option<String>) {
247        let name = name.into();
248        if let Some(default) = default {
249            self.arg_defaults.insert(name, default);
250        }
251    }
252
253    /// Set an ENV variable
254    pub fn set_env(&mut self, name: impl Into<String>, value: impl Into<String>) {
255        self.env_vars.insert(name.into(), value.into());
256    }
257
258    /// Get the effective ARG values (build-time overrides defaults)
259    pub fn effective_args(&self) -> HashMap<String, String> {
260        let mut result = self.arg_defaults.clone();
261        for (k, v) in &self.build_args {
262            result.insert(k.clone(), v.clone());
263        }
264        result
265    }
266
267    /// Expand variables in a string using the current context
268    pub fn expand(&self, input: &str) -> String {
269        expand_variables(input, &self.effective_args(), &self.env_vars)
270    }
271
272    /// Expand variables in a list of strings
273    pub fn expand_list(&self, inputs: &[String]) -> Vec<String> {
274        expand_variables_in_list(inputs, &self.effective_args(), &self.env_vars)
275    }
276}
277
278#[cfg(test)]
279mod tests {
280    use super::*;
281
282    #[test]
283    fn test_simple_variable() {
284        let mut args = HashMap::new();
285        args.insert("VERSION".to_string(), "1.0".to_string());
286        let env = HashMap::new();
287
288        assert_eq!(expand_variables("$VERSION", &args, &env), "1.0");
289        assert_eq!(expand_variables("${VERSION}", &args, &env), "1.0");
290        assert_eq!(expand_variables("v$VERSION", &args, &env), "v1.0");
291        assert_eq!(
292            expand_variables("v${VERSION}-release", &args, &env),
293            "v1.0-release"
294        );
295    }
296
297    #[test]
298    fn test_undefined_variable() {
299        let args = HashMap::new();
300        let env = HashMap::new();
301
302        assert_eq!(expand_variables("$UNDEFINED", &args, &env), "");
303        assert_eq!(expand_variables("${UNDEFINED}", &args, &env), "");
304    }
305
306    #[test]
307    fn test_default_value_colon_minus() {
308        let args = HashMap::new();
309        let env = HashMap::new();
310
311        // Unset variable with default
312        assert_eq!(expand_variables("${VERSION:-1.0}", &args, &env), "1.0");
313
314        // Set variable ignores default
315        let mut args = HashMap::new();
316        args.insert("VERSION".to_string(), "2.0".to_string());
317        assert_eq!(expand_variables("${VERSION:-1.0}", &args, &env), "2.0");
318
319        // Empty variable uses default (colon variant)
320        let mut args = HashMap::new();
321        args.insert("VERSION".to_string(), "".to_string());
322        assert_eq!(expand_variables("${VERSION:-1.0}", &args, &env), "1.0");
323    }
324
325    #[test]
326    fn test_default_value_minus() {
327        let args = HashMap::new();
328        let env = HashMap::new();
329
330        // Unset variable with default
331        assert_eq!(expand_variables("${VERSION-1.0}", &args, &env), "1.0");
332
333        // Empty variable keeps empty (non-colon variant)
334        let mut args = HashMap::new();
335        args.insert("VERSION".to_string(), "".to_string());
336        assert_eq!(expand_variables("${VERSION-1.0}", &args, &env), "");
337    }
338
339    #[test]
340    fn test_alternate_value_colon_plus() {
341        let mut args = HashMap::new();
342        let env = HashMap::new();
343
344        // Unset variable - no alternate
345        assert_eq!(expand_variables("${VERSION:+set}", &args, &env), "");
346
347        // Set non-empty variable - use alternate
348        args.insert("VERSION".to_string(), "1.0".to_string());
349        assert_eq!(expand_variables("${VERSION:+set}", &args, &env), "set");
350
351        // Set empty variable - no alternate (colon variant)
352        args.insert("VERSION".to_string(), "".to_string());
353        assert_eq!(expand_variables("${VERSION:+set}", &args, &env), "");
354    }
355
356    #[test]
357    fn test_alternate_value_plus() {
358        let mut args = HashMap::new();
359        let env = HashMap::new();
360
361        // Unset variable - no alternate
362        assert_eq!(expand_variables("${VERSION+set}", &args, &env), "");
363
364        // Set empty variable - use alternate (non-colon variant)
365        args.insert("VERSION".to_string(), "".to_string());
366        assert_eq!(expand_variables("${VERSION+set}", &args, &env), "set");
367    }
368
369    #[test]
370    fn test_env_takes_precedence() {
371        let mut args = HashMap::new();
372        args.insert("VAR".to_string(), "from_arg".to_string());
373
374        let mut env = HashMap::new();
375        env.insert("VAR".to_string(), "from_env".to_string());
376
377        assert_eq!(expand_variables("$VAR", &args, &env), "from_env");
378    }
379
380    #[test]
381    fn test_escaped_dollar() {
382        let args = HashMap::new();
383        let env = HashMap::new();
384
385        assert_eq!(expand_variables("\\$VAR", &args, &env), "$VAR");
386        assert_eq!(expand_variables("$$", &args, &env), "$");
387    }
388
389    #[test]
390    fn test_nested_default() {
391        let mut args = HashMap::new();
392        args.insert("DEFAULT".to_string(), "nested".to_string());
393        let env = HashMap::new();
394
395        // Default value contains a variable reference
396        assert_eq!(
397            expand_variables("${UNSET:-$DEFAULT}", &args, &env),
398            "nested"
399        );
400    }
401
402    #[test]
403    fn test_variable_context() {
404        let mut ctx = VariableContext::with_build_args({
405            let mut m = HashMap::new();
406            m.insert("BUILD_TYPE".to_string(), "release".to_string());
407            m
408        });
409
410        ctx.add_arg("VERSION", Some("1.0".to_string()));
411        ctx.set_env("HOME", "/app".to_string());
412
413        assert_eq!(ctx.expand("$BUILD_TYPE"), "release");
414        assert_eq!(ctx.expand("$VERSION"), "1.0");
415        assert_eq!(ctx.expand("$HOME"), "/app");
416    }
417
418    #[test]
419    fn test_build_arg_overrides_default() {
420        let mut ctx = VariableContext::with_build_args({
421            let mut m = HashMap::new();
422            m.insert("VERSION".to_string(), "2.0".to_string());
423            m
424        });
425
426        ctx.add_arg("VERSION", Some("1.0".to_string()));
427
428        // Build arg should override default
429        assert_eq!(ctx.expand("$VERSION"), "2.0");
430    }
431
432    #[test]
433    fn test_complex_string() {
434        let mut args = HashMap::new();
435        args.insert("APP".to_string(), "myapp".to_string());
436        args.insert("VERSION".to_string(), "1.2.3".to_string());
437        let env = HashMap::new();
438
439        let input = "FROM registry.example.com/${APP}:${VERSION:-latest}";
440        assert_eq!(
441            expand_variables(input, &args, &env),
442            "FROM registry.example.com/myapp:1.2.3"
443        );
444    }
445}