[avb] Check partition before allocating 64KB memory for VBMeta am: cb00b07ed8

Original change: https://android-review.googlesource.com/c/platform/external/avb/+/2388339

Change-Id: I3da16662b3ee3b408c61b549363cc798d041c340
Signed-off-by: Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com>

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;