[avb] Check partition before allocating 64KB memory for VBMetamain-16k-with-phones

This cl is an optimization for load_and_verify_vbmeta(). When
we want to load the VBMeta from the partition `vbmeta` and we
are the main VBMeta, prior to this cl, we first allocated 64KB
memory and then tried to read the partition `vbmeta`.

This cl added a check before this action. After this cl, we
first try to read 1 byte from the partition, if this fails
because no such partition exists, we fallback directly on
the `boot` partition. With this approach, we avoid allocating
the big chunk of memory on heap for VBMeta when the partition
doesn't exist.

This unblocks the cl aosp/2388176 where we don't have the
`vbmeta` partition and also cannot afford allocating the 64KB
memory in pVM firmware.

Bug: 256148034
Test: atest libavb_host_unittest libpvmfw_avb.integration_test
Change-Id: Ia242ddedc8bb5e8a4277a9a766d563d8ae8cdd60

1 files changed, 49 insertions, 22 deletions

diff --git a/libavb/avb_slot_verify.c b/libavb/avb_slot_verify.c
index 8e0721d..3a066f5 100644
--- a/libavb/avb_slot_verify.c
+++ b/libavb/avb_slot_verify.c
@@ -43,6 +43,9 @@
 /* Maximum size of a vbmeta image - 64 KiB. */
 #define VBMETA_MAX_SIZE (64 * 1024)
 
+/* Test buffer used to check the existence of a partition. */
+#define TEST_BUFFER_SIZE 1
+
 static AvbSlotVerifyResult initialize_persistent_digest(
     AvbOps* ops,
     const char* part_name,
@@ -555,6 +558,22 @@ out:
   return ret;
 }
 
+static AvbIOResult read_one_byte_from_partition(AvbOps* ops,
+                                                const char* partition_name) {
+  /* test_buf is a one-byte buffer used to check the existence of the
+   * current partition before allocating the big chunk of memory on heap
+   * for vbmeta later.
+   */
+  uint8_t test_buf[TEST_BUFFER_SIZE];
+  size_t test_num_read;
+  return ops->read_from_partition(ops,
+                                  partition_name,
+                                  0 /* offset */,
+                                  TEST_BUFFER_SIZE,
+                                  test_buf,
+                                  &test_num_read);
+}
+
 static AvbSlotVerifyResult load_and_verify_vbmeta(
     AvbOps* ops,
     const char* const* requested_partitions,
@@ -674,30 +693,38 @@ static AvbSlotVerifyResult load_and_verify_vbmeta(
     }
   }
 
-  vbmeta_buf = avb_malloc(vbmeta_size);
-  if (vbmeta_buf == NULL) {
-    ret = AVB_SLOT_VERIFY_RESULT_ERROR_OOM;
-    goto out;
-  }
+  /* Read one byte from the partition to check the existence of the
+   * partition before allocating the big chunk of memory on heap
+   * for vbmeta later. `io_ret` will be used later to decide whether
+   * to fallback on the `boot` partition.
+   */
+  io_ret = read_one_byte_from_partition(ops, full_partition_name);
+  if (io_ret != AVB_IO_RESULT_ERROR_NO_SUCH_PARTITION) {
+    vbmeta_buf = avb_malloc(vbmeta_size);
+    if (vbmeta_buf == NULL) {
+      ret = AVB_SLOT_VERIFY_RESULT_ERROR_OOM;
+      goto out;
+    }
 
-  if (vbmeta_offset != 0) {
-    avb_debugv("Loading vbmeta struct in footer from partition '",
-               full_partition_name,
-               "'.\n",
-               NULL);
-  } else {
-    avb_debugv("Loading vbmeta struct from partition '",
-               full_partition_name,
-               "'.\n",
-               NULL);
-  }
+    if (vbmeta_offset != 0) {
+      avb_debugv("Loading vbmeta struct in footer from partition '",
+                 full_partition_name,
+                 "'.\n",
+                 NULL);
+    } else {
+      avb_debugv("Loading vbmeta struct from partition '",
+                 full_partition_name,
+                 "'.\n",
+                 NULL);
+    }
 
-  io_ret = ops->read_from_partition(ops,
-                                    full_partition_name,
-                                    vbmeta_offset,
-                                    vbmeta_size,
-                                    vbmeta_buf,
-                                    &vbmeta_num_read);
+    io_ret = ops->read_from_partition(ops,
+                                      full_partition_name,
+                                      vbmeta_offset,
+                                      vbmeta_size,
+                                      vbmeta_buf,
+                                      &vbmeta_num_read);
+  }
   if (io_ret == AVB_IO_RESULT_ERROR_OOM) {
     ret = AVB_SLOT_VERIFY_RESULT_ERROR_OOM;
     goto out;