Save each stage in multistage dockerfiles as a tarball (#244)

* resolve basenames in dockerfile to fix multistage bug

* WIP

* WIP

* Save dockerfile stages as tarballs

* added unit tests

* fix unit tests

cmd/executor/cmd/root.go

@@ -85,7 +85,7 @@ var RootCmd = &cobra.Command{
 			logrus.Error(err)
 			os.Exit(1)
 		}
-		ref, image, err := executor.DoBuild(executor.KanikoBuildArgs{
+		image, err := executor.DoBuild(executor.KanikoBuildArgs{
 			DockerfilePath: absouteDockerfilePath(),
 			SrcContext:     srcContext,
 			SnapshotMode:   snapshotMode,
@@ -98,7 +98,7 @@ var RootCmd = &cobra.Command{
 			os.Exit(1)
 		}
 
-		if err := executor.DoPush(ref, image, destinations, tarPath); err != nil {
+		if err := executor.DoPush(image, destinations, tarPath); err != nil {
 			logrus.Error(err)
 			os.Exit(1)
 		}

integration/dockerfiles/Dockerfile_test_multistage

@@ -1,10 +1,10 @@
-FROM gcr.io/distroless/base:latest
+FROM gcr.io/distroless/base:latest as base
 COPY . .
 
 FROM scratch as second
 ENV foopath context/foo
 COPY --from=0 $foopath context/b* /foo/
 
-FROM gcr.io/distroless/base:latest
+FROM base
 ARG file
 COPY --from=second /foo $file

pkg/constants/constants.go

@@ -40,6 +40,10 @@ const (
 	// for example, a tarball from a GCS bucket will be unpacked here
 	BuildContextDir = "/kaniko/buildcontext/"
 
+	// KanikoIntermediateStagesDir is where we will store intermediate stages
+	// as tarballs in case they are needed later on
+	KanikoIntermediateStagesDir = "/kaniko/stages"
+
 	// Various snapshot modes:
 	SnapshotModeTime = "time"
 	SnapshotModeFull = "full"

pkg/dockerfile/dockerfile.go

@@ -18,20 +18,13 @@ package dockerfile
 
 import (
 	"bytes"
-	"net/http"
-	"path/filepath"
-	"strconv"
-	"strings"
-
 	"github.com/GoogleContainerTools/kaniko/pkg/constants"
 	"github.com/GoogleContainerTools/kaniko/pkg/util"
 	"github.com/docker/docker/builder/dockerfile/instructions"
 	"github.com/docker/docker/builder/dockerfile/parser"
-	"github.com/google/go-containerregistry/pkg/authn"
+	"path/filepath"
-	"github.com/google/go-containerregistry/pkg/name"
+	"strconv"
-	"github.com/google/go-containerregistry/pkg/v1"
+	"strings"
-	"github.com/google/go-containerregistry/pkg/v1/empty"
-	"github.com/google/go-containerregistry/pkg/v1/remote"
 )
 
 // Parse parses the contents of a Dockerfile and returns a list of commands
@@ -89,29 +82,14 @@ func ParseCommands(cmdArray []string) ([]instructions.Command, error) {
 
 // Dependencies returns a list of files in this stage that will be needed in later stages
 func Dependencies(index int, stages []instructions.Stage, buildArgs *BuildArgs) ([]string, error) {
-	var dependencies []string
+	dependencies := []string{}
 	for stageIndex, stage := range stages {
 		if stageIndex <= index {
 			continue
 		}
-		var sourceImage v1.Image
+		sourceImage, err := util.RetrieveSourceImage(stageIndex, buildArgs.ReplacementEnvs(nil), stages)
-		if stage.BaseName == constants.NoBaseImage {
+		if err != nil {
-			sourceImage = empty.Image
+			return nil, err
-		} else {
-			// Initialize source image
-			ref, err := name.ParseReference(stage.BaseName, name.WeakValidation)
-			if err != nil {
-				return nil, err
-
-			}
-			auth, err := authn.DefaultKeychain.Resolve(ref.Context().Registry)
-			if err != nil {
-				return nil, err
-			}
-			sourceImage, err = remote.Image(ref, remote.WithAuth(auth), remote.WithTransport(http.DefaultTransport))
-			if err != nil {
-				return nil, err
-			}
 		}
 		imageConfig, err := sourceImage.ConfigFile()
 		if err != nil {

pkg/dockerfile/dockerfile_test.go

@@ -85,7 +85,7 @@ func Test_Dependencies(t *testing.T) {
 			helloPath,
 			testDir,
 		},
-		nil,
+		{},
 	}
 
 	for index := range stages {
@@ -125,7 +125,7 @@ func Test_DependenciesWithArg(t *testing.T) {
 			helloPath,
 			testDir,
 		},
-		nil,
+		{},
 	}
 	buildArgs := NewBuildArgs([]string{fmt.Sprintf("hienv=%s", helloPath)})
 

pkg/executor/executor.go

@@ -26,7 +26,6 @@ import (
 	"strconv"
 
 	"github.com/GoogleContainerTools/kaniko/pkg/snapshot"
-
 	"github.com/google/go-containerregistry/pkg/authn"
 	"github.com/google/go-containerregistry/pkg/name"
 	"github.com/google/go-containerregistry/pkg/v1"
@@ -55,94 +54,75 @@ type KanikoBuildArgs struct {
 	Reproducible   bool
 }
 
-func DoBuild(k KanikoBuildArgs) (name.Reference, v1.Image, error) {
+func DoBuild(k KanikoBuildArgs) (v1.Image, error) {
 	// Parse dockerfile and unpack base image to root
 	d, err := ioutil.ReadFile(k.DockerfilePath)
 	if err != nil {
-		return nil, nil, err
+		return nil, err
 	}
 
 	stages, err := dockerfile.Parse(d)
 	if err != nil {
-		return nil, nil, err
+		return nil, err
 	}
 	dockerfile.ResolveStages(stages)
 
 	hasher, err := getHasher(k.SnapshotMode)
 	if err != nil {
-		return nil, nil, err
+		return nil, err
 	}
 	for index, stage := range stages {
-		baseImage, err := util.ResolveEnvironmentReplacement(stage.BaseName, k.Args, false)
-		if err != nil {
-			return nil, nil, err
-		}
 		finalStage := index == len(stages)-1
 		// Unpack file system to root
-		logrus.Infof("Unpacking filesystem of %s...", baseImage)
+		sourceImage, err := util.RetrieveSourceImage(index, k.Args, stages)
-		var sourceImage v1.Image
+		if err != nil {
-		var ref name.Reference
+			return nil, err
-		if baseImage == constants.NoBaseImage {
-			logrus.Info("No base image, nothing to extract")
-			sourceImage = empty.Image
-		} else {
-			// Initialize source image
-			ref, err = name.ParseReference(baseImage, name.WeakValidation)
-			if err != nil {
-				return nil, nil, err
-			}
-			auth, err := authn.DefaultKeychain.Resolve(ref.Context().Registry)
-			if err != nil {
-				return nil, nil, err
-			}
-			sourceImage, err = remote.Image(ref, remote.WithAuth(auth), remote.WithTransport(http.DefaultTransport))
-			if err != nil {
-				return nil, nil, err
-			}
 		}
 		if err := util.GetFSFromImage(sourceImage); err != nil {
-			return nil, nil, err
+			return nil, err
 		}
 		l := snapshot.NewLayeredMap(hasher)
 		snapshotter := snapshot.NewSnapshotter(l, constants.RootDir)
 		// Take initial snapshot
 		if err := snapshotter.Init(); err != nil {
-			return nil, nil, err
+			return nil, err
 		}
 		imageConfig, err := sourceImage.ConfigFile()
-		if baseImage == constants.NoBaseImage {
+		if sourceImage == empty.Image {
 			imageConfig.Config.Env = constants.ScratchEnvVars
 		}
 		if err != nil {
-			return nil, nil, err
+			return nil, err
 		}
 		if err := resolveOnBuild(&stage, &imageConfig.Config); err != nil {
-			return nil, nil, err
+			return nil, err
 		}
 		buildArgs := dockerfile.NewBuildArgs(k.Args)
 		for index, cmd := range stage.Commands {
 			finalCmd := index == len(stage.Commands)-1
 			dockerCommand, err := commands.GetCommand(cmd, k.SrcContext)
 			if err != nil {
-				return nil, nil, err
+				return nil, err
 			}
 			if dockerCommand == nil {
 				continue
 			}
 			if err := dockerCommand.ExecuteCommand(&imageConfig.Config, buildArgs); err != nil {
-				return nil, nil, err
+				return nil, err
 			}
-			if !finalStage || (k.SingleSnapshot && !finalCmd) {
+			// 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 && k.SingleSnapshot) {
 				continue
 			}
 			// Now, we get the files to snapshot from this command and take the snapshot
 			snapshotFiles := dockerCommand.FilesToSnapshot()
-			if k.SingleSnapshot && finalCmd {
+			if finalCmd {
 				snapshotFiles = nil
 			}
 			contents, err := snapshotter.TakeSnapshot(snapshotFiles)
 			if err != nil {
-				return nil, nil, err
+				return nil, err
 			}
 			util.MoveVolumeWhitelistToWhitelist()
 			if contents == nil {
@@ -155,7 +135,7 @@ func DoBuild(k KanikoBuildArgs) (name.Reference, v1.Image, error) {
 			}
 			layer, err := tarball.LayerFromOpener(opener)
 			if err != nil {
-				return nil, nil, err
+				return nil, err
 			}
 			sourceImage, err = mutate.Append(sourceImage,
 				mutate.Addendum{
@@ -167,36 +147,37 @@ func DoBuild(k KanikoBuildArgs) (name.Reference, v1.Image, error) {
 				},
 			)
 			if err != nil {
-				return nil, nil, err
+				return nil, err
 			}
 		}
+		sourceImage, err = mutate.Config(sourceImage, imageConfig.Config)
+		if err != nil {
+			return nil, err
+		}
 		if finalStage {
-			sourceImage, err = mutate.Config(sourceImage, imageConfig.Config)
-			if err != nil {
-				return nil, nil, err
-			}
-
 			if k.Reproducible {
 				sourceImage, err = mutate.Canonical(sourceImage)
 				if err != nil {
-					return nil, nil, err
+					return nil, err
 				}
 			}
-
+			return sourceImage, nil
-			return ref, sourceImage, nil
+		}
+		if err := saveStageAsTarball(index, sourceImage); err != nil {
+			return nil, err
 		}
 		if err := saveStageDependencies(index, stages, buildArgs.Clone()); err != nil {
-			return nil, nil, err
+			return nil, err
 		}
 		// Delete the filesystem
 		if err := util.DeleteFilesystem(); err != nil {
-			return nil, nil, err
+			return nil, err
 		}
 	}
-	return nil, nil, err
+	return nil, err
 }
 
-func DoPush(ref name.Reference, image v1.Image, destinations []string, tarPath string) error {
+func DoPush(image v1.Image, destinations []string, tarPath string) error {
 	// continue pushing unless an error occurs
 	for _, destination := range destinations {
 		// Push the image
@@ -263,6 +244,19 @@ func saveStageDependencies(index int, stages []instructions.Stage, buildArgs *do
 	return nil
 }
 
+func saveStageAsTarball(stageIndex int, image v1.Image) error {
+	destRef, err := name.NewTag("temp/tag", name.WeakValidation)
+	if err != nil {
+		return err
+	}
+	if err := os.MkdirAll(constants.KanikoIntermediateStagesDir, 0750); err != nil {
+		return err
+	}
+	tarPath := filepath.Join(constants.KanikoIntermediateStagesDir, strconv.Itoa(stageIndex))
+	logrus.Infof("Storing source image from stage %d at path %s", stageIndex, tarPath)
+	return tarball.WriteToFile(tarPath, destRef, image, nil)
+}
+
 func getHasher(snapshotMode string) (func(string) (string, error), error) {
 	if snapshotMode == constants.SnapshotModeTime {
 		logrus.Info("Only file modification time will be considered when snapshotting")

pkg/util/image_util.go

@@ -0,0 +1,82 @@
+/*
+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 util
+
+import (
+	"github.com/GoogleContainerTools/kaniko/pkg/constants"
+	"github.com/docker/docker/builder/dockerfile/instructions"
+	"github.com/google/go-containerregistry/pkg/authn"
+	"github.com/google/go-containerregistry/pkg/name"
+	"github.com/google/go-containerregistry/pkg/v1"
+	"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/sirupsen/logrus"
+	"net/http"
+	"path/filepath"
+	"strconv"
+)
+
+var (
+	// For testing
+	retrieveRemoteImage = remoteImage
+	retrieveTarImage    = tarballImage
+)
+
+// RetrieveSourceImage returns the base image of the stage at index
+func RetrieveSourceImage(index int, buildArgs []string, stages []instructions.Stage) (v1.Image, error) {
+	currentStage := stages[index]
+	currentBaseName, err := ResolveEnvironmentReplacement(currentStage.BaseName, buildArgs, false)
+	if err != nil {
+		return nil, err
+	}
+	// First, check if the base image is a scratch image
+	if currentBaseName == constants.NoBaseImage {
+		logrus.Info("No base image, nothing to extract")
+		return empty.Image, nil
+	}
+	// 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 i > index {
+			continue
+		}
+		if stage.Name == currentBaseName {
+			return retrieveTarImage(i)
+		}
+	}
+	// Otherwise, initialize image as usual
+	return retrieveRemoteImage(currentBaseName)
+}
+
+func tarballImage(index int) (v1.Image, error) {
+	tarPath := filepath.Join(constants.KanikoIntermediateStagesDir, strconv.Itoa(index))
+	logrus.Infof("Base image from previous stage %d found, using saved tar at path %s", index, tarPath)
+	return tarball.ImageFromPath(tarPath, nil)
+}
+
+func remoteImage(image string) (v1.Image, error) {
+	ref, err := name.ParseReference(image, name.WeakValidation)
+	if err != nil {
+		return nil, err
+	}
+	auth, err := authn.DefaultKeychain.Resolve(ref.Context().Registry)
+	if err != nil {
+		return nil, err
+	}
+	return remote.Image(ref, remote.WithAuth(auth), remote.WithTransport(http.DefaultTransport))
+}

pkg/util/image_util_test.go

@@ -0,0 +1,97 @@
+/*
+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 util
+
+import (
+	"bytes"
+	"github.com/GoogleContainerTools/kaniko/testutil"
+	"github.com/docker/docker/builder/dockerfile/instructions"
+	"github.com/docker/docker/builder/dockerfile/parser"
+	"github.com/google/go-containerregistry/pkg/v1"
+	"github.com/google/go-containerregistry/pkg/v1/empty"
+	"testing"
+)
+
+var (
+	dockerfile = `
+	FROM gcr.io/distroless/base:latest as base
+	COPY . .
+	
+	FROM scratch as second
+	ENV foopath context/foo
+	COPY --from=0 $foopath context/b* /foo/
+	
+	FROM base
+	ARG file
+	COPY --from=second /foo $file`
+)
+
+func Test_StandardImage(t *testing.T) {
+	stages, err := parse(dockerfile)
+	if err != nil {
+		t.Error(err)
+	}
+	original := retrieveRemoteImage
+	defer func() {
+		retrieveRemoteImage = original
+	}()
+	mock := func(image string) (v1.Image, error) {
+		return nil, nil
+	}
+	retrieveRemoteImage = mock
+	actual, err := RetrieveSourceImage(0, nil, stages)
+	testutil.CheckErrorAndDeepEqual(t, false, err, nil, actual)
+}
+func Test_ScratchImage(t *testing.T) {
+	stages, err := parse(dockerfile)
+	if err != nil {
+		t.Error(err)
+	}
+	actual, err := RetrieveSourceImage(1, nil, stages)
+	expected := empty.Image
+	testutil.CheckErrorAndDeepEqual(t, false, err, expected, actual)
+}
+
+func Test_TarImage(t *testing.T) {
+	stages, err := parse(dockerfile)
+	if err != nil {
+		t.Error(err)
+	}
+	original := retrieveTarImage
+	defer func() {
+		retrieveTarImage = original
+	}()
+	mock := func(index int) (v1.Image, error) {
+		return nil, nil
+	}
+	retrieveTarImage = mock
+	actual, err := RetrieveSourceImage(2, nil, stages)
+	testutil.CheckErrorAndDeepEqual(t, false, err, nil, actual)
+}
+
+// parse parses the contents of a Dockerfile and returns a list of commands
+func parse(s string) ([]instructions.Stage, error) {
+	p, err := parser.Parse(bytes.NewReader([]byte(s)))
+	if err != nil {
+		return nil, err
+	}
+	stages, _, err := instructions.Parse(p.AST)
+	if err != nil {
+		return nil, err
+	}
+	return stages, err
+}