QcomModulePkg: BootImagev3: Add support for fastboot boot

Originally, in the fastboot boot path (CmdBoot()), we perform the
validation for boot-image first and then load the other images.
However, since the validation of boot header (CheckImageHeader())
is now unified for validating also the vendor-boot header, boot
header validation cannot be performed unless vendor-boot is also
loaded. Hence, remove the check in the CmdBoot() path and let it
get natually checked while loading other images through
LoadImageAndAuth().

Change-Id: I009db80656dc8365e42c50821f8e3b978ba1211d

2 files changed, 73 insertions, 54 deletions

diff --git a/QcomModulePkg/Library/FastbootLib/FastbootCmds.c b/QcomModulePkg/Library/FastbootLib/FastbootCmds.c
index 59f293ed4b..06312b9b65 100644
--- a/QcomModulePkg/Library/FastbootLib/FastbootCmds.c
+++ b/QcomModulePkg/Library/FastbootLib/FastbootCmds.c
@@ -2453,9 +2453,9 @@ STATIC VOID
 CmdBoot (CONST CHAR8 *Arg, VOID *Data, UINT32 Size)
 {
   boot_img_hdr *hdr = Data;
+  boot_img_hdr_v3 *HdrV3 = Data;
   EFI_STATUS Status = EFI_SUCCESS;
   UINT32 ImageSizeActual = 0;
-  UINT32 ImageHdrSize = BOOT_IMG_MAX_PAGE_SIZE;
   UINT32 PageSize = 0;
   UINT32 SigActual = SIGACTUAL;
   CHAR8 Resp[MAX_RSP_SIZE];
@@ -2487,29 +2487,17 @@ CmdBoot (CONST CHAR8 *Arg, VOID *Data, UINT32 Size)
     return;
   }
 
-  hdr->cmdline[BOOT_ARGS_SIZE - 1] = '\0';
-  SetBootDevImage ();
-
-  /* TODO: Fix the arguments */
-  Status = CheckImageHeader (Data, ImageHdrSize, NULL, 0, &ImageSizeActual,
-                             &PageSize, FALSE);
-  if (Status != EFI_SUCCESS) {
-    AsciiSPrint (Resp, sizeof (Resp), "Invalid Boot image Header: %r", Status);
-    FastbootFail (Resp);
-    goto out;
+  if (hdr->header_version <= BOOT_HEADER_VERSION_TWO) {
+    hdr->cmdline[BOOT_ARGS_SIZE - 1] = '\0';
+  } else {
+    HdrV3->cmdline[BOOT_ARGS_SIZE + BOOT_EXTRA_ARGS_SIZE - 1] = '\0';
   }
 
-  if (ImageSizeActual > Size) {
-    FastbootFail ("BootImage is Incomplete");
-    goto out;
-  }
-  if ((MaxDownLoadSize - (ImageSizeActual - SigActual)) < PageSize) {
-    FastbootFail ("BootImage: Size os greater than boot image buffer can hold");
-    goto out;
-  }
+  SetBootDevImage ();
 
   Info.Images[0].ImageBuffer = Data;
-  Info.Images[0].ImageSize = ImageSizeActual;
+  /* The actual image size will be updated in LoadImageAndAuth */
+  Info.Images[0].ImageSize = Size;
   Info.Images[0].Name = "boot";
   Info.NumLoadedImages = 1;
   Info.MultiSlotBoot = PartitionHasMultiSlot (L"boot");
@@ -2521,6 +2509,7 @@ CmdBoot (CONST CHAR8 *Arg, VOID *Data, UINT32 Size)
       goto out;
     }
   }
+
   Status = LoadImageAndAuth (&Info);
   if (Status != EFI_SUCCESS) {
     AsciiSPrint (Resp, sizeof (Resp),
@@ -2529,6 +2518,17 @@ CmdBoot (CONST CHAR8 *Arg, VOID *Data, UINT32 Size)
     goto out;
   }
 
+  ImageSizeActual = Info.Images[0].ImageSize;
+
+  if (ImageSizeActual > Size) {
+    FastbootFail ("BootImage is Incomplete");
+    goto out;
+  }
+  if ((MaxDownLoadSize - (ImageSizeActual - SigActual)) < PageSize) {
+    FastbootFail ("BootImage: Size is greater than boot image buffer can hold");
+    goto out;
+  }
+
   /* Exit keys' detection firstly */
   ExitMenuKeysDetection ();
 
diff --git a/QcomModulePkg/Library/avb/VerifiedBoot.c b/QcomModulePkg/Library/avb/VerifiedBoot.c
index 2f112d8e32..cad4575beb 100644
--- a/QcomModulePkg/Library/avb/VerifiedBoot.c
+++ b/QcomModulePkg/Library/avb/VerifiedBoot.c
@@ -349,7 +349,7 @@ LoadVendorBootImageHeader (BootInfo *Info,
 }
 
 STATIC EFI_STATUS
-LoadBootImageNoAuth (BootInfo *Info, UINT32 *PageSize)
+LoadBootImageNoAuth (BootInfo *Info, UINT32 *PageSize, BOOLEAN *FastbootPath)
 {
   EFI_STATUS Status = EFI_SUCCESS;
   VOID *ImageHdrBuffer = NULL;
@@ -357,25 +357,33 @@ LoadBootImageNoAuth (BootInfo *Info, UINT32 *PageSize)
   UINT32 ImageSizeActual = 0;
   VOID *VendorImageHdrBuffer = NULL;
   UINT32 VendorImageHdrSize = 0;
+  BOOLEAN BootIntoRecovery = FALSE;
+  BOOLEAN BootImageLoaded;
 
   /** The Images[0].ImageBuffer would have been loaded with the boot image
    *  already if we are coming from fastboot boot path. Ignore loading it
    *  again.
    **/
-  if (Info->Images[0].ImageBuffer != NULL &&
-     Info->Images[0].ImageSize > 0) {
-    return EFI_SUCCESS;
-  }
+  BootImageLoaded = (Info->Images[0].ImageBuffer != NULL) &&
+                    (Info->Images[0].ImageSize > 0);
+  *FastbootPath = BootImageLoaded;
 
-  Status = LoadImageHeader (Info->Pname, &ImageHdrBuffer, &ImageHdrSize);
-  if (Status != EFI_SUCCESS ||
-      ImageHdrBuffer ==  NULL) {
-    DEBUG ((EFI_D_ERROR, "ERROR: Failed to load image header: %r\n", Status));
-    return Status;
-  } else if (ImageHdrSize < sizeof (boot_img_hdr)) {
-    DEBUG ((EFI_D_ERROR,
-            "ERROR: Invalid image header size: %u\n", ImageHdrSize));
-    return EFI_BAD_BUFFER_SIZE;
+  if (BootImageLoaded) {
+    ImageHdrBuffer = Info->Images[0].ImageBuffer;
+    ImageHdrSize = BOOT_IMG_MAX_PAGE_SIZE;
+  } else {
+    Status = LoadImageHeader (Info->Pname, &ImageHdrBuffer, &ImageHdrSize);
+    if (Status != EFI_SUCCESS ||
+        ImageHdrBuffer ==  NULL) {
+      DEBUG ((EFI_D_ERROR, "ERROR: Failed to load image header: %r\n", Status));
+      return Status;
+    } else if (ImageHdrSize < sizeof (boot_img_hdr)) {
+      DEBUG ((EFI_D_ERROR,
+              "ERROR: Invalid image header size: %u\n", ImageHdrSize));
+      return EFI_BAD_BUFFER_SIZE;
+    }
+
+    BootIntoRecovery = Info->BootIntoRecovery;
   }
 
   Info->HeaderVersion = ((boot_img_hdr *)(ImageHdrBuffer))->header_version;
@@ -398,29 +406,31 @@ LoadBootImageNoAuth (BootInfo *Info, UINT32 *PageSize)
    */
   Status = CheckImageHeader (ImageHdrBuffer, ImageHdrSize,
                              VendorImageHdrBuffer, VendorImageHdrSize,
-                             &ImageSizeActual, PageSize,
-                             Info->BootIntoRecovery);
+                             &ImageSizeActual, PageSize, BootIntoRecovery);
   if (Status != EFI_SUCCESS) {
     DEBUG ((EFI_D_ERROR, "Invalid boot image header:%r\n", Status));
     goto Err;
   }
 
-  Status = LoadImage (Info->Pname, (VOID **)&(Info->Images[0].ImageBuffer),
-                      ImageSizeActual, *PageSize);
-  if (Status != EFI_SUCCESS) {
-    DEBUG ((EFI_D_ERROR, "ERROR: Failed to load image from partition: %r\n",
-            Status));
-    goto Err;
+  if (!BootImageLoaded) {
+    Status = LoadImage (Info->Pname, (VOID **)&(Info->Images[0].ImageBuffer),
+                        ImageSizeActual, *PageSize);
+    if (Status != EFI_SUCCESS) {
+      DEBUG ((EFI_D_ERROR, "ERROR: Failed to load image from partition: %r\n",
+              Status));
+      goto Err;
+    }
+
+    Info->Images[0].Name = AllocateZeroPool (StrLen (Info->Pname) + 1);
+    if (!Info->Images[0].Name) {
+      Status = EFI_OUT_OF_RESOURCES;
+      goto ErrImg;
+    }
+    UnicodeStrToAsciiStr (Info->Pname, Info->Images[0].Name);
+    Info->NumLoadedImages = 1;
   }
 
   Info->Images[0].ImageSize = ImageSizeActual;
-  Info->Images[0].Name = AllocateZeroPool (StrLen (Info->Pname) + 1);
-  if (!Info->Images[0].Name) {
-    Status = EFI_OUT_OF_RESOURCES;
-    goto ErrImg;
-  }
-  UnicodeStrToAsciiStr (Info->Pname, Info->Images[0].Name);
-  Info->NumLoadedImages = 1;
 
   if (Info->HeaderVersion >= BOOT_HEADER_VERSION_THREE) {
     Status = NoAVBLoadVendorBootImage (Info);
@@ -434,11 +444,13 @@ LoadBootImageNoAuth (BootInfo *Info, UINT32 *PageSize)
   return EFI_SUCCESS;
 
 ErrImgName:
-  if (Info->Images[0].Name) {
+  if (!BootImageLoaded &&
+      Info->Images[0].Name) {
     FreePool (Info->Images[0].Name);
   }
 ErrImg:
-  if (Info->Images[0].ImageBuffer) {
+  if (!BootImageLoaded &&
+      Info->Images[0].ImageBuffer) {
     UINT32 ImageSize =
       ADD_OF (ROUND_TO_PAGE (ImageSizeActual, (*PageSize - 1)), *PageSize);
     FreePages (Info->Images[0].ImageBuffer,
@@ -450,7 +462,8 @@ Err:
                ALIGN_PAGES (BOOT_IMG_MAX_PAGE_SIZE, ALIGNMENT_MASK_4KB));
   }
 ErrV3:
-  if (ImageHdrBuffer) {
+  if (!BootImageLoaded &&
+      ImageHdrBuffer) {
     FreePages (ImageHdrBuffer,
                ALIGN_PAGES (BOOT_IMG_MAX_PAGE_SIZE, ALIGNMENT_MASK_4KB));
   }
@@ -465,8 +478,9 @@ LoadImageNoAuth (BootInfo *Info)
   CHAR16 Pname[MAX_GPT_NAME_SIZE];
   UINTN *ImgIdx = &Info->NumLoadedImages;
   UINT32 PageSize = 0;
+  BOOLEAN FastbootPath;
 
-  Status = LoadBootImageNoAuth (Info, &PageSize);
+  Status = LoadBootImageNoAuth (Info, &PageSize, &FastbootPath);
   if (Status != EFI_SUCCESS) {
     return Status;
   }
@@ -537,8 +551,12 @@ Err:
   }
 
   /* Images[0] needs to be freed in a special way as it was allocated
-   * using AllocPages().
+   * using AllocPages(). Although, ignore if we are coming from a
+   * fastboot boot path.
    */
+  if (FastbootPath)
+    goto err_out;
+
   if (Info->Images[0].ImageBuffer) {
     UINT32 ImageSize =
       ADD_OF (ROUND_TO_PAGE (Info->Images[0].ImageSize,
@@ -552,6 +570,7 @@ Err:
     FreePool (Info->Images[0].Name);
   }
 
+err_out:
   return Status;
 }