From da0c7cffd4145e61bd4acc9c123a04ba53fd0bdd Mon Sep 17 00:00:00 2001
From: Kelvin Zhang <zhangkelvin@google.com>
Date: Wed, 25 Jan 2023 14:50:33 -0800
Subject: Truncate vbmeta size to partition size

Currently, avb assumes all vbmeta* partitions to have size of 64K, which
is maximum size possible. This wastes some memory as avb need to
allocate sufficient memory to hold the entire partition data. Instead,
try to detect the partition size, and only allocate minimum memory
needed to hold the partition data.

Bug: 266757931
Change-Id: If622da2a59861d077d40437c5b0ddb578d7be0c0
---
 libavb/avb_slot_verify.c | 40 ++++++++++++++++++++++------------------
 1 file changed, 22 insertions(+), 18 deletions(-)

diff --git a/libavb/avb_slot_verify.c b/libavb/avb_slot_verify.c
index 3a066f5..a548c80 100644
--- a/libavb/avb_slot_verify.c
+++ b/libavb/avb_slot_verify.c
@@ -558,22 +558,6 @@ 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,
@@ -691,14 +675,34 @@ static AvbSlotVerifyResult load_and_verify_vbmeta(
         vbmeta_size = footer.vbmeta_size;
       }
     }
+  } else {
+    uint64_t partition_size = 0;
+    io_ret =
+        ops->get_size_of_partition(ops, full_partition_name, &partition_size);
+    if (io_ret == AVB_IO_RESULT_OK) {
+      if (partition_size < vbmeta_size && partition_size > 0) {
+        avb_debugv(full_partition_name,
+                   ": Using partition size as vbmeta size\n",
+                   NULL);
+        vbmeta_size = partition_size;
+      }
+    } else {
+      avb_debugv(full_partition_name, ": Failed to get partition size\n", NULL);
+      // libavb might fall back to other partitions if current vbmeta partition
+      // isn't found. So AVB_IO_RESULT_ERROR_NO_SUCH_PARTITION is recoverable,
+      // but other errors are not.
+      if (io_ret != AVB_IO_RESULT_ERROR_NO_SUCH_PARTITION) {
+        ret = AVB_SLOT_VERIFY_RESULT_ERROR_IO;
+        goto out;
+      }
+    }
   }
 
-  /* Read one byte from the partition to check the existence of the
+  /* Use result from previous I/O operation 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) {
-- 
cgit v1.2.3