QcomModulePkg: Get boot and vendor_boot actual size for image verification

For optimising boot up time, using image actual size instead of the
whole partition size for boot and vendor_boot image verification.

Change-Id: I89edd58ade815f10d5b28c39bd8da084a41f1294

3 files changed, 201 insertions, 3 deletions

diff --git a/QcomModulePkg/Library/avb/libavb/avb_ops.c b/QcomModulePkg/Library/avb/libavb/avb_ops.c
index fd5c828b82..f26e4b186a 100644..100755
--- a/QcomModulePkg/Library/avb/libavb/avb_ops.c
+++ b/QcomModulePkg/Library/avb/libavb/avb_ops.c
@@ -22,7 +22,7 @@
  * SOFTWARE.
  */
 
-/* Copyright (c) 2017-2020, The Linux Foundation. All rights reserved.
+/* Copyright (c) 2017-2021, The Linux Foundation. All rights reserved.
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are
@@ -51,6 +51,7 @@
 */
 
 #include "Board.h"
+#include "BootImage.h"
 #include "BootLinux.h"
 #include "LinuxLoaderLib.h"
 #include "OEMPublicKey.h"
@@ -63,6 +64,16 @@
 #include <Library/MemoryAllocationLib.h>
 #include <Uefi.h>
 
+#define ADD_TO_SIZE(OutSizeNumBytes, PageSize, tmp_size) do { \
+	OutSizeNumBytes = ADD_OF (OutSizeNumBytes, \
+														ROUND_TO_PAGE (tmp_size, PageSize - 1)); \
+	if (!OutSizeNumBytes) { \
+		DEBUG ((EFI_D_ERROR, \
+		"Integer Overflow: OutSizeNumBytes=%u, tmp_size=%u\n", \
+		OutSizeNumBytes, ROUND_TO_PAGE (tmp_size, PageSize - 1))); \
+		return AVB_IO_RESULT_ERROR_OOM; } \
+	} while(0)
+
 STATIC AvbIOResult GetHandleInfo(const char *Partition, HandleInfo *HandleInfo)
 {
 	EFI_STATUS Status = EFI_SUCCESS;
@@ -573,6 +584,146 @@ AvbIOResult AvbGetUniqueGuidForPartition(AvbOps *Ops, const char *PartitionName,
 	return AVB_IO_RESULT_ERROR_NO_SUCH_PARTITION;
 }
 
+AvbIOResult AvbGetActualSizeOfBootImage(AvbOps *Ops,
+				const char *PARTITION, uint64_t *OutSizeNumBytes)
+{
+	AvbIOResult Result = AVB_IO_RESULT_OK;
+	HandleInfo HandleInfoList[1];
+	EFI_BLOCK_IO_PROTOCOL *BlockIo = NULL;
+	EFI_STATUS Status = EFI_SUCCESS;
+	UINT32 PageSize = 0;
+	VOID *ImageHdrBuffer = NULL;
+	UINT32 ImageHdrSize = 0;
+	CHAR16 PartitionName[MAX_GPT_NAME_SIZE];
+
+	// Boot Image header information variables
+	UINT32 HeaderVersion = 0;
+	UINT32 HeaderSize = 0;
+
+	if (Ops == NULL || PARTITION == NULL || OutSizeNumBytes == NULL) {
+		DEBUG((EFI_D_ERROR, "AvbGetSizeOfPartition invalid parameter pointers\n"));
+		return AVB_IO_RESULT_ERROR_IO;
+	}
+
+	Result = GetHandleInfo(PARTITION, HandleInfoList);
+	if (Result != AVB_IO_RESULT_OK) {
+		DEBUG((EFI_D_ERROR, "AvbGetSizeOfImage: GetHandleInfo failed"));
+		return Result;
+	}
+
+	*OutSizeNumBytes = 0;
+	BlockIo = HandleInfoList[0].BlkIo;
+	PageSize = BlockIo->Media->BlockSize;
+
+	AsciiStrToUnicodeStr (PARTITION, PartitionName);
+	Status = LoadImageHeader (PartitionName, &ImageHdrBuffer, &ImageHdrSize);
+	if (Status != EFI_SUCCESS ||
+			ImageHdrBuffer ==  NULL ||
+			ImageHdrSize < sizeof (boot_img_hdr)) {
+		DEBUG ((EFI_D_ERROR, "ERROR: Failed to load image header: %r\n", Status));
+		goto Out;
+	}
+
+	if (CompareMem ((VOID *)((boot_img_hdr *)(ImageHdrBuffer))->magic,
+					BOOT_MAGIC, BOOT_MAGIC_SIZE) == 0) {
+		HeaderVersion = ((boot_img_hdr *)(ImageHdrBuffer))->header_version;
+	} else if (CompareMem ((VOID *)((vendor_boot_img_hdr_v3 *)
+					(ImageHdrBuffer))->magic,
+					VENDOR_BOOT_MAGIC,VENDOR_BOOT_MAGIC_SIZE) == 0) {
+		HeaderVersion =
+				((vendor_boot_img_hdr_v3 *)(ImageHdrBuffer))->header_version;
+	} else {
+		DEBUG ((EFI_D_ERROR, "Invalid boot image header\n"));
+		goto Out;
+	}
+
+	if (HeaderVersion < BOOT_HEADER_VERSION_THREE) {
+		PageSize = ((boot_img_hdr *)(ImageHdrBuffer))->page_size;
+		HeaderSize = PageSize;
+
+		ADD_TO_SIZE(*OutSizeNumBytes, PageSize,
+							((boot_img_hdr *)(ImageHdrBuffer))->kernel_size);
+		ADD_TO_SIZE(*OutSizeNumBytes, PageSize,
+							((boot_img_hdr *)(ImageHdrBuffer))->ramdisk_size);
+		ADD_TO_SIZE(*OutSizeNumBytes, PageSize,
+							((boot_img_hdr *)(ImageHdrBuffer))->second_size);
+
+		if (HeaderVersion != BOOT_HEADER_VERSION_ZERO) {
+			ADD_TO_SIZE(*OutSizeNumBytes, PageSize,
+						((struct boot_img_hdr_v1 *)(ImageHdrBuffer))->recovery_dtbo_size);
+			HeaderSize = ((struct boot_img_hdr_v1 *)(ImageHdrBuffer))->header_size;
+		}
+
+		ADD_TO_SIZE(*OutSizeNumBytes, PageSize, HeaderSize);
+
+		if (HeaderVersion == BOOT_HEADER_VERSION_TWO) {
+			ADD_TO_SIZE(*OutSizeNumBytes, PageSize,
+								((struct boot_img_hdr_v2 *)(ImageHdrBuffer))->dtb_size);
+		}
+	} else if (HeaderVersion == BOOT_HEADER_VERSION_THREE) {
+		if (!AsciiStrnCmp (PARTITION, "boot", AsciiStrLen ("boot")) ||
+				!AsciiStrnCmp (PARTITION, "recovery", AsciiStrLen ("recovery"))) {
+			ADD_TO_SIZE(*OutSizeNumBytes, PageSize,
+							((boot_img_hdr_v3 *)(ImageHdrBuffer))->header_size);
+			ADD_TO_SIZE(*OutSizeNumBytes, PageSize,
+							((boot_img_hdr_v3 *)(ImageHdrBuffer))->kernel_size);
+			ADD_TO_SIZE(*OutSizeNumBytes, PageSize,
+							((boot_img_hdr_v3 *)(ImageHdrBuffer))->ramdisk_size);
+		} else if (!AsciiStrnCmp (PARTITION, "vendor_boot",
+								AsciiStrLen ("vendor_boot"))) {
+			PageSize = ((vendor_boot_img_hdr_v3 *)(ImageHdrBuffer))->page_size;
+			ADD_TO_SIZE(*OutSizeNumBytes, PageSize,
+						((vendor_boot_img_hdr_v3 *)(ImageHdrBuffer))->header_size);
+			ADD_TO_SIZE(*OutSizeNumBytes, PageSize,
+						((vendor_boot_img_hdr_v3 *)(ImageHdrBuffer))->vendor_ramdisk_size);
+			ADD_TO_SIZE(*OutSizeNumBytes, PageSize,
+						((vendor_boot_img_hdr_v3 *)(ImageHdrBuffer))->dtb_size);
+		}
+	} else if (HeaderVersion == BOOT_HEADER_VERSION_FOUR) {
+		if (!AsciiStrnCmp (PARTITION, "boot", AsciiStrLen ("boot")) ||
+				!AsciiStrnCmp (PARTITION, "recovery", AsciiStrLen ("recovery"))) {
+			ADD_TO_SIZE(*OutSizeNumBytes, PageSize,
+							((boot_img_hdr_v4 *)(ImageHdrBuffer))->header_size);
+			ADD_TO_SIZE(*OutSizeNumBytes, PageSize,
+							((boot_img_hdr_v4 *)(ImageHdrBuffer))->kernel_size);
+			ADD_TO_SIZE(*OutSizeNumBytes, PageSize,
+							((boot_img_hdr_v4 *)(ImageHdrBuffer))->ramdisk_size);
+
+			ADD_TO_SIZE(*OutSizeNumBytes, PageSize,
+							((boot_img_hdr_v4 *)(ImageHdrBuffer))->signature_size);
+		} else  if (!AsciiStrnCmp (PARTITION, "vendor_boot",
+								AsciiStrLen ("vendor_boot"))) {
+			PageSize = ((vendor_boot_img_hdr_v4 *)(ImageHdrBuffer))->page_size;
+
+			ADD_TO_SIZE(*OutSizeNumBytes, PageSize,
+							((vendor_boot_img_hdr_v4 *)(ImageHdrBuffer))->header_size);
+
+			ADD_TO_SIZE(*OutSizeNumBytes, PageSize,
+							((vendor_boot_img_hdr_v4 *)
+							(ImageHdrBuffer))->vendor_ramdisk_size);
+
+			ADD_TO_SIZE(*OutSizeNumBytes, PageSize,
+							((vendor_boot_img_hdr_v4 *)
+							(ImageHdrBuffer))->VendorRamdiskTableSize);
+			ADD_TO_SIZE(*OutSizeNumBytes, PageSize,
+							((vendor_boot_img_hdr_v4 *)
+							(ImageHdrBuffer))->VendorRamdiskTableEntrySize);
+			ADD_TO_SIZE(*OutSizeNumBytes, PageSize,
+							((vendor_boot_img_hdr_v4 *)
+							(ImageHdrBuffer))->VendorBootconfigSize);
+			ADD_TO_SIZE(*OutSizeNumBytes, PageSize,
+							((vendor_boot_img_hdr_v4 *)(ImageHdrBuffer))->dtb_size);
+		}
+	}
+
+Out:
+	if (*OutSizeNumBytes == 0) {
+		return AVB_IO_RESULT_ERROR_RANGE_OUTSIDE_PARTITION;
+	} else {
+		return AVB_IO_RESULT_OK;
+	}
+}
+
 AvbIOResult AvbGetSizeOfPartition(AvbOps *Ops, const char *Partition, uint64_t *OutSizeNumBytes)
 {
 	AvbIOResult Result = AVB_IO_RESULT_OK;
@@ -620,6 +771,7 @@ AvbOps *AvbOpsNew(VOID *UserData)
 	Ops->read_is_device_unlocked = AvbReadIsDeviceUnlocked;
 	Ops->get_unique_guid_for_partition = AvbGetUniqueGuidForPartition;
 	Ops->get_size_of_partition = AvbGetSizeOfPartition;
+	Ops->get_actual_size_of_boot_image = AvbGetActualSizeOfBootImage;
 
 out:
 	return Ops;
diff --git a/QcomModulePkg/Library/avb/libavb/avb_ops.h b/QcomModulePkg/Library/avb/libavb/avb_ops.h
index 76cc9d47af..51055ff949 100644..100755
--- a/QcomModulePkg/Library/avb/libavb/avb_ops.h
+++ b/QcomModulePkg/Library/avb/libavb/avb_ops.h
@@ -22,6 +22,34 @@
  * SOFTWARE.
  */
 
+/* Copyright (c) 2021, The Linux Foundation. All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are
+ * met:
+ * * Redistributions of source code must retain the above copyright
+ *  notice, this list of conditions and the following disclaimer.
+ *  * Redistributions in binary form must reproduce the above
+ * copyright notice, this list of conditions and the following
+ * disclaimer in the documentation and/or other materials provided
+ *  with the distribution.
+ *   * Neither the name of The Linux Foundation nor the names of its
+ * contributors may be used to endorse or promote products derived
+ * from this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED "AS IS" AND ANY EXPRESS OR IMPLIED
+ * WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
+ * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NON-INFRINGEMENT
+ * ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS
+ * BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR
+ * BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
+ * WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
+ * OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN
+ * IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+*/
+
 #if !defined(AVB_INSIDE_LIBAVB_H) && !defined(AVB_COMPILATION)
 #error "Never include this file directly, include libavb.h instead."
 #endif
@@ -228,6 +256,16 @@ struct AvbOps {
                                        const char* partition,
                                        uint64_t* out_size_num_bytes);
 
+  /* Gets the size of a imgae with the name in |partition|
+   * (NUL-terminated UTF-8 string). Returns the value in
+   * |out_size_num_bytes|.
+   *
+   * Returns AVB_IO_RESULT_OK on success, otherwise an error code.
+   */
+  AvbIOResult (*get_actual_size_of_boot_image)(AvbOps* ops,
+                                       const char* PARTITION,
+                                       uint64_t* out_size_num_bytes);
+
   /* Like validate_vbmeta_public_key() but for when the flag
    * AVB_SLOT_VERIFY_FLAGS_NO_VBMETA_PARTITION is being used. The name of the
    * partition to get the public key for is passed in |partition_name|.
diff --git a/QcomModulePkg/Library/avb/libavb/avb_slot_verify.c b/QcomModulePkg/Library/avb/libavb/avb_slot_verify.c
index 7967f007bd..e29b60c190 100644..100755
--- a/QcomModulePkg/Library/avb/libavb/avb_slot_verify.c
+++ b/QcomModulePkg/Library/avb/libavb/avb_slot_verify.c
@@ -22,7 +22,7 @@
  * SOFTWARE.
  */
 
-/* Copyright (c) 2017-2020, The Linux Foundation. All rights reserved.
+/* Copyright (c) 2017-2021, The Linux Foundation. All rights reserved.
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are
@@ -59,6 +59,7 @@
 #include "avb_util.h"
 #include "avb_vbmeta_image.h"
 #include "avb_version.h"
+#include <Library/BaseLib.h>
 
 /* Maximum allow length (in bytes) of a partition name, including
  * ab_suffix.
@@ -169,7 +170,14 @@ static AvbSlotVerifyResult load_and_verify_hash_partition(
    */
   image_size = hash_desc.image_size;
   if (allow_verification_error) {
-    io_ret = ops->get_size_of_partition (ops, part_name, &image_size);
+    if (!AsciiStrnCmp (part_name, "boot", AsciiStrLen ("boot")) ||
+      !AsciiStrnCmp (part_name, "recovery", AsciiStrLen ("recovery")) ||
+      !AsciiStrnCmp (part_name, "vendor_boot", AsciiStrLen ("vendor_boot"))) {
+      io_ret = ops->get_actual_size_of_boot_image (ops, part_name, &image_size);
+    } else {
+      io_ret = ops->get_size_of_partition (ops, part_name, &image_size);
+    }
+
     if (io_ret == AVB_IO_RESULT_ERROR_OOM) {
       ret = AVB_SLOT_VERIFY_RESULT_ERROR_OOM;
       goto out;