Add support for GRUB2 on ARM64 architecture

ARM64 uses UEFI as a reference boot environment, so add all
the relevant bits to support GRUB2 on ARM64.

BUG=b/229863697
TEST=presubmit
RELEASE_NOTE=None

Change-Id: Ia334b15e5f5ec218fac93e33ff661029c6d30e1c
Reviewed-on: https://cos-review.googlesource.com/c/third_party/platform/crosutils/+/32093
Tested-by: Cusky Presubmit Bot <presubmit@cos-infra-prod.iam.gserviceaccount.com>
Reviewed-by: Robert Kolchmeyer <rkolchmeyer@google.com>
Reviewed-by: Vaibhav Rustagi <vaibhavrustagi@google.com>
diff --git a/bin/cros_make_image_bootable b/bin/cros_make_image_bootable
index 5ff60aa..542229f 100755
--- a/bin/cros_make_image_bootable
+++ b/bin/cros_make_image_bootable
@@ -384,6 +384,8 @@
     kernel_part="--kernel_partition='${FLAGS_output_dir}/vmlinuz.image'"
     # Install syslinux on the EFI System Partition.
     kernel_part="${kernel_part} --install_syslinux"
+  elif [[ "${FLAGS_arch}" = "arm64" ]]; then
+    kernel_part="--kernel_partition='${FLAGS_output_dir}/vmlinuz.image'"
   elif [[ "${FLAGS_arch}" = "arm" || "${FLAGS_arch}" = "mips" ]]; then
     # These flags are not used for ARM / MIPS update_bootloaders.sh
     kernel_part=""
diff --git a/build_library/create_legacy_bootloader_templates.sh b/build_library/create_legacy_bootloader_templates.sh
index 28f692a..99f4123 100755
--- a/build_library/create_legacy_bootloader_templates.sh
+++ b/build_library/create_legacy_bootloader_templates.sh
@@ -272,6 +272,57 @@
   fi
   info "Emitted ${FLAGS_to}/efi/boot/grub.cfg"
   exit 0
+elif [[ "${FLAGS_arch}" = "arm64" ]]; then
+  sudo mkdir -p "${FLAGS_to}"/efi/boot
+
+  # Templated variables:
+  #  DMTABLEA, DMTABLEB -> '0 xxxx verity ... '
+  # This should be replaced during postinst when updating the ESP.
+  cat <<EOF | sudo dd of="${FLAGS_to}/efi/boot/grub.cfg" 2>/dev/null
+defaultA=0
+defaultB=1
+gptpriority \$grubdisk ${partition_num_kern_a} prioA
+gptpriority \$grubdisk ${partition_num_kern_b} prioB
+
+if [ \$prioA -lt \$prioB ]; then
+  set default=\$defaultB
+else
+  set default=\$defaultA
+fi
+
+set timeout=0
+
+# NOTE: These magic grub variables are a Chrome OS hack. They are not portable.
+
+menuentry "local image A" {
+  linux /syslinux/vmlinuz.A ${common_args} cros_efi \
+      root=/dev/\$linuxpartA
+}
+
+menuentry "local image B" {
+  linux /syslinux/vmlinuz.B ${common_args} cros_efi \
+      root=/dev/\$linuxpartB
+}
+
+menuentry "verified image A" {
+  linux /syslinux/vmlinuz.A ${common_args} ${verity_common} \
+      cros_efi root=${ROOTDEV} dm="DMTABLEA"
+}
+
+menuentry "verified image B" {
+  linux /syslinux/vmlinuz.B ${common_args} ${verity_common} \
+      cros_efi root=${ROOTDEV} dm="DMTABLEB"
+}
+
+EOF
+  if [[ ${FLAGS_enable_rootfs_verification} -eq ${FLAGS_TRUE} ]]; then
+    sudo sed -i \
+      -e '/^defaultA=/s:=.*:=2:' \
+      -e '/^defaultB=/s:=.*:=3:' \
+      "${FLAGS_to}/efi/boot/grub.cfg"
+  fi
+  info "Emitted ${FLAGS_to}/efi/boot/grub.cfg"
+  exit 0
 fi
 
 info "The target platform does not use bootloader templates."
diff --git a/update_bootloaders.sh b/update_bootloaders.sh
index 9b49ae7..ff9c1b2 100755
--- a/update_bootloaders.sh
+++ b/update_bootloaders.sh
@@ -126,6 +126,40 @@
   }
 fi
 
+if ! type -p update_arm64_bootloaders; then
+  update_arm64_bootloaders() {
+    local old_root="$1"  # e.g., /dev/sd%D%P or %U+1
+    local kernel_cmdline="$2"
+    local esp_fs_dir="$3"
+    local template_dir="$4"
+    local to="$5"
+
+    # Pull out the dm="" values
+    dm_table=
+    if echo "$kernel_cmdline" | grep -q 'dm="'; then
+      dm_table=$(echo "$kernel_cmdline" | sed -s 's/.*dm="\([^"]*\)".*/\1/')
+    fi
+
+    # Discover last known partition numbers.
+    local partition_num_root_a="$(get_layout_partition_number \
+      "${FLAGS_image_type}" ROOT-A)"
+    local partition_num_root_b="$(get_layout_partition_number \
+      "${FLAGS_image_type}" ROOT-B)"
+    root_a_uuid="PARTUUID=$(part_index_to_uuid "$to" ${partition_num_root_a})"
+    root_b_uuid="PARTUUID=$(part_index_to_uuid "$to" ${partition_num_root_b})"
+
+    # Rewrite grub table
+    grub_dm_table_a=${dm_table//${old_root}/${root_a_uuid}}
+    grub_dm_table_b=${dm_table//${old_root}/${root_b_uuid}}
+    sudo sed -e "s|DMTABLEA|${grub_dm_table_a}|g" \
+        -e "s|DMTABLEB|${grub_dm_table_b}|g" \
+        -e "s|/dev/\\\$linuxpartA|${root_a_uuid}|g" \
+        -e "s|/dev/\\\$linuxpartB|${root_b_uuid}|g" \
+        "${template_dir}"/efi/boot/grub.cfg |
+        sudo dd of="${esp_fs_dir}"/efi/boot/grub.cfg status=none
+  }
+fi
+
 ESP_DEV_OURS=
 ESP_DEV=
 if [[ ! -e "${FLAGS_to}" ]]; then
@@ -216,6 +250,26 @@
     # mount again for cleanup to free resource gracefully
     sudo mount -o ro "${ESP_DEV}" "${ESP_FS_DIR}"
   fi
+elif [[ "${FLAGS_arch}" = "arm64" ]]; then
+  set -x
+  # Populate the EFI bootloader configuration
+  sudo mkdir -p "${ESP_FS_DIR}/efi/boot"
+
+  # Extract kernel flags
+  kernel_cfg=
+  old_root='PARTUUID=%U/PARTNROFF=1'
+  if [[ -n "${FLAGS_kernel_cmdline}" ]]; then
+    info "Using supplied kernel_cmdline to update templates."
+    kernel_cfg="${FLAGS_kernel_cmdline}"
+  elif [[ -n "${FLAGS_kernel_partition}" ]]; then
+    info "Extracting the kernel command line from ${FLAGS_kernel_partition}"
+    kernel_cfg=$(dump_kernel_config "${FLAGS_kernel_partition}")
+  fi
+  update_arm64_bootloaders "${old_root}" \
+                         "${kernel_cfg}" \
+                         "${ESP_FS_DIR}" \
+                         "${FLAGS_from}" \
+                         "${FLAGS_to}"
 elif [[ "${FLAGS_arch}" = "arm" || "${FLAGS_arch}" = "mips" ]]; then
   # Copy u-boot script to ESP partition
   if [ -r "${FLAGS_from}/boot-A.scr.uimg" ]; then