QcomModulePkg: BootImagev3: Update DTB offsets

(cherry picked from commit 9ba062d261320af32e475fcba48dcb6eee94a51a)
https://source.codeaurora.org/quic/la/abl/tianocore/edk2/log/?h=LA.UM.9.12.r1-08000-SMxx50.0

For boot-image header version three, the DTB is packaged as
a part of vendor-boot image. Hence, update the image buffer,
DTB size and offset accordingly to get the device tree and
parse it.

Also, replace the variable DtbSize with ImageSize in
DTBImgCheckAndAppendDT() as it's a little misleading.
Technically, it's used to hold the size of the Image
so that the DT can be extracted properly from it.

Change-Id: I19cb04ce60f00968baebc7e12adadb8dbad68b04

Signed-off-by: Amit Pundir <amit.pundir@linaro.org>

2 files changed, 39 insertions, 18 deletions

diff --git a/QcomModulePkg/Library/BootLib/BootLinux.c b/QcomModulePkg/Library/BootLib/BootLinux.c
index aa261dcc41..c2c4f5e836 100644
--- a/QcomModulePkg/Library/BootLib/BootLinux.c
+++ b/QcomModulePkg/Library/BootLib/BootLinux.c
@@ -390,13 +390,15 @@ DTBImgCheckAndAppendDT (BootInfo *Info, BootParamlist *BootParamlistPtr)
   UINT32 HeaderVersion = 0;
   struct boot_img_hdr_v1 *BootImgHdrV1;
   struct boot_img_hdr_v2 *BootImgHdrV2;
+  vendor_boot_img_hdr_v3 *VendorBootImgHdrV3;
   UINT32 NumHeaderPages;
   UINT32 NumKernelPages;
   UINT32 NumSecondPages;
   UINT32 NumRamdiskPages;
+  UINT32 NumVendorRamdiskPages;
   UINT32 NumRecoveryDtboPages;
   VOID* ImageBuffer = NULL;
-  UINT32 DtbSize = 0;
+  UINT32 ImageSize = 0;
 
   if (Info == NULL ||
       BootParamlistPtr == NULL) {
@@ -407,7 +409,7 @@ DTBImgCheckAndAppendDT (BootInfo *Info, BootParamlist *BootParamlistPtr)
   ImageBuffer = BootParamlistPtr->ImageBuffer +
                         BootParamlistPtr->PageSize +
                         BootParamlistPtr->PatchedKernelHdrSize;
-  DtbSize = BootParamlistPtr->KernelSize;
+  ImageSize = BootParamlistPtr->KernelSize;
   HeaderVersion = Info->HeaderVersion;
 
   if (HeaderVersion > BOOT_HEADER_VERSION_ONE) {
@@ -429,26 +431,40 @@ DTBImgCheckAndAppendDT (BootInfo *Info, BootParamlist *BootParamlistPtr)
         NumSecondPages =
                 GetNumberOfPages (BootParamlistPtr->SecondSize,
                         BootParamlistPtr->PageSize);
-        NumRecoveryDtboPages =
-                GetNumberOfPages (BootImgHdrV1->recovery_dtbo_size,
-                        BootParamlistPtr->PageSize);
-        BootParamlistPtr->DtbOffset =
-                BootParamlistPtr->PageSize *
-                        (NumHeaderPages + NumKernelPages + NumRamdiskPages
-                                + NumSecondPages + NumRecoveryDtboPages);
-        DtbSize = BootImgHdrV2->dtb_size + BootParamlistPtr->DtbOffset;
-        ImageBuffer = BootParamlistPtr->ImageBuffer;
-  }
 
+       if (HeaderVersion  == BOOT_HEADER_VERSION_TWO) {
+          NumRecoveryDtboPages =
+                           GetNumberOfPages (BootImgHdrV1->recovery_dtbo_size,
+                           BootParamlistPtr->PageSize);
+          BootParamlistPtr->DtbOffset = BootParamlistPtr->PageSize *
+                           (NumHeaderPages + NumKernelPages + NumRamdiskPages +
+                            NumSecondPages + NumRecoveryDtboPages);
+          ImageSize = BootImgHdrV2->dtb_size + BootParamlistPtr->DtbOffset;
+          ImageBuffer = BootParamlistPtr->ImageBuffer;
+        } else {
+          VendorBootImgHdrV3 = BootParamlistPtr->VendorImageBuffer;
+
+          NumVendorRamdiskPages = GetNumberOfPages (
+                                           BootParamlistPtr->VendorRamdiskSize,
+                                           BootParamlistPtr->PageSize);
+          BootParamlistPtr->DtbOffset = BootParamlistPtr->PageSize *
+                           (NumHeaderPages + NumVendorRamdiskPages);
+          ImageSize = VendorBootImgHdrV3->dtb_size +
+                      BootParamlistPtr->DtbOffset;
+
+          // DTB is a part of vendor-boot image
+          ImageBuffer = BootParamlistPtr->VendorImageBuffer;
+        }
+  }
   DtboImgInvalid = LoadAndValidateDtboImg (Info, BootParamlistPtr);
   if (!DtboImgInvalid) {
     // appended device tree
     Dtb = DeviceTreeAppended (ImageBuffer,
-                             DtbSize,
+                             ImageSize,
                              BootParamlistPtr->DtbOffset,
                              (VOID *)BootParamlistPtr->DeviceTreeLoadAddr);
     if (!Dtb) {
-      if (BootParamlistPtr->DtbOffset >= DtbSize) {
+      if (BootParamlistPtr->DtbOffset >= ImageSize) {
         DEBUG ((EFI_D_ERROR, "Dtb offset goes beyond the image size\n"));
         return EFI_BAD_BUFFER_SIZE;
       }
@@ -457,7 +473,7 @@ DTBImgCheckAndAppendDT (BootInfo *Info, BootParamlist *BootParamlistPtr)
                      BootParamlistPtr->DtbOffset);
 
       if (!fdt_check_header (SingleDtHdr)) {
-        if ((DtbSize - BootParamlistPtr->DtbOffset) <
+        if ((ImageSize - BootParamlistPtr->DtbOffset) <
             fdt_totalsize (SingleDtHdr)) {
           DEBUG ((EFI_D_ERROR, "Dtb offset goes beyond the image size\n"));
           return EFI_BAD_BUFFER_SIZE;
@@ -507,7 +523,7 @@ DTBImgCheckAndAppendDT (BootInfo *Info, BootParamlist *BootParamlistPtr)
   } else {
     /*It is the case of DTB overlay Get the Soc specific dtb */
     SocDtb = GetSocDtb (ImageBuffer,
-         DtbSize,
+         ImageSize,
          BootParamlistPtr->DtbOffset,
          (VOID *)BootParamlistPtr->DeviceTreeLoadAddr);
 
@@ -1069,9 +1085,12 @@ UpdateBootParamsSizeAndCmdLine (BootInfo *Info, BootParamlist *BootParamlistPtr)
     return Status;
   }
 
+  BootParamlistPtr->VendorImageBuffer = VendorBootImgHdrV3;
+  BootParamlistPtr->VendorImageSize = VendorBootImgSize;
   BootParamlistPtr->KernelSize = BootImgHdrV3->kernel_size;
-  BootParamlistPtr->RamdiskSize = BootImgHdrV3->ramdisk_size +
-                                  VendorBootImgHdrV3->vendor_ramdisk_size;
+  BootParamlistPtr->RamdiskSize = BootImgHdrV3->ramdisk_size;
+  BootParamlistPtr->VendorRamdiskSize =
+                    VendorBootImgHdrV3->vendor_ramdisk_size;
   BootParamlistPtr->PageSize = VendorBootImgHdrV3->page_size;
   BootParamlistPtr->SecondSize = 0;
 
diff --git a/QcomModulePkg/Library/BootLib/LocateDeviceTree.c b/QcomModulePkg/Library/BootLib/LocateDeviceTree.c
index 13701e81ba..3d4db2ccb0 100644
--- a/QcomModulePkg/Library/BootLib/LocateDeviceTree.c
+++ b/QcomModulePkg/Library/BootLib/LocateDeviceTree.c
@@ -914,6 +914,8 @@ cleanup:
 /*
  * For Header Version 2, the arguments Kernel and KernelSize will be
  * the entire bootimage and the bootimage size.
+ * For Header Version 3, Kernel holds the base of the vendor-boot
+ * image and KernelSize holds its size.
  */
 VOID *
 GetSocDtb (VOID *Kernel, UINT32 KernelSize, UINT32 DtbOffset, VOID *DtbLoadAddr)