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