From 91009aea8c077369b11dda29c6d1d969567b5051 Mon Sep 17 00:00:00 2001
From: Yiran Wang <yiran@uber.com>
Date: Wed, 10 Jun 2020 22:33:19 -0700
Subject: [PATCH 1/2] Only cleanup fs when necessary

---
 bin/makisu/cmd/build.go    |  8 ++------
 lib/builder/build_plan.go  | 15 +++++++++++----
 lib/builder/build_stage.go |  3 ---
 3 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/bin/makisu/cmd/build.go b/bin/makisu/cmd/build.go
index 3817850b..bedc9ca8 100644
--- a/bin/makisu/cmd/build.go
+++ b/bin/makisu/cmd/build.go
@@ -236,19 +236,15 @@ func (cmd *buildCmd) Build(contextDir string) error {
 	}
 	defer buildContext.Cleanup()
 
-	// Make sure sandbox is cleaned after build.
-	// Optionally remove everything before and after build.
+	// Optionally stash everything before build and restore after build.
 	defer storage.CleanupSandbox(cmd.storageDir)
-	if cmd.allowModifyFS {
-		if cmd.preserveRoot {
+	if cmd.allowModifyFS && cmd.preserveRoot {
 			rootPreserver, err := storage.NewRootPreserver("/", cmd.storageDir, pathutils.DefaultBlacklist)
 			if err != nil {
 				return fmt.Errorf("failed to preserve root: %s", err)
 			}
 			defer rootPreserver.RestoreRoot()
 		}
-		buildContext.MemFS.Remove()
-		defer buildContext.MemFS.Remove()
 	}
 
 	// Create and execute build plan.
diff --git a/lib/builder/build_plan.go b/lib/builder/build_plan.go
index d022be3e..df88c2a5 100644
--- a/lib/builder/build_plan.go
+++ b/lib/builder/build_plan.go
@@ -234,6 +234,17 @@ func (plan *BuildPlan) Execute() (*image.DistributionManifest, error) {
 }
 
 func (plan *BuildPlan) executeStage(stage *buildStage, lastStage, copiedFrom bool) error {
+	if stage.opts.requireOnDisk {
+		if !stage.opts.allowModifyFS {
+			return fmt.Errorf("fs not allowed to be modified")
+		}
+		// Clean up local fs before build if current stage relies on local fs
+		// (i.e. contains RUN, or is source of COPY --from).
+		if err := stage.cleanup(); err != nil {
+			return fmt.Errorf("cleanup stage %s before build: %s", stage.alias, err)
+		}
+	}
+
 	if err := stage.build(plan.cacheMgr, lastStage, copiedFrom); err != nil {
 		return fmt.Errorf("build stage %s: %s", stage.alias, err)
 	}
@@ -245,10 +256,6 @@ func (plan *BuildPlan) executeStage(stage *buildStage, lastStage, copiedFrom boo
 		if err := stage.checkpoint(plan.copyFromDirs[stage.alias]); err != nil {
 			return fmt.Errorf("checkpoint stage %s: %s", stage.alias, err)
 		}
-
-		if err := stage.cleanup(); err != nil {
-			return fmt.Errorf("cleanup stage %s: %s", stage.alias, err)
-		}
 	}
 
 	return nil
diff --git a/lib/builder/build_stage.go b/lib/builder/build_stage.go
index 3c830acd..6f6bf8dc 100644
--- a/lib/builder/build_stage.go
+++ b/lib/builder/build_stage.go
@@ -175,9 +175,6 @@ func (stage *buildStage) build(cacheMgr cache.Manager, lastStage, copiedFrom boo
 	for i, node := range stage.nodes {
 		// Build current step from the previous image config (possibly cached).
 		modifyFS := stage.opts.requireOnDisk || copiedFrom
-		if modifyFS && !stage.opts.allowModifyFS {
-			return fmt.Errorf("fs not allowed to be modified")
-		}
 		skipBuild := i < stage.latestFetched() && i > 0
 		lastStep := i == len(stage.nodes)-1
 		forceCommit := i == 0 || (lastStage && lastStep) || stage.opts.forceCommit

From 6ee0c8ed6105165175d91da4cae076bc39ea5d73 Mon Sep 17 00:00:00 2001
From: Yiran Wang <yiran@uber.com>
Date: Wed, 10 Jun 2020 22:39:13 -0700
Subject: [PATCH 2/2] Fix build

---
 bin/makisu/cmd/build.go | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/bin/makisu/cmd/build.go b/bin/makisu/cmd/build.go
index bedc9ca8..436dee0c 100644
--- a/bin/makisu/cmd/build.go
+++ b/bin/makisu/cmd/build.go
@@ -239,12 +239,11 @@ func (cmd *buildCmd) Build(contextDir string) error {
 	// Optionally stash everything before build and restore after build.
 	defer storage.CleanupSandbox(cmd.storageDir)
 	if cmd.allowModifyFS && cmd.preserveRoot {
-			rootPreserver, err := storage.NewRootPreserver("/", cmd.storageDir, pathutils.DefaultBlacklist)
-			if err != nil {
-				return fmt.Errorf("failed to preserve root: %s", err)
-			}
-			defer rootPreserver.RestoreRoot()
+		rootPreserver, err := storage.NewRootPreserver("/", cmd.storageDir, pathutils.DefaultBlacklist)
+		if err != nil {
+			return fmt.Errorf("failed to preserve root: %s", err)
 		}
+		defer rootPreserver.RestoreRoot()
 	}
 
 	// Create and execute build plan.