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