mz440: cleanup kaniko workspace on failure too (#453)

* mz440: cleanup kaniko dir

* mz440: cleanup on failure too

* ai review

* address review items

Author: mzihlmann

README.md

@@ -967,7 +967,7 @@ ie. in order to support custom built kaniko images.
 
 #### Flag `--cleanup`
 
-Set this flag to clean the filesystem at the end of the build.
+Set this flag to clean the filesystem and kaniko's working directory at the end of the build.
 
 #### Flag `--compression`
 

pkg/config/init.go

@@ -17,9 +17,11 @@ limitations under the License.
 package config
 
 import (
+	"errors"
 	"fmt"
 	"os"
 	"path/filepath"
+	"strings"
 
 	"github.com/osscontainertools/kaniko/pkg/constants"
 	"github.com/sirupsen/logrus"
@@ -94,3 +96,68 @@ func init() {
 	RootDir = constants.RootDir
 	MountInfoPath = constants.MountInfoPath
 }
+
+// Same as os.RemoveAll, but asserts that we don't delete / or /kaniko.
+// This should be impossible at runtime and would indicate a programming mistake.
+func safeRemove(target string) error {
+	targetInfo, err := os.Stat(target)
+	if errors.Is(err, os.ErrNotExist) {
+		return nil
+	} else if err != nil {
+		return fmt.Errorf("failed to stat %q: %w", target, err)
+	}
+	rootInfo, err := os.Stat(RootDir)
+	if err != nil {
+		return fmt.Errorf("failed to stat %q: %w", RootDir, err)
+	}
+	kanikoInfo, err := os.Stat(KanikoDir)
+	if err != nil {
+		return fmt.Errorf("failed to stat %q: %w", KanikoDir, err)
+	}
+	if os.SameFile(targetInfo, rootInfo) {
+		logrus.Fatalf("refusing to remove /")
+	}
+	if os.SameFile(targetInfo, kanikoInfo) {
+		logrus.Fatalf("refusing to remove %q", KanikoDir)
+	}
+	if !strings.HasPrefix(target, KanikoDir+"/") {
+		logrus.Fatalf("refusing to remove %q outside %q", target, KanikoDir)
+	}
+	return os.RemoveAll(target)
+}
+
+func Cleanup() error {
+	err := safeRemove(DockerfilePath)
+	if err != nil {
+		return err
+	}
+	err = safeRemove(KanikoIntermediateStagesDir)
+	if err != nil {
+		return err
+	}
+	err = safeRemove(BuildContextDir)
+	if err != nil {
+		return err
+	}
+	if EnvBoolDefault("FF_KANIKO_NEW_CACHE_LAYOUT", true) {
+		err = safeRemove(KanikoInterStageDepsDir)
+		if err != nil {
+			return err
+		}
+		err = safeRemove(KanikoLayersDir)
+		if err != nil {
+			return err
+		}
+	}
+	err = safeRemove(KanikoSecretsDir)
+	if err != nil {
+		return err
+	}
+	_, err = os.Stat(KanikoSwapDir)
+	if err == nil {
+		return fmt.Errorf("expected directory %q to not exist, but it does", KanikoSwapDir)
+	} else if !errors.Is(err, os.ErrNotExist) {
+		return fmt.Errorf("failed to stat %q: %w", KanikoSwapDir, err)
+	}
+	return nil
+}

pkg/executor/build.go

@@ -724,7 +724,7 @@ func CalculateDependencies(stages []config.KanikoStage, opts *config.KanikoOptio
 }
 
 // DoBuild executes building the Dockerfile
-func DoBuild(opts *config.KanikoOptions) (v1.Image, error) {
+func DoBuild(opts *config.KanikoOptions) (image v1.Image, retErr error) {
 	t := timing.Start("Total Build Time")
 	digestToCacheKey := make(map[string]string)
 	stageIdxToDigest := make(map[int]string)
@@ -792,6 +792,29 @@ func DoBuild(opts *config.KanikoOptions) (v1.Image, error) {
 		}
 	}
 
+	if opts.Cleanup {
+		defer assignIfNil(&retErr, func() error {
+			if err = util.DeleteFilesystem(); err != nil {
+				return err
+			}
+			err = config.Cleanup()
+			if err != nil {
+				return err
+			}
+			if opts.PreserveContext {
+				if tarball == "" {
+					return fmt.Errorf("context snapshot is missing")
+				}
+				_, err := util.UnpackLocalTarArchive(tarball, config.RootDir)
+				if err != nil {
+					return fmt.Errorf("failed to unpack context snapshot: %w", err)
+				}
+				logrus.Info("Context restored")
+			}
+			return nil
+		})
+	}
+
 	for _, stage := range kanikoStages {
 		sb, err := newStageBuilder(
 			args, opts, stage,
@@ -859,21 +882,6 @@ func DoBuild(opts *config.KanikoOptions) (v1.Image, error) {
 			if len(opts.Annotations) > 0 {
 				sourceImage = mutate.Annotations(sourceImage, opts.Annotations).(v1.Image)
 			}
-			if opts.Cleanup {
-				if err = util.DeleteFilesystem(); err != nil {
-					return nil, err
-				}
-				if opts.PreserveContext {
-					if tarball == "" {
-						return nil, fmt.Errorf("context snapshot is missing")
-					}
-					_, err := util.UnpackLocalTarArchive(tarball, config.RootDir)
-					if err != nil {
-						return nil, fmt.Errorf("failed to unpack context snapshot: %w", err)
-					}
-					logrus.Info("Context restored")
-				}
-			}
 			timing.DefaultRun.Stop(t)
 			return sourceImage, nil
 		}
@@ -919,6 +927,12 @@ func DoBuild(opts *config.KanikoOptions) (v1.Image, error) {
 	return nil, err
 }
 
+func assignIfNil(dst *error, fn func() error) {
+	if err := fn(); err != nil && *dst == nil {
+		*dst = err
+	}
+}
+
 // filesToSave returns all the files matching the given pattern in deps.
 // If a file is a symlink, it also returns the target file.
 func filesToSave(deps []string) ([]string, error) {