feat: can now resolves args from stage

pkg/dockerfile/dockerfile.go

@@ -57,24 +57,31 @@ 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)
+
+	stagesResolvedArgs, err := resolveStagesArgs(stages, args)
+	if 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(stagesResolvedArgs[index].Name) > 0 {
+			logrus.Infof("Resolved base name %s to %s", stagesResolvedArgs[index].BaseName, stagesResolvedArgs[index].Name)
 		}
-		stage.Name = resolvedBaseName
-		logrus.Infof("Resolved base name %s to %s", stage.BaseName, stage.Name)
+		baseImageIndex := baseImageIndex(index, stagesResolvedArgs)
 		kanikoStages = append(kanikoStages, config.KanikoStage{
 			Stage:                  stage,
-			BaseImageIndex:         baseImageIndex(index, stages),
-			BaseImageStoredLocally: (baseImageIndex(index, stages) != -1),
-			SaveStage:              saveStage(index, stages),
+			BaseImageIndex:         baseImageIndex,
+			BaseImageStoredLocally: (baseImageIndex != -1),
+			SaveStage:              saveStage(index, stagesResolvedArgs),
 			Final:                  index == targetStage,
 			MetaArgs:               metaArgs,
 			Index:                  index,
@@ -87,6 +94,25 @@ func Stages(opts *config.KanikoOptions) ([]config.KanikoStage, error) {
 	return kanikoStages, nil
 }
 
+// unifyArgs returns the unified args between metaArgs and --buildArgs
+// the chosen rule is that --buildArgs overrides metaArgs
+func unifyArgs(metaArgs []instructions.ArgCommand, buildArgs []string) []string {
+	argsMap := make(map[string]string)
+	for _, a := range metaArgs {
+		argsMap[a.Key] = *a.Value
+	}
+	splitter := "="
+	for _, a := range buildArgs {
+		s := strings.Split(a, splitter)
+		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 +254,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) ([]instructions.Stage, error) {
+	var stagesResolvedArgs []instructions.Stage
+	for _, s := range stages {
+		resolvedBaseName, err := util.ResolveEnvironmentReplacement(s.BaseName, args, false)
+		if err != nil {
+			return nil, errors.Wrap(err, fmt.Sprintf("resolving base name %s", s.BaseName))
+		}
+		s.BaseName = resolvedBaseName
+		stagesResolvedArgs = append(stagesResolvedArgs, s)
+	}
+	return stagesResolvedArgs, 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,72 @@ 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
+	`
+	stages, _, err := Parse([]byte(dockerfile))
+	if err != nil {
+		t.Fatal(err)
+	}
+	resolveStages(stages)
+
+	metaArgImage := "ubuntu:16.04"
+	buildArgImage := "\"alpine:3.11\""
+
+	lastVariants := []string{"dev", "prod"}
+
+	newArgCommand := func(key, val string) instructions.ArgCommand {
+		return instructions.ArgCommand{
+			KeyValuePairOptional: instructions.KeyValuePairOptional{Key: key, Value: &val},
+		}
+	}
+
+	for _, lastVariant := range lastVariants {
+		buildArgs := []string{fmt.Sprintf("IMAGE=%s", buildArgImage), fmt.Sprintf("LAST_STAGE_VARIANT=%s", lastVariant)}
+		metaArgs := []instructions.ArgCommand{newArgCommand("IMAGE", metaArgImage)}
+		args := unifyArgs(metaArgs, buildArgs)
+		resolvedArgs, err := resolveStagesArgs(stages, args)
+		if 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:   resolvedArgs[0].SourceCode,
+				actualBaseName:     resolvedArgs[0].BaseName,
+				expectedSourceCode: "FROM ${IMAGE} as base",
+				expectedBaseName:   buildArgImage,
+			},
+			{
+				name:               "Test_BuildArg_From_Last_Stage",
+				actualSourceCode:   resolvedArgs[len(resolvedArgs)-1].SourceCode,
+				actualBaseName:     resolvedArgs[len(resolvedArgs)-1].BaseName,
+				expectedSourceCode: "FROM base-${LAST_STAGE_VARIANT}",
+				expectedBaseName:   fmt.Sprintf("base-%s", lastVariant),
+			},
+		}
+		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)
+			})
+		}
+	}
+}