mount_gpt_image.sh: Load partition numbers from the GPT

The recent break (https://chromium-review.googlesource.com/c/441109/)
was caused because the failing parrot test
(https://uberchromegw.corp.google.com/i/chromeos/builders/parrot-paladin/builds/22380/steps/VMTest%20%28attempt%202%29/logs/stdio)
copies a test qemu image out of the build area without also including
the partition script.

The mount_gpt_image.sh script tries to use --from to load the partition
script. When this fails (because it was not copied with the image), it
tries to use --rootfs_mountpt -- which also fails, because the rootfs
isn't actually mounted. This means all of the PARTITION_NUM_* variables
are not defined, which causes the parrot test error.

In general we handled this in the past by using cgpt_py to get the
partition number from the disk_layout.json, but it looks like we can
just read using ${GPT} directly. This avoids needing to supply the image
type needed to satisfy cgpt_py, since we're using the GPT directly.

BUG=b:35045101
TEST=Build parrot qemu test image, verify mount_gpt_image.sh against an
image, and verify the commands succeed:

  b=parrot
  ./build_packages --board=$b
  ./build_image --board=$b test
  ./image_to_vm.sh --board=$b --from=../build/images/$b/latest --test_image
  T=$(mktemp -d);mkdir -p $T/{s,m}
  cp ../build/images/$b/latest/chromiumos_qemu_image.bin $T
  ./mount_gpt_image.sh --from=$T --image=chromiumos_qemu_image.bin \
      --rootfs_mountpt=$T/m --stateful_mountpt=$T/s --safe
  ./mount_gpt_image.sh --from=$T --image=chromiumos_qemu_image.bin \
      --rootfs_mountpt=$T/m --stateful_mountpt=$T/s --safe -u

Change-Id: I90f444c1dd2077474a122d3fc93ca57282f94be5
Reviewed-on: https://chromium-review.googlesource.com/455576
Commit-Ready: Ian Coolidge <icoolidge@google.com>
Tested-by: Ian Coolidge <icoolidge@google.com>
Reviewed-by: Joshua Emele <jemele@google.com>
Reviewed-by: Mike Frysinger <vapier@chromium.org>
diff --git a/build_library/disk_layout_util.sh b/build_library/disk_layout_util.sh
index 88759b1..c69ec0d 100644
--- a/build_library/disk_layout_util.sh
+++ b/build_library/disk_layout_util.sh
@@ -158,6 +158,13 @@
   cgpt_py readlabel "${image_type}" "${DISK_LAYOUT_PATH}" "${part_id}"
 }
 
+get_image_partition_number() {
+  local image="$1"
+  local label="$2"
+  local part=$("${GPT}" find -n -l "${label}" "${image}")
+  echo "${part}"
+}
+
 get_layout_partition_number() {
   local image_type=$1
   local part_label=$2
diff --git a/mount_gpt_image.sh b/mount_gpt_image.sh
index e459a71..630823a 100755
--- a/mount_gpt_image.sh
+++ b/mount_gpt_image.sh
@@ -119,7 +119,7 @@
   return 0
 }
 
-load_partition_script() {
+load_image_partition_numbers() {
   local partition_script="${FLAGS_from}/${FLAGS_partition_script}"
   # Attempt to load the partition script from the rootfs when not found in the
   # FLAGS_from directory.
@@ -129,7 +129,24 @@
   if [[ -f "${partition_script}" ]]; then
     . "${partition_script}"
     load_partition_vars
+    return
   fi
+
+  # Without a partition script, infer numbers from the payload image.
+  local image
+  if [[ -b "${FLAGS_from}" ]]; then
+    image="${FLAGS_from}"
+  else
+    image="${FLAGS_from}/${FLAGS_image}"
+    if [[ ! -f "${image}" ]]; then
+      die "Image ${image} does not exist."
+    fi
+  fi
+  PARTITION_NUM_STATE="$(get_image_partition_number "${image}" "STATE")"
+  PARTITION_NUM_ROOT_A="$(get_image_partition_number "${image}" "ROOT-A")"
+  PARTITION_NUM_OEM="$(get_image_partition_number "${image}" "OEM")"
+  PARTITION_NUM_EFI_SYSTEM="$(get_image_partition_number "${image}" \
+    "EFI-SYSTEM")"
 }
 
 # Common unmounts for either a device or directory
@@ -160,8 +177,6 @@
     fi
   fi
 
-  load_partition_script
-
   # Unmount in reverse order: EFI, OEM, stateful and rootfs.
   local var_name mountpoint fs_format fs_options
   local part_label part_num part_offset part_size data_size
@@ -349,8 +364,6 @@
     ro_rw="ro"
   fi
 
-  load_partition_script
-
   # If all the filesystems are mounted via the kernel, just issue the remount
   # command with the proper flags.
   local var_name mountpoint
@@ -396,6 +409,9 @@
 OEM_MOUNTPOINT="${FLAGS_rootfs_mountpt}/usr/share/oem"
 EFI_SYSTEM_MOUNTPOINT="${FLAGS_esp_mountpt}"
 
+# Read the image partition numbers from the GPT.
+load_image_partition_numbers
+
 # Perform desired operation.
 if [[ ${FLAGS_unmount} -eq ${FLAGS_TRUE} ]]; then
   unmount_image