feat: can now resolves args from stage

pkg/dockerfile/dockerfile.go

@@ -57,23 +57,29 @@ func Stages(opts *config.KanikoOptions) ([]config.KanikoStage, error) {
 	if err != nil {
 		return nil, errors.Wrap(err, "parsing dockerfile")
 	}
+
 	targetStage, err := targetStage(stages, opts.Target)
 	if err != nil {
 		return nil, err
 	}
 	resolveStages(stages)
+
+	args := unifyArgs(metaArgs, opts.BuildArgs)
+
+	if err := resolveStagesArgs(stages, args); err != nil {
+		return nil, errors.Wrap(err, "resolving args")
+	}
+
 	var kanikoStages []config.KanikoStage
 	for index, stage := range stages {
-		resolvedBaseName, err := util.ResolveEnvironmentReplacement(stage.BaseName, opts.BuildArgs, false)
-		if err != nil {
-			return nil, errors.Wrap(err, "resolving base name")
+		if len(stage.Name) > 0 {
+			logrus.Infof("Resolved base name %s to %s", stage.BaseName, stage.Name)
 		}
-		stage.Name = resolvedBaseName
-		logrus.Infof("Resolved base name %s to %s", stage.BaseName, stage.Name)
+		baseImageIndex := baseImageIndex(index, stages)
 		kanikoStages = append(kanikoStages, config.KanikoStage{
 			Stage:                  stage,
-			BaseImageIndex:         baseImageIndex(index, stages),
-			BaseImageStoredLocally: (baseImageIndex(index, stages) != -1),
+			BaseImageIndex:         baseImageIndex,
+			BaseImageStoredLocally: (baseImageIndex != -1),
 			SaveStage:              saveStage(index, stages),
 			Final:                  index == targetStage,
 			MetaArgs:               metaArgs,
@@ -87,6 +93,29 @@ func Stages(opts *config.KanikoOptions) ([]config.KanikoStage, error) {
 	return kanikoStages, nil
 }
 
+// unifyArgs returns the unified args between metaArgs and --build-arg
+// by default --build-arg overrides metaArgs except when --build-arg is empty
+func unifyArgs(metaArgs []instructions.ArgCommand, buildArgs []string) []string {
+	argsMap := make(map[string]string)
+	for _, a := range metaArgs {
+		if a.Value != nil {
+			argsMap[a.Key] = *a.Value
+		}
+	}
+	splitter := "="
+	for _, a := range buildArgs {
+		s := strings.Split(a, splitter)
+		if len(s) > 1 && s[1] != "" {
+			argsMap[s[0]] = s[1]
+		}
+	}
+	var args []string
+	for k, v := range argsMap {
+		args = append(args, fmt.Sprintf("%s=%s", k, v))
+	}
+	return args
+}
+
 // baseImageIndex returns the index of the stage the current stage is built off
 // returns -1 if the current stage isn't built off a previous stage
 func baseImageIndex(currentStage int, stages []instructions.Stage) int {
@@ -228,6 +257,21 @@ func resolveStages(stages []instructions.Stage) {
 	}
 }
 
+// resolveStagesArgs resolves all the args from list of stages
+// it returns a list of stages with all resolved args
+func resolveStagesArgs(stages []instructions.Stage, args []string) error {
+	for i, s := range stages {
+		resolvedBaseName, err := util.ResolveEnvironmentReplacement(s.BaseName, args, false)
+		if err != nil {
+			return errors.Wrap(err, fmt.Sprintf("resolving base name %s", s.BaseName))
+		}
+		if s.BaseName != resolvedBaseName {
+			stages[i].BaseName = resolvedBaseName
+		}
+	}
+	return nil
+}
+
 // ParseCommands parses an array of commands into an array of instructions.Command; used for onbuild
 func ParseCommands(cmdArray []string) ([]instructions.Command, error) {
 	var cmds []instructions.Command

pkg/dockerfile/dockerfile_test.go

@@ -17,6 +17,7 @@ limitations under the License.
 package dockerfile
 
 import (
+	"fmt"
 	"io/ioutil"
 	"os"
 	"strconv"
@@ -32,10 +33,10 @@ func Test_Stages_ArgValueWithQuotes(t *testing.T) {
 	ARG IMAGE="ubuntu:16.04"
 	FROM ${IMAGE}
 	RUN echo hi > /hi
-	
+
 	FROM scratch AS second
 	COPY --from=0 /hi /hi2
-	
+
 	FROM scratch
 	COPY --from=second /hi2 /hi3
 	`
@@ -193,10 +194,10 @@ func Test_resolveStages(t *testing.T) {
 	dockerfile := `
 	FROM scratch
 	RUN echo hi > /hi
-	
+
 	FROM scratch AS second
 	COPY --from=0 /hi /hi2
-	
+
 	FROM scratch AS tHiRd
 	COPY --from=second /hi2 /hi3
 	COPY --from=1 /hi2 /hi3
@@ -230,10 +231,10 @@ func Test_targetStage(t *testing.T) {
 	dockerfile := `
 	FROM scratch
 	RUN echo hi > /hi
-	
+
 	FROM scratch AS second
 	COPY --from=0 /hi /hi2
-	
+
 	FROM scratch
 	COPY --from=second /hi2 /hi3
 	`
@@ -364,3 +365,73 @@ func Test_baseImageIndex(t *testing.T) {
 		})
 	}
 }
+
+func Test_ResolveStagesArgs(t *testing.T) {
+	dockerfile := `
+	ARG IMAGE="ubuntu:16.04"
+	ARG LAST_STAGE_VARIANT
+	FROM ${IMAGE} as base
+	RUN echo hi > /hi
+	FROM base AS base-dev
+	RUN echo dev >> /hi
+	FROM base AS base-prod
+	RUN echo prod >> /hi
+	FROM base-${LAST_STAGE_VARIANT}
+	RUN cat /hi
+	`
+
+	buildArgLastVariants := []string{"dev", "prod"}
+	buildArgImages := []string{"alpine:3.11", ""}
+	var expectedImage string
+
+	for _, buildArgLastVariant := range buildArgLastVariants {
+		for _, buildArgImage := range buildArgImages {
+			if buildArgImage != "" {
+				expectedImage = buildArgImage
+			} else {
+				expectedImage = "ubuntu:16.04"
+			}
+			buildArgs := []string{fmt.Sprintf("IMAGE=%s", buildArgImage), fmt.Sprintf("LAST_STAGE_VARIANT=%s", buildArgLastVariant)}
+
+			stages, metaArgs, err := Parse([]byte(dockerfile))
+			if err != nil {
+				t.Fatal(err)
+			}
+			stagesLen := len(stages)
+			resolveStages(stages)
+
+			args := unifyArgs(metaArgs, buildArgs)
+			if err := resolveStagesArgs(stages, args); err != nil {
+				t.Fatalf("fail to resolves args %v: %v", buildArgs, err)
+			}
+			tests := []struct {
+				name               string
+				actualSourceCode   string
+				actualBaseName     string
+				expectedSourceCode string
+				expectedBaseName   string
+			}{
+				{
+					name:               "Test_BuildArg_From_First_Stage",
+					actualSourceCode:   stages[0].SourceCode,
+					actualBaseName:     stages[0].BaseName,
+					expectedSourceCode: "FROM ${IMAGE} as base",
+					expectedBaseName:   expectedImage,
+				},
+				{
+					name:               "Test_BuildArg_From_Last_Stage",
+					actualSourceCode:   stages[stagesLen-1].SourceCode,
+					actualBaseName:     stages[stagesLen-1].BaseName,
+					expectedSourceCode: "FROM base-${LAST_STAGE_VARIANT}",
+					expectedBaseName:   fmt.Sprintf("base-%s", buildArgLastVariant),
+				},
+			}
+			for _, test := range tests {
+				t.Run(test.name, func(t *testing.T) {
+					testutil.CheckDeepEqual(t, test.expectedSourceCode, test.actualSourceCode)
+					testutil.CheckDeepEqual(t, test.expectedBaseName, test.actualBaseName)
+				})
+			}
+		}
+	}
+}