cos-customizer: Refactor embeddings into main packages

Embeddings need to be compiled for the same platform as the executable
that they are embedded into. This is simpler to do when embeddings are
dependencies of binary targets instead of library targets.

One might say that this can be trivially accomplished with Bazel's
"--platforms" flag. This is not the case, because cos-customizer has two
parts; the part that runs in Cloud Build, and the part that runs on the
preload VM. The Cloud Build part must always be compiled for amd64,
while the preload VM part needs to be compiled for both amd64 and arm64.
The --platforms flag can't handle this kind of configuration. We will
take an alternative approach where binary targets will be duplicated
into amd64 and arm64 variants, and the final cos-customizer container
image will ultimately depend on both kinds of binary targets.

BUG=b/176990931
TEST=./run_tests

Change-Id: I5d92483503d834984bb7df5c78cb7c9590bd708c
Reviewed-on: https://cos-review.googlesource.com/c/cos/tools/+/25830
Cloud-Build: GCB Service account <228075978874@cloudbuild.gserviceaccount.com>
Reviewed-by: Roy Yang <royyang@google.com>
Tested-by: Robert Kolchmeyer <rkolchmeyer@google.com>
diff --git a/src/cmd/provisioner/BUILD.bazel b/src/cmd/provisioner/BUILD.bazel
index 5d3b8ae..997a236 100644
--- a/src/cmd/provisioner/BUILD.bazel
+++ b/src/cmd/provisioner/BUILD.bazel
@@ -13,6 +13,32 @@
 # limitations under the License.
 
 load("@io_bazel_rules_go//go:def.bzl", "go_binary", "go_library")
+load("@io_bazel_rules_go//extras:embed_data.bzl", "go_embed_data")
+
+# Our goal is for this program to be embedded into this Go package. Go embed
+# only allows files in the same package directory to be embedded. So we need to
+# use a "no-op" genrule to place this binary in the same directory as the
+# package source.
+genrule(
+    name = "handle_disk_layout.bin",
+    srcs = ["//src/cmd/handle_disk_layout:handle_disk_layout_bin"],
+    outs = ["_handle_disk_layout.bin"],
+    cmd = "cp $< $@",
+)
+
+genrule(
+    name = "veritysetup.img",
+    srcs = ["//:veritysetup.tar"],
+    outs = ["_veritysetup.img"],
+    cmd = "cp $< $@",
+)
+
+genrule(
+    name = "docker_credential_gcr",
+    srcs = ["@com_github_googlecloudplatform_docker_credential_gcr//:docker-credential-gcr"],
+    outs = ["docker-credential-gcr"],
+    cmd = "cp $< $@",
+)
 
 go_library(
     name = "provisioner_lib",
@@ -21,6 +47,11 @@
         "resume.go",
         "run.go",
     ],
+    embedsrcs = [
+        ":handle_disk_layout.bin",
+        ":veritysetup.img",
+        ":docker_credential_gcr",
+    ],
     importpath = "cos.googlesource.com/cos/tools.git/src/cmd/provisioner",
     visibility = ["//visibility:private"],
     deps = [
diff --git a/src/cmd/provisioner/main.go b/src/cmd/provisioner/main.go
index 983d341..7d37e91 100644
--- a/src/cmd/provisioner/main.go
+++ b/src/cmd/provisioner/main.go
@@ -18,6 +18,7 @@
 
 import (
 	"context"
+	_ "embed"
 	"flag"
 	"log"
 	"os"
@@ -28,6 +29,15 @@
 	"cos.googlesource.com/cos/tools.git/src/pkg/provisioner"
 )
 
+//go:embed _handle_disk_layout.bin
+var handleDiskLayoutBin []byte
+
+//go:embed _veritysetup.img
+var veritySetupImage []byte
+
+//go:embed docker-credential-gcr
+var dockerCredentialGCR []byte
+
 var (
 	stateDir = flag.String("state-dir", "/var/lib/.cos-customizer", "Absolute path to the directory to use for provisioner state. "+
 		"This directory is used for persisting internal state across reboots, unpacking inputs, and running provisioning scripts. "+
@@ -48,14 +58,17 @@
 		os.Exit(int(subcommands.ExitFailure))
 	}
 	deps := provisioner.Deps{
-		GCSClient:    gcsClient,
-		TarCmd:       "tar",
-		SystemctlCmd: "systemctl",
-		RootdevCmd:   "rootdev",
-		CgptCmd:      "cgpt",
-		Resize2fsCmd: "resize2fs",
-		E2fsckCmd:    "e2fsck",
-		RootDir:      "/",
+		GCSClient:           gcsClient,
+		TarCmd:              "tar",
+		SystemctlCmd:        "systemctl",
+		RootdevCmd:          "rootdev",
+		CgptCmd:             "cgpt",
+		Resize2fsCmd:        "resize2fs",
+		E2fsckCmd:           "e2fsck",
+		RootDir:             "/",
+		DockerCredentialGCR: dockerCredentialGCR,
+		VeritySetupImage:    veritySetupImage,
+		HandleDiskLayoutBin: handleDiskLayoutBin,
 	}
 	var exitCode int
 	ret := subcommands.Execute(ctx, deps, &exitCode)
diff --git a/src/pkg/provisioner/BUILD.bazel b/src/pkg/provisioner/BUILD.bazel
index d19dd8d..cf40efa 100644
--- a/src/pkg/provisioner/BUILD.bazel
+++ b/src/pkg/provisioner/BUILD.bazel
@@ -13,32 +13,6 @@
 # limitations under the License.
 
 load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
-load("@io_bazel_rules_go//extras:embed_data.bzl", "go_embed_data")
-
-# Our goal is for this program to be embedded into this Go package. Go embed
-# only allows files in the same package directory to be embedded. So we need to
-# use a "no-op" genrule to place this binary in the same directory as the
-# package source.
-genrule(
-    name = "handle_disk_layout.bin",
-    srcs = ["//src/cmd/handle_disk_layout:handle_disk_layout_bin"],
-    outs = ["_handle_disk_layout.bin"],
-    cmd = "cp $< $@",
-)
-
-genrule(
-    name = "veritysetup.img",
-    srcs = ["//:veritysetup.tar"],
-    outs = ["_veritysetup.img"],
-    cmd = "cp $< $@",
-)
-
-genrule(
-    name = "docker_credential_gcr",
-    srcs = ["@com_github_googlecloudplatform_docker_credential_gcr//:docker-credential-gcr"],
-    outs = ["docker-credential-gcr"],
-    cmd = "cp $< $@",
-)
 
 go_library(
     name = "provisioner",
@@ -56,11 +30,6 @@
         "state.go",
         "systemd.go",
     ],
-    embedsrcs = [
-        ":handle_disk_layout.bin",
-        ":veritysetup.img",
-        ":docker_credential_gcr",
-    ],
     importpath = "cos.googlesource.com/cos/tools.git/src/pkg/provisioner",
     visibility = ["//visibility:public"],
     deps = [
diff --git a/src/pkg/provisioner/config.go b/src/pkg/provisioner/config.go
index 3621d45..a0ccf09 100644
--- a/src/pkg/provisioner/config.go
+++ b/src/pkg/provisioner/config.go
@@ -94,6 +94,9 @@
 type stepDeps struct {
 	// GCSClient is used to access Google Cloud Storage.
 	GCSClient *storage.Client
+	// VeritySetupImage is an embedded Docker image that contains the
+	// "veritysetup" tool.
+	VeritySetupImage []byte
 }
 
 type step interface {
diff --git a/src/pkg/provisioner/disk_layout.go b/src/pkg/provisioner/disk_layout.go
index d45eed5..5cb418a 100644
--- a/src/pkg/provisioner/disk_layout.go
+++ b/src/pkg/provisioner/disk_layout.go
@@ -15,7 +15,6 @@
 package provisioner
 
 import (
-	_ "embed"
 	"errors"
 	"fmt"
 	"io"
@@ -34,9 +33,6 @@
 	"cos.googlesource.com/cos/tools.git/src/pkg/utils"
 )
 
-//go:embed _handle_disk_layout.bin
-var handleDiskLayoutBin []byte
-
 func switchRoot(deps Deps, runState *state) (err error) {
 	if !runState.data.Config.BootDisk.ReclaimSDA3 {
 		log.Println("ReclaimSDA3 is not set, not switching root device")
@@ -105,7 +101,7 @@
 	if err := mountFunc("", filepath.Join(deps.RootDir, "tmp"), "", unix.MS_REMOUNT|unix.MS_NOSUID|unix.MS_NODEV, ""); err != nil {
 		return fmt.Errorf("error remounting /tmp as exec: %v", err)
 	}
-	if err := ioutil.WriteFile(filepath.Join(deps.RootDir, "tmp", "handle_disk_layout.bin"), handleDiskLayoutBin, 0744); err != nil {
+	if err := ioutil.WriteFile(filepath.Join(deps.RootDir, "tmp", "handle_disk_layout.bin"), deps.HandleDiskLayoutBin, 0744); err != nil {
 		return err
 	}
 	data := fmt.Sprintf(`[Unit]
diff --git a/src/pkg/provisioner/provisioner.go b/src/pkg/provisioner/provisioner.go
index 0fc4491..954a5eb 100644
--- a/src/pkg/provisioner/provisioner.go
+++ b/src/pkg/provisioner/provisioner.go
@@ -19,7 +19,6 @@
 import (
 	"bufio"
 	"context"
-	_ "embed"
 	"errors"
 	"fmt"
 	"io/ioutil"
@@ -34,9 +33,6 @@
 	"golang.org/x/sys/unix"
 )
 
-//go:embed docker-credential-gcr
-var dockerCredentialGCR []byte
-
 // ErrRebootRequired indicates that a reboot is necessary for provisioning to
 // continue.
 var ErrRebootRequired = errors.New("reboot required to continue provisioning")
@@ -103,12 +99,12 @@
 	return parsedOptions, nil
 }
 
-func setup(runState *state, rootDir string, systemd *systemdClient) error {
+func setup(runState *state, deps Deps, systemd *systemdClient) error {
 	log.Println("Setting up environment...")
 	if err := systemd.stop("update-engine.service"); err != nil {
 		return err
 	}
-	if err := mountFunc("tmpfs", filepath.Join(rootDir, "root"), "tmpfs", 0, ""); err != nil {
+	if err := mountFunc("tmpfs", filepath.Join(deps.RootDir, "root"), "tmpfs", 0, ""); err != nil {
 		return fmt.Errorf("error mounting tmpfs at /root: %v", err)
 	}
 	binPath := filepath.Join(runState.dir, "bin")
@@ -119,7 +115,7 @@
 		}
 	}
 	if _, err := os.Stat(dockerCredentialGCRPath); os.IsNotExist(err) {
-		if err := ioutil.WriteFile(dockerCredentialGCRPath, dockerCredentialGCR, 0744); err != nil {
+		if err := ioutil.WriteFile(dockerCredentialGCRPath, deps.DockerCredentialGCR, 0744); err != nil {
 			return err
 		}
 	}
@@ -133,7 +129,7 @@
 	if err := mountFunc(binPath, binPath, "ext4", unix.MS_BIND, ""); err != nil {
 		return fmt.Errorf("error bind mounting %q: %v", dockerCredentialGCRPath, err)
 	}
-	opts, err := mountOptions(rootDir, binPath)
+	opts, err := mountOptions(deps.RootDir, binPath)
 	if err != nil {
 		return err
 	}
@@ -338,6 +334,15 @@
 	// RootDir is the path to the root file system. Should be "/" in all real
 	// runtime situations.
 	RootDir string
+	// DockerCredentialGCR is an embedded docker-credential-gcr program to use as a Docker
+	// credential helper.
+	DockerCredentialGCR []byte
+	// VeritySetupImage is an embedded Docker image that contains the
+	// "veritysetup" tool.
+	VeritySetupImage []byte
+	// HandleDiskLayoutBin is an embedded program for reformatting a COS disk
+	// image.
+	HandleDiskLayoutBin []byte
 }
 
 func run(ctx context.Context, deps Deps, runState *state) (err error) {
@@ -345,10 +350,13 @@
 	if err := repartitionBootDisk(deps, runState); err != nil {
 		return err
 	}
-	if err := setup(runState, deps.RootDir, systemd); err != nil {
+	if err := setup(runState, deps, systemd); err != nil {
 		return err
 	}
-	stepDeps := stepDeps{GCSClient: deps.GCSClient}
+	stepDeps := stepDeps{
+		GCSClient:        deps.GCSClient,
+		VeritySetupImage: deps.VeritySetupImage,
+	}
 	if err := executeSteps(ctx, runState, stepDeps); err != nil {
 		return err
 	}
diff --git a/src/pkg/provisioner/provisioner_test.go b/src/pkg/provisioner/provisioner_test.go
index 3a8308f..92af977 100644
--- a/src/pkg/provisioner/provisioner_test.go
+++ b/src/pkg/provisioner/provisioner_test.go
@@ -27,6 +27,8 @@
 	"golang.org/x/sys/unix"
 )
 
+const trueExecutable = "#!/bin/bash\ntrue"
+
 func testDataDir(t *testing.T) string {
 	t.Helper()
 	path, err := filepath.Abs("testdata")
@@ -68,9 +70,12 @@
 		t.Fatal(err)
 	}
 	deps := Deps{
-		GCSClient:    nil,
-		TarCmd:       "",
-		SystemctlCmd: "",
+		GCSClient:           nil,
+		TarCmd:              "",
+		SystemctlCmd:        "",
+		DockerCredentialGCR: []byte(trueExecutable),
+		VeritySetupImage:    []byte(trueExecutable),
+		HandleDiskLayoutBin: []byte(trueExecutable),
 	}
 	config := Config{}
 	if err := Run(ctx, deps, dir, config); err != errStateAlreadyExists {
@@ -109,10 +114,13 @@
 			defer os.RemoveAll(tempDir)
 			gcs := fakes.GCSForTest(t)
 			deps := Deps{
-				GCSClient:    gcs.Client,
-				TarCmd:       "tar",
-				SystemctlCmd: "/bin/true",
-				RootDir:      tempDir,
+				GCSClient:           gcs.Client,
+				TarCmd:              "tar",
+				SystemctlCmd:        "/bin/true",
+				RootDir:             tempDir,
+				DockerCredentialGCR: []byte(trueExecutable),
+				VeritySetupImage:    []byte(trueExecutable),
+				HandleDiskLayoutBin: []byte(trueExecutable),
 			}
 			stateDir := filepath.Join(tempDir, "var", "lib", ".cos-customizer")
 			if err := stubMountInfo(filepath.Join(tempDir, "proc", "self", "mountinfo"), filepath.Join(stateDir, "bin")); err != nil {
@@ -181,10 +189,13 @@
 				gcs.Objects[name] = data
 			}
 			deps := Deps{
-				GCSClient:    gcs.Client,
-				TarCmd:       "tar",
-				SystemctlCmd: "/bin/true",
-				RootDir:      tempDir,
+				GCSClient:           gcs.Client,
+				TarCmd:              "tar",
+				SystemctlCmd:        "/bin/true",
+				RootDir:             tempDir,
+				DockerCredentialGCR: []byte(trueExecutable),
+				VeritySetupImage:    []byte(trueExecutable),
+				HandleDiskLayoutBin: []byte(trueExecutable),
 			}
 			stateDir := filepath.Join(tempDir, "var", "lib", ".cos-customizer")
 			if err := stubMountInfo(filepath.Join(tempDir, "proc", "self", "mountinfo"), filepath.Join(stateDir, "bin")); err != nil {
@@ -261,10 +272,13 @@
 				gcs.Objects[name] = data
 			}
 			deps := Deps{
-				GCSClient:    gcs.Client,
-				TarCmd:       "tar",
-				SystemctlCmd: "/bin/true",
-				RootDir:      tempDir,
+				GCSClient:           gcs.Client,
+				TarCmd:              "tar",
+				SystemctlCmd:        "/bin/true",
+				RootDir:             tempDir,
+				DockerCredentialGCR: []byte(trueExecutable),
+				VeritySetupImage:    []byte(trueExecutable),
+				HandleDiskLayoutBin: []byte(trueExecutable),
 			}
 			stateDir := filepath.Join(tempDir, "var", "lib", ".cos-customizer")
 			if err := stubMountInfo(filepath.Join(tempDir, "proc", "self", "mountinfo"), filepath.Join(stateDir, "bin")); err != nil {
diff --git a/src/pkg/provisioner/seal_oem_step.go b/src/pkg/provisioner/seal_oem_step.go
index ce8e91a..9c1415e 100644
--- a/src/pkg/provisioner/seal_oem_step.go
+++ b/src/pkg/provisioner/seal_oem_step.go
@@ -16,7 +16,6 @@
 
 import (
 	"context"
-	_ "embed"
 	"io/ioutil"
 	"log"
 	"os"
@@ -25,16 +24,13 @@
 	"cos.googlesource.com/cos/tools.git/src/pkg/tools"
 )
 
-//go:embed _veritysetup.img
-var veritysetupImg []byte
-
 type SealOEMStep struct{}
 
 func (s *SealOEMStep) run(ctx context.Context, runState *state, deps *stepDeps) error {
 	log.Println("Sealing the OEM partition with dm-verity")
 	veritysetupImgPath := filepath.Join(runState.dir, "veritysetup.img")
 	if _, err := os.Stat(veritysetupImgPath); os.IsNotExist(err) {
-		if err := ioutil.WriteFile(veritysetupImgPath, veritysetupImg, 0644); err != nil {
+		if err := ioutil.WriteFile(veritysetupImgPath, deps.VeritySetupImage, 0644); err != nil {
 			return err
 		}
 	}