Added a KanikoStage type for each stage of a Dockerfile

I added a KanikoStage to hold each stage of the Dockerfile along with
information about each stage that would be useful later on.

The new KanikoStage type holds the stage itself, along with some
additional information:

1. FinalStage -- whether the current stage is the final stage
2. BaseImageStoredLocally/BaseImageIndex -- whether the base image for
this stage is stored locally, and if so what the index of the base image
is
3. SaveStage -- whether this stage needs to be saved for use in a future
stage

This is the first part of a larger refactor for building stages, which
will later make it easier to add layer caching.

cmd/executor/cmd/root.go

@@ -22,9 +22,9 @@ import (
 	"strings"
 
 	"github.com/GoogleContainerTools/kaniko/pkg/buildcontext"
+	"github.com/GoogleContainerTools/kaniko/pkg/config"
 	"github.com/GoogleContainerTools/kaniko/pkg/constants"
 	"github.com/GoogleContainerTools/kaniko/pkg/executor"
-	"github.com/GoogleContainerTools/kaniko/pkg/options"
 	"github.com/GoogleContainerTools/kaniko/pkg/util"
 	"github.com/genuinetools/amicontained/container"
 	"github.com/pkg/errors"
@@ -33,7 +33,7 @@ import (
 )
 
 var (
-	opts     = &options.KanikoOptions{}
+	opts     = &config.KanikoOptions{}
 	logLevel string
 	force    bool
 )

pkg/config/args.go

@@ -14,7 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-package options
+package config
 
 import (
 	"strings"

pkg/config/options.go

@@ -14,7 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-package options
+package config
 
 // KanikoOptions are options that are set by command line arguments
 type KanikoOptions struct {

pkg/config/stage.go

@@ -0,0 +1,28 @@
+/*
+Copyright 2018 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package config
+
+import "github.com/moby/buildkit/frontend/dockerfile/instructions"
+
+// KanikoStage wraps a stage of the Dockerfile and provides extra information
+type KanikoStage struct {
+	instructions.Stage
+	FinalStage             bool
+	BaseImageStoredLocally bool
+	BaseImageIndex         int
+	SaveStage              bool
+}

pkg/dockerfile/dockerfile.go

@@ -23,26 +23,61 @@ import (
 	"strconv"
 	"strings"
 
+	"github.com/GoogleContainerTools/kaniko/pkg/config"
+	"github.com/GoogleContainerTools/kaniko/pkg/util"
 	"github.com/moby/buildkit/frontend/dockerfile/instructions"
 	"github.com/moby/buildkit/frontend/dockerfile/parser"
+	"github.com/pkg/errors"
 )
 
-// Stages reads the Dockerfile, validates it's contents, and returns stages
+// Stages parses a Dockerfile and returns an array of KanikoStage
-func Stages(dockerfilePath, target string) ([]instructions.Stage, error) {
+func Stages(opts *config.KanikoOptions) ([]config.KanikoStage, error) {
-	d, err := ioutil.ReadFile(dockerfilePath)
+	d, err := ioutil.ReadFile(opts.DockerfilePath)
 	if err != nil {
-		return nil, err
+		return nil, errors.Wrap(err, fmt.Sprintf("reading dockerfile at path %s", opts.DockerfilePath))
 	}
-
 	stages, err := Parse(d)
+	if err != nil {
+		return nil, errors.Wrap(err, "parsing dockerfile")
+	}
+	targetStage, err := targetStage(stages, opts.Target)
 	if err != nil {
 		return nil, err
 	}
-	if err := ValidateTarget(stages, target); err != nil {
+	resolveStages(stages)
-		return nil, err
+	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")
+		}
+		stage.Name = resolvedBaseName
+		kanikoStages = append(kanikoStages, config.KanikoStage{
+			Stage:                  stage,
+			BaseImageIndex:         baseImageIndex(opts, index, stages),
+			BaseImageStoredLocally: (baseImageIndex(opts, index, stages) != -1),
+			SaveStage:              saveStage(index, stages),
+			FinalStage:             index == targetStage,
+		})
+		if index == targetStage {
+			break
+		}
 	}
-	ResolveStages(stages)
+	return kanikoStages, nil
-	return stages, nil
+}
+
+// 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(opts *config.KanikoOptions, currentStage int, stages []instructions.Stage) int {
+	for i, stage := range stages {
+		if i > currentStage {
+			break
+		}
+		if stage.Name == stages[currentStage].BaseName {
+			return i
+		}
+	}
+	return -1
 }
 
 // Parse parses the contents of a Dockerfile and returns a list of commands
@@ -58,21 +93,22 @@ func Parse(b []byte) ([]instructions.Stage, error) {
 	return stages, err
 }
 
-func ValidateTarget(stages []instructions.Stage, target string) error {
+// targetStage returns the index of the target stage kaniko is trying to build
+func targetStage(stages []instructions.Stage, target string) (int, error) {
 	if target == "" {
-		return nil
+		return len(stages) - 1, nil
 	}
-	for _, stage := range stages {
+	for i, stage := range stages {
 		if stage.Name == target {
-			return nil
+			return i, nil
 		}
 	}
-	return fmt.Errorf("%s is not a valid target build stage", target)
+	return -1, fmt.Errorf("%s is not a valid target build stage", target)
 }
 
-// ResolveStages resolves any calls to previous stages with names to indices
+// resolveStages resolves any calls to previous stages with names to indices
 // Ex. --from=second_stage should be --from=1 for easier processing later on
-func ResolveStages(stages []instructions.Stage) {
+func resolveStages(stages []instructions.Stage) {
 	nameToIndex := make(map[string]string)
 	for i, stage := range stages {
 		index := strconv.Itoa(i)
@@ -111,7 +147,7 @@ func ParseCommands(cmdArray []string) ([]instructions.Command, error) {
 }
 
 // SaveStage returns true if the current stage will be needed later in the Dockerfile
-func SaveStage(index int, stages []instructions.Stage) bool {
+func saveStage(index int, stages []instructions.Stage) bool {
 	for stageIndex, stage := range stages {
 		if stageIndex <= index {
 			continue

pkg/dockerfile/dockerfile_test.go

@@ -17,17 +17,15 @@ limitations under the License.
 package dockerfile
 
 import (
-	"io/ioutil"
-	"os"
-	"path/filepath"
 	"strconv"
 	"testing"
 
+	"github.com/GoogleContainerTools/kaniko/pkg/config"
 	"github.com/GoogleContainerTools/kaniko/testutil"
 	"github.com/moby/buildkit/frontend/dockerfile/instructions"
 )
 
-func Test_ResolveStages(t *testing.T) {
+func Test_resolveStages(t *testing.T) {
 	dockerfile := `
 	FROM scratch
 	RUN echo hi > /hi
@@ -42,7 +40,7 @@ func Test_ResolveStages(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	ResolveStages(stages)
+	resolveStages(stages)
 	for index, stage := range stages {
 		if index == 0 {
 			continue
@@ -55,7 +53,7 @@ func Test_ResolveStages(t *testing.T) {
 	}
 }
 
-func Test_ValidateTarget(t *testing.T) {
+func Test_targetStage(t *testing.T) {
 	dockerfile := `
 	FROM scratch
 	RUN echo hi > /hi
@@ -71,70 +69,44 @@ func Test_ValidateTarget(t *testing.T) {
 		t.Fatal(err)
 	}
 	tests := []struct {
-		name      string
+		name        string
-		target    string
+		target      string
-		shouldErr bool
+		targetIndex int
+		shouldErr   bool
 	}{
 		{
-			name:      "test valid target",
+			name:        "test valid target",
-			target:    "second",
+			target:      "second",
-			shouldErr: false,
+			targetIndex: 1,
+			shouldErr:   false,
 		},
 		{
-			name:      "test invalid target",
+			name:        "test no target",
-			target:    "invalid",
+			target:      "",
-			shouldErr: true,
+			targetIndex: 2,
+			shouldErr:   false,
+		},
+		{
+			name:        "test invalid target",
+			target:      "invalid",
+			targetIndex: -1,
+			shouldErr:   true,
 		},
 	}
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
-			actualErr := ValidateTarget(stages, test.target)
+			target, err := targetStage(stages, test.target)
-			testutil.CheckError(t, test.shouldErr, actualErr)
+			testutil.CheckError(t, test.shouldErr, err)
+			if !test.shouldErr {
+				if target != test.targetIndex {
+					t.Errorf("got incorrect target, expected %d got %d", test.targetIndex, target)
+				}
+			}
 		})
 	}
 }
 
 func Test_SaveStage(t *testing.T) {
-	tempDir, err := ioutil.TempDir("", "")
-	if err != nil {
-		t.Fatalf("couldn't create temp dir: %v", err)
-	}
-	defer os.RemoveAll(tempDir)
-	files := map[string]string{
-		"Dockerfile": `
-		FROM scratch
-		RUN echo hi > /hi
-		
-		FROM scratch AS second
-		COPY --from=0 /hi /hi2
-	
-		FROM second
-		RUN xxx
-		
-		FROM scratch
-		COPY --from=second /hi2 /hi3
-
-		FROM ubuntu:16.04 AS base
-		ENV DEBIAN_FRONTEND noninteractive
-		ENV LC_ALL C.UTF-8
-
-		FROM base AS development
-		ENV PS1 " 🐳 \[\033[1;36m\]\W\[\033[0;35m\] # \[\033[0m\]"
-
-		FROM development AS test
-		ENV ORG_ENV UnitTest
-
-		FROM base AS production
-		COPY . /code
-		`,
-	}
-	if err := testutil.SetupFiles(tempDir, files); err != nil {
-		t.Fatalf("couldn't create dockerfile: %v", err)
-	}
-	stages, err := Stages(filepath.Join(tempDir, "Dockerfile"), "")
-	if err != nil {
-		t.Fatalf("couldn't retrieve stages from Dockerfile: %v", err)
-	}
 	tests := []struct {
 		name     string
 		index    int
@@ -171,10 +143,51 @@ func Test_SaveStage(t *testing.T) {
 			expected: false,
 		},
 	}
+	stages, err := Parse([]byte(testutil.Dockerfile))
+	if err != nil {
+		t.Fatalf("couldn't retrieve stages from Dockerfile: %v", err)
+	}
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
-			actual := SaveStage(test.index, stages)
+			actual := saveStage(test.index, stages)
 			testutil.CheckErrorAndDeepEqual(t, false, nil, test.expected, actual)
 		})
 	}
 }
+
+func Test_baseImageIndex(t *testing.T) {
+	tests := []struct {
+		name         string
+		currentStage int
+		expected     int
+	}{
+		{
+			name:         "stage that is built off of a previous stage",
+			currentStage: 2,
+			expected:     1,
+		},
+		{
+			name:         "another stage that is built off of a previous stage",
+			currentStage: 5,
+			expected:     4,
+		},
+		{
+			name:         "stage that isn't built off of a previous stage",
+			currentStage: 4,
+			expected:     -1,
+		},
+	}
+
+	stages, err := Parse([]byte(testutil.Dockerfile))
+	if err != nil {
+		t.Fatalf("couldn't retrieve stages from Dockerfile: %v", err)
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			actual := baseImageIndex(&config.KanikoOptions{}, test.currentStage, stages)
+			if actual != test.expected {
+				t.Fatalf("unexpected result, expected %d got %d", test.expected, actual)
+			}
+		})
+	}
+}

pkg/executor/build.go

@@ -29,20 +29,19 @@ import (
 	"github.com/google/go-containerregistry/pkg/v1"
 	"github.com/google/go-containerregistry/pkg/v1/mutate"
 	"github.com/google/go-containerregistry/pkg/v1/tarball"
-	"github.com/moby/buildkit/frontend/dockerfile/instructions"
 	"github.com/sirupsen/logrus"
 
 	"github.com/GoogleContainerTools/kaniko/pkg/commands"
+	"github.com/GoogleContainerTools/kaniko/pkg/config"
 	"github.com/GoogleContainerTools/kaniko/pkg/constants"
 	"github.com/GoogleContainerTools/kaniko/pkg/dockerfile"
-	"github.com/GoogleContainerTools/kaniko/pkg/options"
 	"github.com/GoogleContainerTools/kaniko/pkg/snapshot"
 	"github.com/GoogleContainerTools/kaniko/pkg/util"
 )
 
-func DoBuild(opts *options.KanikoOptions) (v1.Image, error) {
+func DoBuild(opts *config.KanikoOptions) (v1.Image, error) {
 	// Parse dockerfile and unpack base image to root
-	stages, err := dockerfile.Stages(opts.DockerfilePath, opts.Target)
+	stages, err := dockerfile.Stages(opts)
 	if err != nil {
 		return nil, err
 	}
@@ -52,9 +51,8 @@ func DoBuild(opts *options.KanikoOptions) (v1.Image, error) {
 		return nil, err
 	}
 	for index, stage := range stages {
-		finalStage := finalStage(index, opts.Target, stages)
 		// Unpack file system to root
-		sourceImage, err := util.RetrieveSourceImage(index, opts.BuildArgs, stages)
+		sourceImage, err := util.RetrieveSourceImage(stage, opts.BuildArgs)
 		if err != nil {
 			return nil, err
 		}
@@ -89,7 +87,7 @@ func DoBuild(opts *options.KanikoOptions) (v1.Image, error) {
 			}
 			// Don't snapshot if it's not the final stage and not the final command
 			// Also don't snapshot if it's the final stage, not the final command, and single snapshot is set
-			if (!finalStage && !finalCmd) || (finalStage && !finalCmd && opts.SingleSnapshot) {
+			if (!stage.FinalStage && !finalCmd) || (stage.FinalStage && !finalCmd && opts.SingleSnapshot) {
 				continue
 			}
 			// Now, we get the files to snapshot from this command and take the snapshot
@@ -131,7 +129,7 @@ func DoBuild(opts *options.KanikoOptions) (v1.Image, error) {
 		if err != nil {
 			return nil, err
 		}
-		if finalStage {
+		if stage.FinalStage {
 			if opts.Reproducible {
 				sourceImage, err = mutate.Canonical(sourceImage)
 				if err != nil {
@@ -140,7 +138,7 @@ func DoBuild(opts *options.KanikoOptions) (v1.Image, error) {
 			}
 			return sourceImage, nil
 		}
-		if dockerfile.SaveStage(index, stages) {
+		if stage.SaveStage {
 			if err := saveStageAsTarball(index, sourceImage); err != nil {
 				return nil, err
 			}
@@ -156,16 +154,6 @@ func DoBuild(opts *options.KanikoOptions) (v1.Image, error) {
 	return nil, err
 }
 
-func finalStage(index int, target string, stages []instructions.Stage) bool {
-	if index == len(stages)-1 {
-		return true
-	}
-	if target == "" {
-		return false
-	}
-	return target == stages[index].Name
-}
-
 func extractImageToDependecyDir(index int, image v1.Image) error {
 	dependencyDir := filepath.Join(constants.KanikoDir, strconv.Itoa(index))
 	if err := os.MkdirAll(dependencyDir, 0755); err != nil {
@@ -199,7 +187,7 @@ func getHasher(snapshotMode string) (func(string) (string, error), error) {
 	return nil, fmt.Errorf("%s is not a valid snapshot mode", snapshotMode)
 }
 
-func resolveOnBuild(stage *instructions.Stage, config *v1.Config) error {
+func resolveOnBuild(stage *config.KanikoStage, config *v1.Config) error {
 	if config.OnBuild == nil {
 		return nil
 	}

pkg/executor/push.go

@@ -21,7 +21,7 @@ import (
 	"fmt"
 	"net/http"
 
-	"github.com/GoogleContainerTools/kaniko/pkg/options"
+	"github.com/GoogleContainerTools/kaniko/pkg/config"
 	"github.com/GoogleContainerTools/kaniko/pkg/version"
 	"github.com/google/go-containerregistry/pkg/authn"
 	"github.com/google/go-containerregistry/pkg/authn/k8schain"
@@ -43,7 +43,7 @@ func (w *withUserAgent) RoundTrip(r *http.Request) (*http.Response, error) {
 }
 
 // DoPush is responsible for pushing image to the destinations specified in opts
-func DoPush(image v1.Image, opts *options.KanikoOptions) error {
+func DoPush(image v1.Image, opts *config.KanikoOptions) error {
 	if opts.NoPush {
 		logrus.Info("Skipping push to container registry due to --no-push flag")
 		return nil

pkg/util/image_util.go

@@ -27,9 +27,9 @@ import (
 	"github.com/google/go-containerregistry/pkg/v1/empty"
 	"github.com/google/go-containerregistry/pkg/v1/remote"
 	"github.com/google/go-containerregistry/pkg/v1/tarball"
-	"github.com/moby/buildkit/frontend/dockerfile/instructions"
 	"github.com/sirupsen/logrus"
 
+	"github.com/GoogleContainerTools/kaniko/pkg/config"
 	"github.com/GoogleContainerTools/kaniko/pkg/constants"
 )
 
@@ -40,9 +40,8 @@ var (
 )
 
 // RetrieveSourceImage returns the base image of the stage at index
-func RetrieveSourceImage(index int, buildArgs []string, stages []instructions.Stage) (v1.Image, error) {
+func RetrieveSourceImage(stage config.KanikoStage, buildArgs []string) (v1.Image, error) {
-	currentStage := stages[index]
+	currentBaseName, err := ResolveEnvironmentReplacement(stage.BaseName, buildArgs, false)
-	currentBaseName, err := ResolveEnvironmentReplacement(currentStage.BaseName, buildArgs, false)
 	if err != nil {
 		return nil, err
 	}
@@ -53,14 +52,10 @@ func RetrieveSourceImage(index int, buildArgs []string, stages []instructions.St
 	}
 	// Next, check if the base image of the current stage is built from a previous stage
 	// If so, retrieve the image from the stored tarball
-	for i, stage := range stages {
+	if stage.BaseImageStoredLocally {
-		if i > index {
+		return retrieveTarImage(stage.BaseImageIndex)
-			continue
-		}
-		if stage.Name == currentBaseName {
-			return retrieveTarImage(i)
-		}
 	}
+
 	// Otherwise, initialize image as usual
 	return retrieveRemoteImage(currentBaseName)
 }

pkg/util/image_util_test.go

@@ -20,6 +20,7 @@ import (
 	"bytes"
 	"testing"
 
+	"github.com/GoogleContainerTools/kaniko/pkg/config"
 	"github.com/GoogleContainerTools/kaniko/testutil"
 	"github.com/google/go-containerregistry/pkg/v1"
 	"github.com/google/go-containerregistry/pkg/v1/empty"
@@ -54,7 +55,9 @@ func Test_StandardImage(t *testing.T) {
 		return nil, nil
 	}
 	retrieveRemoteImage = mock
-	actual, err := RetrieveSourceImage(0, nil, stages)
+	actual, err := RetrieveSourceImage(config.KanikoStage{
+		Stage: stages[0],
+	}, nil)
 	testutil.CheckErrorAndDeepEqual(t, false, err, nil, actual)
 }
 func Test_ScratchImage(t *testing.T) {
@@ -62,7 +65,9 @@ func Test_ScratchImage(t *testing.T) {
 	if err != nil {
 		t.Error(err)
 	}
-	actual, err := RetrieveSourceImage(1, nil, stages)
+	actual, err := RetrieveSourceImage(config.KanikoStage{
+		Stage: stages[1],
+	}, nil)
 	expected := empty.Image
 	testutil.CheckErrorAndDeepEqual(t, false, err, expected, actual)
 }
@@ -80,7 +85,11 @@ func Test_TarImage(t *testing.T) {
 		return nil, nil
 	}
 	retrieveTarImage = mock
-	actual, err := RetrieveSourceImage(2, nil, stages)
+	actual, err := RetrieveSourceImage(config.KanikoStage{
+		BaseImageStoredLocally: true,
+		BaseImageIndex:         0,
+		Stage:                  stages[2],
+	}, nil)
 	testutil.CheckErrorAndDeepEqual(t, false, err, nil, actual)
 }
 

testutil/constants.go

@@ -0,0 +1,47 @@
+/*
+Copyright 2018 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package testutil
+
+const (
+	// Dockerfile is used for unit testing
+	Dockerfile = `
+	FROM scratch
+	RUN echo hi > /hi
+	
+	FROM scratch AS second
+	COPY --from=0 /hi /hi2
+
+	FROM second
+	RUN xxx
+	
+	FROM scratch
+	COPY --from=second /hi2 /hi3
+
+	FROM ubuntu:16.04 AS base
+	ENV DEBIAN_FRONTEND noninteractive
+	ENV LC_ALL C.UTF-8
+
+	FROM base AS development
+	ENV PS1 " 🐳 \[\033[1;36m\]\W\[\033[0;35m\] # \[\033[0m\]"
+
+	FROM development AS test
+	ENV ORG_ENV UnitTest
+
+	FROM base AS production
+	COPY . /code
+	`
+)