Extract only output paths from the action workspace

enterprise/server/remote_execution/containers/firecracker/firecracker.go

@@ -1690,7 +1690,8 @@ func (c *FirecrackerContainer) copyOutputsToWorkspace(ctx context.Context) error
 		}
 		defer m.Unmount()
 	} else {
-		if err := ext4.ImageToDirectory(ctx, workspaceExt4Path, wsDir); err != nil {
+		outputPaths := workspacePathsToExtract(c.task)
+		if err := ext4.ImageToDirectory(ctx, workspaceExt4Path, wsDir, outputPaths); err != nil {
 			return err
 		}
 	}
@@ -1699,10 +1700,6 @@ func (c *FirecrackerContainer) copyOutputsToWorkspace(ctx context.Context) error
 		if err != nil {
 			return err
 		}
-		// Skip filesystem layerfs write-layer files.
-		if strings.HasPrefix(path, "bbvmroot/") || strings.HasPrefix(path, "bbvmwork/") {
-			return nil
-		}
 		targetLocation := filepath.Join(c.actionWorkingDir, path)
 		alreadyExists, err := disk.FileExists(ctx, targetLocation)
 		if err != nil {
@@ -2870,3 +2867,25 @@ func getRandomNUMANode() (int, error) {
 
 	return nodes[rand.IntN(len(nodes))], nil
 }
+
+// Returns the paths relative to the workspace root that should be copied back
+// to the action workspace directory after execution has completed.
+//
+// For performance reasons, we only extract the action's declared outputs,
+// unless the action is running with preserve-workspace=true.
+func workspacePathsToExtract(task *repb.ExecutionTask) []string {
+	if platform.IsTrue(platform.FindEffectiveValue(task, platform.PreserveWorkspacePropertyName)) {
+		return []string{"/"}
+	}
+
+	// Special files
+	// TODO: declare this list as a constant somewhere?
+	paths := []string{".BUILDBUDDY_DO_NOT_RECYCLE"}
+
+	// Declared paths
+	paths = append(paths, task.GetCommand().GetOutputDirectories()...)
+	paths = append(paths, task.GetCommand().GetOutputFiles()...)
+	paths = append(paths, task.GetCommand().GetOutputPaths()...)
+
+	return paths
+}

enterprise/server/remote_execution/containers/firecracker/firecracker_test.go

@@ -1404,7 +1404,11 @@ func TestFirecrackerRun_ReapOrphanedZombieProcess(t *testing.T) {
 		},
 		ExecutorConfig: getExecutorConfig(t),
 	}
-	c, err := firecracker.NewContainer(ctx, env, &repb.ExecutionTask{}, opts)
+	c, err := firecracker.NewContainer(ctx, env, &repb.ExecutionTask{
+		Command: &repb.Command{
+			OutputPaths: []string{"sh.pid", "sleep.pid"},
+		},
+	}, opts)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -1815,7 +1819,11 @@ func TestFirecrackerExecWithRecycledWorkspaceWithDocker(t *testing.T) {
 		},
 		ExecutorConfig: getExecutorConfig(t),
 	}
-	c, err := firecracker.NewContainer(ctx, env, &repb.ExecutionTask{}, opts)
+	c, err := firecracker.NewContainer(ctx, env, &repb.ExecutionTask{
+		Command: &repb.Command{
+			OutputPaths: []string{"preserves.txt"},
+		},
+	}, opts)
 	require.NoError(t, err)
 	err = container.PullImageIfNecessary(ctx, env, c, oci.Credentials{}, opts.ContainerImage)
 	require.NoError(t, err)
@@ -1988,7 +1996,11 @@ func TestFirecrackerRun_Timeout_DebugOutputIsAvailable(t *testing.T) {
 		},
 		ExecutorConfig: getExecutorConfig(t),
 	}
-	c, err := firecracker.NewContainer(ctx, env, &repb.ExecutionTask{}, opts)
+	c, err := firecracker.NewContainer(ctx, env, &repb.ExecutionTask{
+		Command: &repb.Command{
+			OutputPaths: []string{"output.txt"},
+		},
+	}, opts)
 	require.NoError(t, err)
 
 	cmd := &repb.Command{Arguments: []string{"sh", "-c", `
@@ -2032,7 +2044,11 @@ func TestFirecrackerExec_Timeout_DebugOutputIsAvailable(t *testing.T) {
 		},
 		ExecutorConfig: getExecutorConfig(t),
 	}
-	c, err := firecracker.NewContainer(ctx, env, &repb.ExecutionTask{}, opts)
+	c, err := firecracker.NewContainer(ctx, env, &repb.ExecutionTask{
+		Command: &repb.Command{
+			OutputPaths: []string{"output.txt"},
+		},
+	}, opts)
 	require.NoError(t, err)
 	err = container.PullImageIfNecessary(ctx, env, c, oci.Credentials{}, opts.ContainerImage)
 	require.NoError(t, err)

enterprise/server/remote_execution/overlayfs/overlayfs_test.go

@@ -62,7 +62,7 @@ func TestOverlayWorkspace(t *testing.T) {
 	// removed in t.Cleanup() above, otherwise we get 'device or resource busy'
 	// when attempting to clean up the workspace.
 	t.Cleanup(func() { fileAOverlay.Close() })
-	// Remove dir1/b.txt and dir1/child1/.
+	// Remove dir1's children (b.txt and child1/), but keep dir1 itself.
 	err = os.RemoveAll(filepath.Join(ws, "dir1/b.txt"))
 	require.NoError(t, err)
 	err = os.RemoveAll(filepath.Join(ws, "dir1/child1/"))
@@ -101,6 +101,7 @@ func TestOverlayWorkspace(t *testing.T) {
 	lowerdir := ws + ".lower"
 	testfs.AssertExactFileContents(t, lowerdir, map[string]string{
 		"a.txt":                "A-MODIFIED",
+		"dir1":                 testfs.EmptyDir,
 		"dir2/d.txt":           "D",
 		"dir3/e.txt/child.txt": "child-contents",
 	})

enterprise/server/remote_execution/platform/platform.go

@@ -74,7 +74,7 @@ const (
 	RecycleRunnerPropertyName            = "recycle-runner"
 	AffinityRoutingPropertyName          = "affinity-routing"
 	RunnerRecyclingMaxWaitPropertyName   = "runner-recycling-max-wait"
-	preserveWorkspacePropertyName        = "preserve-workspace"
+	PreserveWorkspacePropertyName        = "preserve-workspace"
 	nonrootWorkspacePropertyName         = "nonroot-workspace"
 	overlayfsWorkspacePropertyName       = "overlayfs-workspace"
 	cleanWorkspaceInputsPropertyName     = "clean-workspace-inputs"
@@ -351,7 +351,7 @@ func ParseProperties(task *repb.ExecutionTask) (*Properties, error) {
 		RunnerRecyclingMaxWait:    runnerRecyclingMaxWait,
 		EnableVFS:                 vfsEnabled,
 		IncludeSecrets:            boolProp(m, IncludeSecretsPropertyName, false),
-		PreserveWorkspace:         boolProp(m, preserveWorkspacePropertyName, false),
+		PreserveWorkspace:         boolProp(m, PreserveWorkspacePropertyName, false),
 		OverlayfsWorkspace:        boolProp(m, overlayfsWorkspacePropertyName, false),
 		NonrootWorkspace:          boolProp(m, nonrootWorkspacePropertyName, false),
 		CleanWorkspaceInputs:      stringProp(m, cleanWorkspaceInputsPropertyName, ""),

enterprise/server/util/ext4/ext4.go

@@ -10,6 +10,7 @@ import (
 	"os"
 	"os/exec"
 	"path/filepath"
+	"strings"
 	"syscall"
 
 	"github.com/buildbuddy-io/buildbuddy/server/util/log"
@@ -166,21 +167,30 @@ func isDirEmpty(dir string) (bool, error) {
 }
 
 // ImageToDirectory unpacks an ext4 image into outputDir, which must be empty.
-func ImageToDirectory(ctx context.Context, inputFile, outputDir string) error {
+// Only the given paths are unpacked. Paths are unpacked recursively, which
+// means that they can reference either directories or files. Non-existent
+// paths are silently ignored.
+func ImageToDirectory(ctx context.Context, inputFile, outputDir string, paths []string) error {
 	empty, err := isDirEmpty(outputDir)
 	if err != nil {
 		return err
 	}
 	if !empty {
 		return status.FailedPreconditionError("Unpacking image in non-empty directory is unsupported.")
 	}
-	args := []string{
-		"/sbin/debugfs",
-		inputFile,
-		"-R",
-		fmt.Sprintf("rdump \"/\" \"%s\"", outputDir),
+	requests := make([]string, 0, len(paths))
+	for _, p := range paths {
+		p = filepath.Clean(filepath.Join(outputDir, p))
+		p = strings.TrimPrefix(p, filepath.Clean(outputDir))
+		parent := filepath.Dir(p)
+		if err := os.MkdirAll(filepath.Join(outputDir, parent), 0755); err != nil {
+			return status.InternalErrorf("make parent dir %q: %s", parent, err)
+		}
+		requests = append(requests, fmt.Sprintf("rdump %q %q", p, filepath.Join(outputDir, parent)))
 	}
-	if out, err := exec.CommandContext(ctx, args[0], args[1:]...).CombinedOutput(); err != nil {
+	cmd := exec.CommandContext(ctx, "/sbin/debugfs", inputFile)
+	cmd.Stdin = strings.NewReader(strings.Join(requests, "\n"))
+	if out, err := cmd.CombinedOutput(); err != nil {
 		return status.InternalErrorf("%s: %s", err, out)
 	}
 	return nil

enterprise/server/util/ext4/ext4_test.go

@@ -61,7 +61,7 @@ func TestE2E(t *testing.T) {
 
 	// Create a new (empty) directory and unpack the image file to the new directory.
 	newRoot := testfs.MakeTempDir(t)
-	if err := ext4.ImageToDirectory(ctx, imageFile, newRoot); err != nil {
+	if err := ext4.ImageToDirectory(ctx, imageFile, newRoot, []string{"/"}); err != nil {
 		t.Fatal(err)
 	}
 
@@ -103,3 +103,160 @@ func TestDirectoryToImageAutoSize_DanglingSymlink(t *testing.T) {
 
 	require.NoError(t, err)
 }
+
+func TestImageToDirectory(t *testing.T) {
+	for _, test := range []struct {
+		name     string
+		image    map[string]string
+		paths    []string
+		expected map[string]string
+	}{
+		{
+			name: "passing the root path extracts the whole image",
+			image: map[string]string{
+				"foo/bar/hello.txt": "hello",
+				"world.txt":         "world",
+			},
+			paths: []string{"/"},
+			expected: map[string]string{
+				"foo/bar/hello.txt": "hello",
+				"world.txt":         "world",
+				// ext4 recovery dir
+				"lost+found": testfs.EmptyDir,
+			},
+		},
+		{
+			name: "file paths are allowed",
+			image: map[string]string{
+				"ignore.txt":         "IGNORE_ME",
+				"foo/hello.txt":      "hello",
+				"foo/ignore.txt":     "hello",
+				"foo/bar/ignore.txt": "IGNORE_ME",
+			},
+			paths: []string{"/foo/hello.txt"},
+			expected: map[string]string{
+				"foo/hello.txt": "hello",
+			},
+		},
+		{
+			name: "directory paths are allowed",
+			image: map[string]string{
+				"parent1/child1/hello.txt":  "hello",
+				"parent1/child2/ignore.txt": "IGNORE_ME",
+				"parent2/world.txt":         "world",
+			},
+			paths: []string{"/parent1/child1", "/parent2/"},
+			expected: map[string]string{
+				"parent1/child1/hello.txt": "hello",
+				"parent2/world.txt":        "world",
+			},
+		},
+		{
+			name: "nonexistent paths are silently ignored",
+			image: map[string]string{
+				"parent1/child1/hello.txt": "hello",
+				"world.txt":                "world",
+			},
+			paths: []string{
+				"/does_not_exist",
+				"/parent1/does_not_exist",
+				"/parent1/child1/hello.txt",
+			},
+			expected: map[string]string{
+				"parent1/child1/hello.txt": "hello",
+			},
+		},
+		{
+			name: "empty directories are copied",
+			image: map[string]string{
+				"parent1/emptydir1": testfs.EmptyDir,
+				"parent2/emptydir2": testfs.EmptyDir,
+				"emptydir3":         testfs.EmptyDir,
+			},
+			paths: []string{"/parent1/emptydir1", "/parent2", "/emptydir3"},
+			expected: map[string]string{
+				"parent1/emptydir1": testfs.EmptyDir,
+				"parent2/emptydir2": testfs.EmptyDir,
+				"emptydir3":         testfs.EmptyDir,
+			},
+		},
+		{
+			name: "duplicate paths are allowed",
+			image: map[string]string{
+				"parent1/hello.txt": "hello",
+				"parent2/world.txt": "world",
+			},
+			paths: []string{"/parent1", "/parent1/hello.txt", "/parent2", "/parent2"},
+			expected: map[string]string{
+				"parent1/hello.txt": "hello",
+				"parent2/world.txt": "world",
+			},
+		},
+		{
+			name: "paths not beginning with slash are allowed",
+			image: map[string]string{
+				"parent1/hello.txt": "hello",
+				"parent2/world.txt": "world",
+			},
+			paths: []string{"parent1"},
+			expected: map[string]string{
+				"parent1/hello.txt": "hello",
+			},
+		},
+		{
+			name: "paths beginning with dotslash are allowed",
+			image: map[string]string{
+				"parent1/hello.txt": "hello",
+				"parent2/world.txt": "world",
+			},
+			paths: []string{"./parent1"},
+			expected: map[string]string{
+				"parent1/hello.txt": "hello",
+			},
+		},
+		{
+			name: "empty path is the same as root path",
+			image: map[string]string{
+				"foo/bar/hello.txt": "hello",
+				"world.txt":         "world",
+			},
+			paths: []string{""},
+			expected: map[string]string{
+				"foo/bar/hello.txt": "hello",
+				"world.txt":         "world",
+				// ext4 recovery dir
+				"lost+found": testfs.EmptyDir,
+			},
+		},
+		{
+			name: "dot is the same as root path",
+			image: map[string]string{
+				"foo/bar/hello.txt": "hello",
+				"world.txt":         "world",
+			},
+			paths: []string{"."},
+			expected: map[string]string{
+				"foo/bar/hello.txt": "hello",
+				"world.txt":         "world",
+				// ext4 recovery dir
+				"lost+found": testfs.EmptyDir,
+			},
+		},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			ctx := context.Background()
+			inputDir := testfs.MakeTempDir(t)
+			testfs.WriteAllFileContents(t, inputDir, test.image)
+			tmp := testfs.MakeTempDir(t)
+			imagePath := filepath.Join(tmp, "image.ext4")
+			err := ext4.DirectoryToImageAutoSize(ctx, inputDir, imagePath)
+			require.NoError(t, err)
+			outputDir := testfs.MakeTempDir(t)
+
+			err = ext4.ImageToDirectory(ctx, imagePath, outputDir, test.paths)
+			require.NoError(t, err)
+
+			testfs.AssertExactFileContents(t, outputDir, test.expected)
+		})
+	}
+}

server/testutil/testfs/BUILD

@@ -7,6 +7,7 @@ go_library(
     visibility = ["//visibility:public"],
     deps = [
         "//server/util/random",
+        "@com_github_google_go_containerregistry//pkg/v1/empty",
         "@com_github_stretchr_testify//assert",
         "@com_github_stretchr_testify//require",
         "@io_bazel_rules_go//go/runfiles:go_default_library",