cos-customizer: Select correct provisioner arch at runtime

We ship both provisioner_amd64 and provisioner_arm64 in cos-customizer,
and we update startup.yaml to select the correct one at runtime, based
on architecture.

Since arm64 ttys have a different name in /dev, we need to template that
into the customizer systemd service.

This change does change the behavior of cos-customizer. However, an
additional change is needed for users to actually use the arm64 code
path. So the behavioral changes shouldn't be noticeable to users right
now.

BUG=b/176990931
TEST=./run_tests

Change-Id: I8e0cd908126d128ee15f789ba72a61aee2c316e0
Reviewed-on: https://cos-review.googlesource.com/c/cos/tools/+/25901
Cloud-Build: GCB Service account <228075978874@cloudbuild.gserviceaccount.com>
Tested-by: Robert Kolchmeyer <rkolchmeyer@google.com>
Reviewed-by: Roy Yang <royyang@google.com>
Reviewed-by: Arnav Kansal <rnv@google.com>
diff --git a/BUILD.bazel b/BUILD.bazel
index 23a1235..93e78a0 100644
--- a/BUILD.bazel
+++ b/BUILD.bazel
@@ -38,6 +38,7 @@
 
 container_image(
     name = "veritysetup_amd64",
+    architecture = "amd64",
     debs = [
         packages_amd64["coreutils"],
         packages_amd64["tar"],
@@ -65,6 +66,7 @@
 
 container_image(
     name = "veritysetup_arm64",
+    architecture = "arm64",
     debs = [
         packages_arm64["coreutils"],
         packages_arm64["tar"],
diff --git a/src/cmd/cos_customizer/cloudbuild.yaml b/src/cmd/cos_customizer/cloudbuild.yaml
index b23c64a..7639ba1 100644
--- a/src/cmd/cos_customizer/cloudbuild.yaml
+++ b/src/cmd/cos_customizer/cloudbuild.yaml
@@ -33,6 +33,8 @@
 options:
   machineType: 'N1_HIGHCPU_8'
   substitutionOption: 'MUST_MATCH'
+# 25 minute timeout; concurrent arm64/amd64 builds take time
+timeout: '1500s'
 images:
 - 'gcr.io/${_OUTPUT_PROJECT}/cos-customizer:${TAG_NAME}'
 - 'gcr.io/${_OUTPUT_PROJECT}/cos-customizer:latest'
diff --git a/src/data/startup.yaml b/src/data/startup.yaml
index ed69f69..7e0ad8d 100644
--- a/src/data/startup.yaml
+++ b/src/data/startup.yaml
@@ -31,6 +31,23 @@
       return "${PIPESTATUS[0]}"
     }
 
+    # To make sure that ARCH gets set properly, this function cannot run in a
+    # subshell.
+    set_arch() {
+      case "$(uname -m)" in
+        x86_64)
+          ARCH=amd64
+          ;;
+        aarch64)
+          ARCH=arm64
+          ;;
+        *)
+          ARCH="$(uname -m)"
+          ;;
+      esac
+      status echo "Set ARCH to ${ARCH}"
+    }
+
     run_provisioner() {
       status $@ && ret=$? || ret=$?
       if [[ "${ret}" != 0 ]]; then
@@ -55,7 +72,7 @@
         # time to capture the serial logs. Once Daisy is done capturing the
         # serial logs, it will add the "DaisyEnd" metadata key. Let's wait for
         # that key to appear (and shutdown anyway after 5 minutes).
-        /mnt/disks/cidata/metadata_watcher DaisyEnd
+        /mnt/disks/cidata/metadata_watcher_${ARCH} DaisyEnd
         umount /mnt/disks/cidata
         rm -r /mnt/disks || :
         shutdown -h now
@@ -65,17 +82,18 @@
 
     main() {
       status history -c
+      set_arch
       status mkdir -p /mnt/disks/cidata
       status mount /dev/disk/by-label/CIDATA /mnt/disks/cidata
       if [[ ! -d /var/lib/.cos-customizer ]]; then
-        run_provisioner /mnt/disks/cidata/provisioner run --config=/mnt/disks/cidata/config.json
+        run_provisioner /mnt/disks/cidata/provisioner_${ARCH} run --config=/mnt/disks/cidata/config.json
       else
-        run_provisioner /mnt/disks/cidata/provisioner resume
+        run_provisioner /mnt/disks/cidata/provisioner_${ARCH} resume
       fi
     }
 
     main
-- path: /etc/systemd/system/customizer.service
+- path: /etc/systemd/system/customizer@.service
   permissions: 0644
   content: |
     [Unit]
@@ -88,12 +106,17 @@
     RemainAfterExit=yes
     User=root
     ExecStart=/bin/bash /tmp/startup.sh
-    ExecStopPost=/bin/bash -c 'rm /etc/systemd/system/customizer.service'
+    ExecStopPost=/bin/bash -c 'rm /etc/systemd/system/customizer@.service'
     StandardOutput=tty
     StandardError=tty
-    TTYPath=/dev/ttyS2
+    TTYPath=%I
 
 runcmd:
 - echo "Starting startup service..."
 - systemctl daemon-reload
-- systemctl --no-block start customizer.service
+- |
+  if [[ "$(uname -m)" == "aarch64" ]]; then
+    systemctl --no-block start customizer@-dev-ttyAMA2.service
+  else
+    systemctl --no-block start customizer@-dev-ttyS2.service
+  fi
diff --git a/src/pkg/preloader/BUILD.bazel b/src/pkg/preloader/BUILD.bazel
index ca79876..c7e110a 100644
--- a/src/pkg/preloader/BUILD.bazel
+++ b/src/pkg/preloader/BUILD.bazel
@@ -19,16 +19,20 @@
     srcs = [
         "//:src/data/startup.yaml",
         "//src/cmd/provisioner:provisioner_amd64",
+        "//src/cmd/provisioner:provisioner_arm64",
         "//src/cmd/metadata_watcher:metadata_watcher_amd64",
+        "//src/cmd/metadata_watcher:metadata_watcher_arm64",
     ],
     outs = ["cidata.img"],
     cmd = "\
-$(location @dosfstools//:mkfs.fat) -n CIDATA -S 512 -s 8 -C $@ 65536;\
+$(location @dosfstools//:mkfs.fat) -n CIDATA -S 512 -s 8 -C $@ 131072;\
 touch meta-data;\
 $(location @mtools//:mcopy) -i $@ $(location //:src/data/startup.yaml) ::/user-data;\
 $(location @mtools//:mcopy) -i $@ meta-data ::/meta-data;\
-$(location @mtools//:mcopy) -i $@ $(location //src/cmd/provisioner:provisioner_amd64) ::/provisioner;\
-$(location @mtools//:mcopy) -i $@ $(location //src/cmd/metadata_watcher:metadata_watcher_amd64) ::/metadata_watcher;",
+$(location @mtools//:mcopy) -i $@ $(location //src/cmd/provisioner:provisioner_amd64) ::/provisioner_amd64;\
+$(location @mtools//:mcopy) -i $@ $(location //src/cmd/provisioner:provisioner_arm64) ::/provisioner_arm64;\
+$(location @mtools//:mcopy) -i $@ $(location //src/cmd/metadata_watcher:metadata_watcher_amd64) ::/metadata_watcher_amd64;\
+$(location @mtools//:mcopy) -i $@ $(location //src/cmd/metadata_watcher:metadata_watcher_arm64) ::/metadata_watcher_arm64;",
     tools = [
         "@dosfstools//:mkfs.fat",
         "@mtools//:mcopy",
diff --git a/src/pkg/tools/seal_oem_partition.go b/src/pkg/tools/seal_oem_partition.go
index 328c622..5628113 100644
--- a/src/pkg/tools/seal_oem_partition.go
+++ b/src/pkg/tools/seal_oem_partition.go
@@ -21,6 +21,7 @@
 	"log"
 	"os"
 	"os/exec"
+	"runtime"
 	"strconv"
 	"strings"
 
@@ -78,7 +79,7 @@
 			"input: imgPath=%q, error msg: (%v)", imgPath, err)
 	}
 	var idBuf bytes.Buffer
-	cmd = exec.Command("sudo", "docker", "images", "veritysetup:veritysetup", "-q")
+	cmd = exec.Command("sudo", "docker", "images", "veritysetup:veritysetup_"+runtime.GOARCH, "-q")
 	cmd.Stdout = &idBuf
 	cmd.Stderr = os.Stderr
 	if err := cmd.Run(); err != nil {