supermachine 0.7.72

//! Dockerfile build-time variable substitution.
//!
//! Replaces `$VAR` / `${VAR}` (and the `${VAR:-default}` / `${VAR:+alt}`
//! modifiers, with `\$` as a literal `$`) in instruction arguments, using the
//! in-scope `ARG` + `ENV` values — Docker's build-arg semantics. `ENV` takes
//! precedence over `ARG`; both accumulate as the stage is walked, so a later
//! instruction sees values declared by earlier `ARG`/`ENV` lines.
//!
//! Substitution is resolved ONCE up front ([`resolve_stage_instructions`]),
//! producing an instruction list with every `$VAR` already expanded. Everything
//! downstream — the content-addressed layer cache key AND the executor — then
//! operates on the resolved instructions, so a changed `ARG`/`ENV` value flows
//! into the cache key (correct invalidation) and into the executed command.
//!
//! Scope note (v1): the base image's own `ENV` is not seeded into the scope, and
//! global `ARG`s (declared before `FROM`) are made visible to every stage rather
//! than only after an in-stage `ARG name` re-declaration. Both are minor, safe
//! deviations from Docker that cover the common cases; documented for later.

use std::collections::BTreeMap;

use super::dockerfile::{CopyFlags, Instruction, RunMount, ShellOrExec};

/// Expand `$VAR` / `${VAR}` / `${VAR:-default}` / `${VAR:+alt}` in `s`.
/// `lookup(name)` returns the variable's value (`None` = unset). Unset/empty
/// vars expand to `""` (or the `:-` default). `\$` is a literal `$`.
pub(crate) fn expand(s: &str, lookup: &dyn Fn(&str) -> Option<String>) -> String {
    let b = s.as_bytes();
    let mut out = String::with_capacity(s.len());
    let mut i = 0;
    while i < b.len() {
        let c = b[i];
        if c == b'\\' && i + 1 < b.len() && b[i + 1] == b'$' {
            out.push('$');
            i += 2;
            continue;
        }
        if c != b'$' {
            out.push(c as char);
            i += 1;
            continue;
        }
        // c == '$'
        if i + 1 < b.len() && b[i + 1] == b'{' {
            // ${ ... } — find the matching close brace.
            if let Some(close) = find_byte(b, i + 2, b'}') {
                let inner = &s[i + 2..close];
                match expand_braced(inner, lookup) {
                    Some(v) => out.push_str(&v),
                    // Undeclared `${VAR}` with no modifier: pass through
                    // LITERALLY (Docker substitutes only declared ARG/ENV;
                    // the shell sees the rest). Blanking it broke
                    // `RUN X=hi && echo ${X}`.
                    None => {
                        out.push_str("${");
                        out.push_str(inner);
                        out.push('}');
                    }
                }
                i = close + 1;
                continue;
            }
            // Unterminated `${` — emit literally.
            out.push('$');
            i += 1;
            continue;
        }
        // $name
        let start = i + 1;
        let mut j = start;
        while j < b.len() && is_name_byte(b[j], j == start) {
            j += 1;
        }
        if j == start {
            // Lone `$` (e.g. end of string or `$` before non-name char).
            out.push('$');
            i += 1;
            continue;
        }
        let name = &s[start..j];
        match lookup(name) {
            Some(v) => out.push_str(&v),
            // Undeclared `$name`: pass through LITERALLY so a shell-assigned
            // variable in the same RUN works (Docker semantics — substitution
            // applies only to declared ARG/ENV). Blanking it turned
            // `RUN D=/x && mkdir "$D"` into `mkdir ""`.
            None => {
                out.push('$');
                out.push_str(name);
            }
        }
        i = j;
    }
    out
}

/// Expand the inside of `${ ... }`: `name`, `name:-word`, `name:+word`,
/// `name-word`, `name+word`. `word` is itself expanded (one level of nesting).
///
/// Returns `None` for a plain `${name}` (or an unrecognized form) whose `name`
/// is UNDECLARED — the caller then emits the `${name}` text literally, matching
/// Docker (build-time substitution only touches declared ARG/ENV; everything
/// else is left for the shell). The `:-` / `:+` / `-` / `+` modifier forms are
/// always resolved (Docker processes them at build time), so an undeclared
/// `${X:-d}` still yields `d`.
fn expand_braced(inner: &str, lookup: &dyn Fn(&str) -> Option<String>) -> Option<String> {
    // Split name from an optional modifier (`:-`, `:+`, `-`, `+`).
    let name_end = inner.find([':', '-', '+']).unwrap_or(inner.len());
    let name = &inner[..name_end];
    let rest = &inner[name_end..];
    let val = lookup(name);
    let set_nonempty = val.as_deref().map(|v| !v.is_empty()).unwrap_or(false);

    if rest.is_empty() {
        // Plain `${name}`: declared → its value (even if empty); undeclared →
        // None so the caller keeps `${name}` literal.
        return val;
    }
    // Modifier: optional leading ':' then '-' or '+'.
    let rest = rest.strip_prefix(':').unwrap_or(rest);
    let (op, word) = match rest.as_bytes().first() {
        Some(b'-') => ('-', &rest[1..]),
        Some(b'+') => ('+', &rest[1..]),
        // Unrecognized form — declared → its value; undeclared → literal.
        _ => return val,
    };
    Some(match op {
        '-' if set_nonempty => val.unwrap_or_default(),
        '-' => expand(word, lookup),
        '+' if set_nonempty => expand(word, lookup),
        _ => String::new(),
    })
}

fn find_byte(b: &[u8], from: usize, target: u8) -> Option<usize> {
    (from..b.len()).find(|&k| b[k] == target)
}

fn is_name_byte(c: u8, first: bool) -> bool {
    c == b'_' || c.is_ascii_alphabetic() || (!first && c.is_ascii_digit())
}

/// Resolve all `$VAR` substitutions in a stage's instructions, tracking the
/// `ARG`/`ENV` scope as it walks. `global_args` (declared before the first
/// `FROM`) seed the arg scope; `build_args` (`--build-arg`) override `ARG`
/// defaults. Returns the instructions with every supported field expanded.
pub(crate) fn resolve_stage_instructions(
    instrs: &[Instruction],
    global_args: &[(String, Option<String>)],
    build_args: &BTreeMap<String, String>,
) -> Vec<Instruction> {
    let mut env: BTreeMap<String, String> = BTreeMap::new();
    let mut args: BTreeMap<String, String> = BTreeMap::new();
    for (name, default) in global_args {
        if let Some(v) = build_args.get(name).cloned().or_else(|| default.clone()) {
            args.insert(name.clone(), v);
        }
    }

    let mut out = Vec::with_capacity(instrs.len());
    for instr in instrs {
        match instr {
            Instruction::Arg { name, default } => {
                // `--build-arg` wins; else the (expanded) default; else unset→"".
                let resolved = build_args
                    .get(name)
                    .cloned()
                    .or_else(|| default.as_ref().map(|d| expand(d, &lk(&env, &args))));
                args.insert(name.clone(), resolved.clone().unwrap_or_default());
                // Carry the effective value so the cache key reflects it.
                out.push(Instruction::Arg {
                    name: name.clone(),
                    default: resolved,
                });
            }
            Instruction::Env(pairs) => {
                let mut new_pairs = Vec::with_capacity(pairs.len());
                for (k, v) in pairs {
                    let ev = expand(v, &lk(&env, &args));
                    env.insert(k.clone(), ev.clone());
                    new_pairs.push((k.clone(), ev));
                }
                out.push(Instruction::Env(new_pairs));
            }
            other => out.push(expand_instr(other, &lk(&env, &args))),
        }
    }
    out
}

/// Build the `ENV`-then-`ARG` lookup closure for the current scope.
fn lk<'a>(
    env: &'a BTreeMap<String, String>,
    args: &'a BTreeMap<String, String>,
) -> impl Fn(&str) -> Option<String> + 'a {
    move |name: &str| env.get(name).or_else(|| args.get(name)).cloned()
}

/// Expand the substitutable fields of a single (non-ARG, non-ENV) instruction.
fn expand_instr(instr: &Instruction, lookup: &dyn Fn(&str) -> Option<String>) -> Instruction {
    let ev = |s: &str| expand(s, lookup);
    let evv = |v: &[String]| v.iter().map(|s| ev(s)).collect::<Vec<_>>();
    match instr {
        Instruction::Run { run, mounts } => Instruction::Run {
            run: expand_soe(run, lookup),
            mounts: mounts.iter().map(|m| expand_mount(m, lookup)).collect(),
        },
        Instruction::Copy {
            sources,
            dest,
            flags,
        } => Instruction::Copy {
            sources: evv(sources),
            dest: ev(dest),
            flags: expand_flags(flags, lookup),
        },
        Instruction::Add {
            sources,
            dest,
            flags,
        } => Instruction::Add {
            sources: evv(sources),
            dest: ev(dest),
            flags: expand_flags(flags, lookup),
        },
        Instruction::Workdir(d) => Instruction::Workdir(ev(d)),
        Instruction::User(u) => Instruction::User(ev(u)),
        Instruction::Expose(p) => Instruction::Expose(evv(p)),
        Instruction::Label(pairs) => {
            Instruction::Label(pairs.iter().map(|(k, v)| (k.clone(), ev(v))).collect())
        }
        Instruction::Entrypoint(e) => Instruction::Entrypoint(expand_soe(e, lookup)),
        Instruction::Cmd(c) => Instruction::Cmd(expand_soe(c, lookup)),
        Instruction::Volume(v) => Instruction::Volume(evv(v)),
        Instruction::StopSignal(s) => Instruction::StopSignal(ev(s)),
        // SHELL is the interpreter itself — not substituted. ARG/ENV handled
        // by the caller.
        Instruction::Shell(_) | Instruction::Arg { .. } | Instruction::Env(_) => instr.clone(),
    }
}

fn expand_soe(soe: &ShellOrExec, lookup: &dyn Fn(&str) -> Option<String>) -> ShellOrExec {
    match soe {
        ShellOrExec::Shell(s) => ShellOrExec::Shell(expand(s, lookup)),
        ShellOrExec::Exec(a) => ShellOrExec::Exec(a.iter().map(|s| expand(s, lookup)).collect()),
    }
}

fn expand_flags(f: &CopyFlags, lookup: &dyn Fn(&str) -> Option<String>) -> CopyFlags {
    CopyFlags {
        // `--from` is a stage ref, not substituted.
        from: f.from.clone(),
        chown: f.chown.as_ref().map(|s| expand(s, lookup)),
        chmod: f.chmod.as_ref().map(|s| expand(s, lookup)),
    }
}

fn expand_mount(m: &RunMount, lookup: &dyn Fn(&str) -> Option<String>) -> RunMount {
    let e = |o: &Option<String>| o.as_ref().map(|s| expand(s, lookup));
    RunMount {
        kind: m.kind.clone(),
        target: e(&m.target),
        source: e(&m.source),
        id: e(&m.id),
        from: m.from.clone(),
        readonly: m.readonly,
        required: m.required,
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn map(pairs: &[(&str, &str)]) -> BTreeMap<String, String> {
        pairs
            .iter()
            .map(|(k, v)| (k.to_string(), v.to_string()))
            .collect()
    }

    fn lookup_from(m: &BTreeMap<String, String>) -> impl Fn(&str) -> Option<String> + '_ {
        move |n: &str| m.get(n).cloned()
    }

    #[test]
    fn basic_forms() {
        let m = map(&[("V", "1.2"), ("EMPTY", "")]);
        let l = lookup_from(&m);
        assert_eq!(expand("v$V", &l), "v1.2");
        assert_eq!(expand("v${V}x", &l), "v1.2x");
        // A declared-but-empty var still substitutes (to empty).
        assert_eq!(expand("a${EMPTY}b", &l), "ab");
        // UNDECLARED vars pass through literally (Docker: only declared
        // ARG/ENV are substituted; the rest is left for the shell).
        assert_eq!(expand("$MISSING.", &l), "$MISSING.");
        assert_eq!(expand("a${MISSING}b", &l), "a${MISSING}b");
    }

    #[test]
    fn undeclared_vars_pass_through_to_the_shell() {
        // The bug: a shell-assigned variable used in the same RUN got blanked
        // because the builder substituted the undeclared `$VAR` to "" before
        // the shell ran. Docker leaves undeclared vars literal.
        let m = map(&[("V", "1")]);
        let l = lookup_from(&m);
        // Shell-var idiom must survive substitution verbatim.
        assert_eq!(
            expand(r#"D="/root/.cache/x/bin" && mkdir -p "$D""#, &l),
            r#"D="/root/.cache/x/bin" && mkdir -p "$D""#
        );
        assert_eq!(expand("X=hi && echo [$X]", &l), "X=hi && echo [$X]");
        assert_eq!(expand("echo ${X}", &l), "echo ${X}");
        // Declared vars still substitute, even when mixed with undeclared ones.
        assert_eq!(expand("v$V then $X", &l), "v1 then $X");
        assert_eq!(expand("${V}-${X}", &l), "1-${X}");
        // `:-` / `:+` modifiers are still resolved for undeclared names
        // (Docker processes them at build time).
        assert_eq!(expand("${X:-fallback}", &l), "fallback");
        assert_eq!(expand("${X:+set}", &l), "");
        // Escapes still win over substitution.
        assert_eq!(expand(r"\$X stays", &l), "$X stays");
    }

    #[test]
    fn default_and_alt_modifiers() {
        let m = map(&[("SET", "x"), ("EMPTY", "")]);
        let l = lookup_from(&m);
        assert_eq!(expand("${SET:-d}", &l), "x");
        assert_eq!(expand("${EMPTY:-d}", &l), "d");
        assert_eq!(expand("${MISSING:-d}", &l), "d");
        assert_eq!(expand("${SET:+yes}", &l), "yes");
        assert_eq!(expand("${EMPTY:+yes}", &l), "");
        assert_eq!(expand("${MISSING:+yes}", &l), "");
    }

    #[test]
    fn escaped_dollar_is_literal() {
        let m = map(&[("V", "1")]);
        let l = lookup_from(&m);
        assert_eq!(expand(r"price \$5 not $V", &l), "price $5 not 1");
    }

    #[test]
    fn arg_default_then_env_precedence_and_accumulation() {
        // ARG VERSION=1.21  → RUN sees 1.21
        // ENV GOPATH=/go     → later RUN sees /go
        // ENV overrides ARG of the same name for substitution.
        let instrs = vec![
            Instruction::Arg {
                name: "VERSION".into(),
                default: Some("1.21".into()),
            },
            Instruction::run(ShellOrExec::Shell("wget go${VERSION}.tgz".into())),
            Instruction::Env(vec![("DIR".into(), "/opt/${VERSION}".into())]),
            Instruction::run(ShellOrExec::Shell("ls $DIR".into())),
            Instruction::Arg {
                name: "VERSION".into(),
                default: Some("ignored".into()),
            },
            Instruction::Env(vec![("VERSION".into(), "9".into())]),
            Instruction::run(ShellOrExec::Shell("echo ${VERSION}".into())),
        ];
        let out = resolve_stage_instructions(&instrs, &[], &BTreeMap::new());
        let shells: Vec<String> = out
            .iter()
            .filter_map(|i| match i {
                Instruction::Run {
                    run: ShellOrExec::Shell(s),
                    ..
                } => Some(s.clone()),
                _ => None,
            })
            .collect();
        assert_eq!(shells[0], "wget go1.21.tgz");
        assert_eq!(shells[1], "ls /opt/1.21");
        // After `ENV VERSION=9`, ENV beats the re-declared ARG default.
        assert_eq!(shells[2], "echo 9");
    }

    #[test]
    fn build_arg_override_beats_default() {
        let instrs = vec![
            Instruction::Arg {
                name: "TAG".into(),
                default: Some("latest".into()),
            },
            Instruction::run(ShellOrExec::Shell("pull img:$TAG".into())),
        ];
        let out = resolve_stage_instructions(&instrs, &[], &map(&[("TAG", "v2")]));
        match &out[1] {
            Instruction::Run {
                run: ShellOrExec::Shell(s),
                ..
            } => assert_eq!(s, "pull img:v2"),
            _ => panic!(),
        }
    }

    #[test]
    fn global_args_visible_in_stage() {
        let instrs = vec![Instruction::run(ShellOrExec::Shell("echo $G".into()))];
        let out = resolve_stage_instructions(
            &instrs,
            &[("G".into(), Some("glob".into()))],
            &BTreeMap::new(),
        );
        match &out[0] {
            Instruction::Run {
                run: ShellOrExec::Shell(s),
                ..
            } => assert_eq!(s, "echo glob"),
            _ => panic!(),
        }
    }

    #[test]
    fn copy_and_workdir_and_label_expand() {
        let instrs = vec![
            Instruction::Env(vec![("APP".into(), "myapp".into())]),
            Instruction::Workdir("/srv/$APP".into()),
            Instruction::Label(vec![("app".into(), "${APP}-prod".into())]),
        ];
        let out = resolve_stage_instructions(&instrs, &[], &BTreeMap::new());
        assert_eq!(out[1], Instruction::Workdir("/srv/myapp".into()));
        assert_eq!(
            out[2],
            Instruction::Label(vec![("app".into(), "myapp-prod".into())])
        );
    }

    #[test]
    fn exec_form_run_and_adjacent_vars() {
        let m = map(&[("A", "x"), ("B", "y")]);
        let l = lookup_from(&m);
        // Adjacent + braced-adjacent expansions.
        assert_eq!(expand("$A$B", &l), "xy");
        assert_eq!(expand("${A}${B}z", &l), "xyz");
        // Exec-form RUN expands each argv element.
        let instrs = vec![
            Instruction::Env(vec![("BIN".into(), "mytool".into())]),
            Instruction::run(ShellOrExec::Exec(vec![
                "/usr/bin/$BIN".into(),
                "--out=${BIN}.log".into(),
            ])),
        ];
        let out = resolve_stage_instructions(&instrs, &[], &BTreeMap::new());
        assert_eq!(
            out[1],
            Instruction::run(ShellOrExec::Exec(vec![
                "/usr/bin/mytool".into(),
                "--out=mytool.log".into(),
            ]))
        );
    }

    #[test]
    fn copy_chown_and_run_mount_fields_expand() {
        use super::super::dockerfile::{CopyFlags, MountKind, RunMount};
        let instrs = vec![
            Instruction::Arg {
                name: "U".into(),
                default: Some("appuser".into()),
            },
            Instruction::Arg {
                name: "CACHE".into(),
                default: Some("/var/cache/$U".into()),
            },
            Instruction::Copy {
                sources: vec!["src".into()],
                dest: "/dst".into(),
                flags: CopyFlags {
                    from: None,
                    chown: Some("$U:$U".into()),
                    chmod: None,
                },
            },
            Instruction::Run {
                run: ShellOrExec::Shell("make".into()),
                mounts: vec![RunMount {
                    kind: MountKind::Cache,
                    target: Some("${CACHE}".into()),
                    source: None,
                    id: Some("c-$U".into()),
                    from: None,
                    readonly: false,
                    required: false,
                }],
            },
        ];
        let out = resolve_stage_instructions(&instrs, &[], &BTreeMap::new());
        match &out[2] {
            Instruction::Copy { flags, .. } => {
                assert_eq!(flags.chown.as_deref(), Some("appuser:appuser"))
            }
            _ => panic!("expected COPY"),
        }
        match &out[3] {
            Instruction::Run { mounts, .. } => {
                // ${CACHE} itself expanded $U when the ARG was declared.
                assert_eq!(mounts[0].target.as_deref(), Some("/var/cache/appuser"));
                assert_eq!(mounts[0].id.as_deref(), Some("c-appuser"));
            }
            _ => panic!("expected RUN"),
        }
    }

    #[test]
    fn unterminated_brace_and_lone_dollar_are_literal() {
        let m = map(&[("V", "1")]);
        let l = lookup_from(&m);
        assert_eq!(expand("a${V", &l), "a${V");
        assert_eq!(expand("cost: $ 5", &l), "cost: $ 5");
        assert_eq!(expand("trailing$", &l), "trailing$");
    }
}