From 0a8fde33cbe2efbb573ce4060356404ef28d731e Mon Sep 17 00:00:00 2001
From: Talon Bowler <talon.bowler@docker.com>
Date: Tue, 7 May 2024 17:00:46 -0700
Subject: [PATCH] call validateCommandVars from within the command.Expand call

Signed-off-by: Talon Bowler <talon.bowler@docker.com>
---
 frontend/dockerfile/dockerfile2llb/convert.go | 45 ++++++++++---------
 1 file changed, 25 insertions(+), 20 deletions(-)

diff --git a/frontend/dockerfile/dockerfile2llb/convert.go b/frontend/dockerfile/dockerfile2llb/convert.go
index 0d781029ba88b..b4bd8933de8d3 100644
--- a/frontend/dockerfile/dockerfile2llb/convert.go
+++ b/frontend/dockerfile/dockerfile2llb/convert.go
@@ -753,13 +753,15 @@ type dispatchOpt struct {
 }
 
 func dispatch(d *dispatchState, cmd command, opt dispatchOpt) error {
-	env, err := d.state.Env(context.TODO())
-	if err != nil {
-		return err
-	}
+	var err error
 	if ex, ok := cmd.Command.(instructions.SupportsSingleWordExpansion); ok {
 		err := ex.Expand(func(word string) (string, error) {
-			newword, _, err := opt.shlex.ProcessWord(word, env)
+			env, err := d.state.Env(context.TODO())
+			if err != nil {
+				return "", err
+			}
+			newword, unmatched, err := opt.shlex.ProcessWord(word, env)
+			validateCommandVars(cmd, unmatched, &opt)
 			return newword, err
 		})
 		if err != nil {
@@ -768,16 +770,20 @@ func dispatch(d *dispatchState, cmd command, opt dispatchOpt) error {
 	}
 	if ex, ok := cmd.Command.(instructions.SupportsSingleWordExpansionRaw); ok {
 		err := ex.ExpandRaw(func(word string) (string, error) {
+			env, err := d.state.Env(context.TODO())
+			if err != nil {
+				return "", err
+			}
 			lex := shell.NewLex('\\')
 			lex.SkipProcessQuotes = true
-			newword, _, err := lex.ProcessWord(word, env)
+			newword, unmatched, err := lex.ProcessWord(word, env)
+			validateCommandVars(cmd, unmatched, &opt)
 			return newword, err
 		})
 		if err != nil {
 			return err
 		}
 	}
-	validateCommandVar(cmd.Command, env, &opt)
 
 	switch c := cmd.Command.(type) {
 	case *instructions.MaintainerCommand:
@@ -837,7 +843,7 @@ func dispatch(d *dispatchState, cmd command, opt dispatchOpt) error {
 	case *instructions.ShellCommand:
 		err = dispatchShell(d, c)
 	case *instructions.ArgCommand:
-		err = dispatchArg(d, c, opt.metaArgs, &opt)
+		err = dispatchArg(d, c, &opt)
 	case *instructions.CopyCommand:
 		l := opt.buildContext
 		if len(cmd.sources) != 0 {
@@ -1546,7 +1552,7 @@ func dispatchShell(d *dispatchState, c *instructions.ShellCommand) error {
 	return commitToHistory(&d.image, fmt.Sprintf("SHELL %v", c.Shell), false, nil, d.epoch)
 }
 
-func dispatchArg(d *dispatchState, c *instructions.ArgCommand, metaArgs []instructions.KeyValuePairOptional, opt *dispatchOpt) error {
+func dispatchArg(d *dispatchState, c *instructions.ArgCommand, opt *dispatchOpt) error {
 	commitStrs := make([]string, 0, len(c.Args))
 	for _, arg := range c.Args {
 		opt.argCmdVars[arg.Key] = struct{}{}
@@ -1560,7 +1566,7 @@ func dispatchArg(d *dispatchState, c *instructions.ArgCommand, metaArgs []instru
 
 		skipArgInfo := false // skip the arg info if the arg is inherited from global scope
 		if buildArg.Value == nil {
-			for _, ma := range metaArgs {
+			for _, ma := range opt.metaArgs {
 				if ma.Key == buildArg.Key {
 					buildArg.Value = ma.Value
 					skipArgInfo = true
@@ -1573,6 +1579,8 @@ func dispatchArg(d *dispatchState, c *instructions.ArgCommand, metaArgs []instru
 		if buildArg.Value != nil {
 			if _, ok := nonEnvArgs[buildArg.Key]; !ok {
 				d.state = d.state.AddEnv(buildArg.Key, *buildArg.Value)
+			} else {
+				d.state = d.state.AddEnv(buildArg.Key, "")
 			}
 			ai.value = *buildArg.Value
 		}
@@ -2046,16 +2054,13 @@ func validateStageNames(stages []instructions.Stage, warn linter.LintWarnFunc) {
 	}
 }
 
-func validateCommandVar(cmd instructions.Command, env []string, opt *dispatchOpt) {
-	if cmdstr, ok := cmd.(fmt.Stringer); ok {
-		_, unmatched, _ := opt.shlex.ProcessWord(cmdstr.String(), env)
-		for arg := range unmatched {
-			_, argCmdOk := opt.argCmdVars[arg]
-			_, nonEnvOk := nonEnvArgs[arg]
-			if !argCmdOk && !nonEnvOk {
-				msg := linter.RuleUndefinedVar.Format(arg)
-				linter.RuleUndefinedVar.Run(opt.warn, cmd.Location(), msg)
-			}
+func validateCommandVars(cmd instructions.Command, unmatched map[string]struct{}, opt *dispatchOpt) {
+	for cmdVar := range unmatched {
+		_, argCmdOk := opt.argCmdVars[cmdVar]
+		_, nonEnvOk := nonEnvArgs[cmdVar]
+		if !argCmdOk && !nonEnvOk {
+			msg := linter.RuleUndefinedVar.Format(cmdVar)
+			linter.RuleUndefinedVar.Run(opt.lintWarn, cmd.Location(), msg)
 		}
 	}
 }