QcomModulePkg: Add support for AVB2.0 when A/B feature is disabled.

Add support in AVB when multislot is disabled.
Remove explicit checks, support handling alternate
slot only when multislot is enabled.
Support loading of recovery image when multislot is
disabled and booting into recovery mode.

Acked-by: Barani Muthukumaran <bmuthuku@qti.qualcomm.com>
Change-Id: I9aa07ec63cd5767c79555a7b93f778ca13f734d6
diff --git a/QcomModulePkg/Library/avb/AvbLib.inf b/QcomModulePkg/Library/avb/AvbLib.inf
index b94fdd5..d0a69ac 100644
--- a/QcomModulePkg/Library/avb/AvbLib.inf
+++ b/QcomModulePkg/Library/avb/AvbLib.inf
@@ -85,7 +85,10 @@
 
 
 [Guids]
-
+	gEfiVbmetaPartitionGuid
+	gEfiBootImgPartitionGuid
+	gEfiDtboPartitionGuid
+	gEfiRecoveryImgPartitionGuid
 
 [Protocols]
 	gQcomQseecomProtocolGuid
diff --git a/QcomModulePkg/Library/avb/VerifiedBoot.c b/QcomModulePkg/Library/avb/VerifiedBoot.c
index 8133dab..053ea3f 100644
--- a/QcomModulePkg/Library/avb/VerifiedBoot.c
+++ b/QcomModulePkg/Library/avb/VerifiedBoot.c
@@ -141,12 +141,6 @@
 {
   EFI_STATUS Status = EFI_SUCCESS;
 
-  /* AVB 2.0 requires multi slot */
-  if (GetAVBVersion () >= AVB_2 && !Info->MultiSlotBoot) {
-    DEBUG ((EFI_D_ERROR, "AVB requires multislot support!\n"));
-    return EFI_LOAD_ERROR;
-  }
-
   Info->BootState = RED;
 
   Status = gBS->LocateProtocol (&gEfiQcomVerifiedBootProtocolGuid, NULL,
@@ -375,7 +369,8 @@
   CHAR8 PnameAscii[MAX_GPT_NAME_SIZE] = {0};
   CHAR8 *SlotSuffix = NULL;
   BOOLEAN AllowVerificationError = IsUnlocked ();
-  CONST CHAR8 *RequestedPartitionAll[] = {"boot", "dtbo", NULL};
+  CONST CHAR8 *RequestedPartitionMission[] = {"boot", "dtbo", NULL};
+  CONST CHAR8 *RequestedPartitionRecovery[] = {"recovery", "dtbo", NULL};
   CONST CHAR8 *CONST *RequestedPartition = NULL;
   UINTN NumRequestedPartition = 0;
   UINT32 ImageHdrSize = 0;
@@ -409,23 +404,41 @@
     goto out;
   }
 
-  UnicodeStrToAsciiStr (Info->Pname, PnameAscii);
-  if ((MAX_SLOT_SUFFIX_SZ + 1) > AsciiStrLen (PnameAscii)) {
-    DEBUG ((EFI_D_ERROR, "ERROR: Can not determine slot suffix\n"));
-    Status = EFI_INVALID_PARAMETER;
-    goto out;
+  if (Info->MultiSlotBoot) {
+    UnicodeStrToAsciiStr (Info->Pname, PnameAscii);
+    if ((MAX_SLOT_SUFFIX_SZ + 1) > AsciiStrLen (PnameAscii)) {
+      DEBUG ((EFI_D_ERROR, "ERROR: Can not determine slot suffix\n"));
+      Status = EFI_INVALID_PARAMETER;
+      goto out;
+    }
+    SlotSuffix = &PnameAscii[AsciiStrLen (PnameAscii) - MAX_SLOT_SUFFIX_SZ + 1];
+  } else {
+     SlotSuffix = "\0";
   }
-  SlotSuffix = &PnameAscii[AsciiStrLen (PnameAscii) - MAX_SLOT_SUFFIX_SZ + 1];
 
   DEBUG ((EFI_D_VERBOSE, "Slot: %a, allow verification error: %a\n", SlotSuffix,
           BooleanString[AllowVerificationError].name));
 
-  RequestedPartition = RequestedPartitionAll;
-  NumRequestedPartition = ARRAY_SIZE (RequestedPartitionAll) - 1;
+  if ((!Info->MultiSlotBoot) &&
+           Info->BootIntoRecovery) {
+     RequestedPartition = RequestedPartitionRecovery;
+     NumRequestedPartition = ARRAY_SIZE (RequestedPartitionRecovery) - 1;
+     if (Info->NumLoadedImages) {
+       /* fastboot boot option, skip Index 0, as boot image already
+        * loaded */
+       RequestedPartition = &RequestedPartitionRecovery[1];
+     }
+  } else {
+     RequestedPartition = RequestedPartitionMission;
+     NumRequestedPartition = ARRAY_SIZE (RequestedPartitionMission) - 1;
+     if (Info->NumLoadedImages) {
+       /* fastboot boot option, skip Index 0, as boot image already
+        * loaded */
+       RequestedPartition = &RequestedPartitionMission[1];
+     }
+  }
+
   if (Info->NumLoadedImages) {
-    /* fastboot boot option, skip Index 0, as boot image already
-     * loaded */
-    RequestedPartition = &RequestedPartitionAll[1];
     NumRequestedPartition--;
   }
 
@@ -485,10 +498,10 @@
     }
   }
 
-  if (Info->NumLoadedImages < (ARRAY_SIZE (RequestedPartitionAll) - 1)) {
+  if (Info->NumLoadedImages < NumRequestedPartition) {
     DEBUG ((EFI_D_ERROR, "ERROR: AvbSlotVerify slot data error: num of "
                          "loaded partitions %d, requested %d\n",
-            Info->NumLoadedImages, ARRAY_SIZE (RequestedPartitionAll) - 1));
+            Info->NumLoadedImages, NumRequestedPartition));
     Status = EFI_LOAD_ERROR;
     goto out;
   }
@@ -506,10 +519,13 @@
   Info->VBData = (VOID *)VBData;
 
   GetPageSize (&ImageHdrSize);
-  GUARD_OUT (GetImage (Info, &ImageBuffer, &ImageSize, "boot"));
+  GUARD_OUT (GetImage (Info, &ImageBuffer, &ImageSize,
+                    ((!Info->MultiSlotBoot) &&
+                     Info->BootIntoRecovery) ?
+                     "recovery" : "boot"));
 
-  Status =
-      CheckImageHeader (ImageBuffer, ImageHdrSize, &ImageSizeActual, &PageSize);
+  Status = CheckImageHeader (ImageBuffer, ImageHdrSize,
+        &ImageSizeActual, &PageSize);
   if (Status != EFI_SUCCESS) {
     DEBUG ((EFI_D_ERROR, "Invalid boot image header:%r\n", Status));
     goto out;
@@ -566,10 +582,16 @@
       avb_free (VBData);
     }
     Info->BootState = RED;
-    HandleActiveSlotUnbootable ();
-    /* HandleActiveSlotUnbootable should have swapped slots and
-    * reboot the device. If no bootable slot found, enter fastboot */
-    DEBUG ((EFI_D_WARN, "No bootable slots found enter fastboot mode\n"));
+    if (Info->MultiSlotBoot) {
+      HandleActiveSlotUnbootable ();
+      /* HandleActiveSlotUnbootable should have swapped slots and
+       * reboot the device. If no bootable slot found, enter fastboot
+       */
+      DEBUG ((EFI_D_WARN, "No bootable slots found enter fastboot mode\n"));
+    } else {
+       DEBUG ((EFI_D_WARN,
+           "Non Multi-slot: Unbootable entering fastboot mode\n"));
+    }
   }
 
   DEBUG ((EFI_D_ERROR, "VB2: boot state: %a(%d)\n", VbSn[Info->BootState].name,
diff --git a/QcomModulePkg/Library/avb/libavb/avb_ops.c b/QcomModulePkg/Library/avb/libavb/avb_ops.c
index 6fd7b31..1c18cd4 100644
--- a/QcomModulePkg/Library/avb/libavb/avb_ops.c
+++ b/QcomModulePkg/Library/avb/libavb/avb_ops.c
@@ -22,7 +22,7 @@
  * SOFTWARE.
  */
 
-/* Copyright (c) 2017, The Linux Foundation. All rights reserved.
+/* Copyright (c) 2017-2018, 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
@@ -103,8 +103,8 @@
 	}
 
 	if (MaxHandles != 1) {
-		// Unable to deterministically load from single
-		// partition
+                /* Unable to deterministically load from single
+                partition */
 		DEBUG((EFI_D_ERROR, "GetHandleInfo: More than one result!\n"));
 		return AVB_IO_RESULT_ERROR_IO;
 	}
@@ -112,6 +112,19 @@
 	return AVB_IO_RESULT_OK;
 }
 
+typedef struct {
+        CONST CHAR8 *Name;
+        EFI_GUID *Guid;
+}AvbPartitionDetails;
+
+static AvbPartitionDetails SupportedPartitions[] =
+{
+  { "vbmeta", &gEfiVbmetaPartitionGuid },
+  { "boot", &gEfiBootImgPartitionGuid },
+  { "dtbo", &gEfiDtboPartitionGuid },
+  { "recovery", &gEfiRecoveryImgPartitionGuid} ,
+};
+
 AvbIOResult AvbReadFromPartition(AvbOps *Ops, const char *Partition, int64_t ReadOffset,
                      size_t NumBytes, void *Buffer, size_t *OutNumRead)
 {
@@ -127,6 +140,12 @@
 	UINT32 LastBlock = 0;
 	UINT32 FullBlock = 0;
 	UINTN StartPageReadSize = 0;
+        UINT32 BlkIOAttrib = 0;
+        PartiSelectFilter HandleFilter;
+        UINT32 MaxHandles = 0;
+        AvbPartitionDetails *List = SupportedPartitions;
+        UINT32 Count = ARRAY_SIZE (SupportedPartitions);
+        EFI_GUID *PType = NULL;
 
 	if (Partition == NULL || Buffer == NULL || OutNumRead == NULL || NumBytes <= 0) {
 		DEBUG((EFI_D_ERROR, "bad input paramaters\n"));
@@ -135,12 +154,56 @@
 	}
 	*OutNumRead = 0;
 
-	Result = GetHandleInfo(Partition, InfoList);
-	if (Result != AVB_IO_RESULT_OK) {
-		DEBUG((EFI_D_ERROR,
-		       "AvbGetSizeOfPartition: GetHandleInfo failed"));
-		goto out;
-	}
+        for (size_t Index = 0; Index < Count; Index++) {
+                if (!AsciiStrCmp (List[Index].Name, Partition)) {
+                             DEBUG ((EFI_D_INFO,
+                                  "Partition found: %a\n", Partition));
+                             PType = List[Index].Guid;
+                 }
+        }
+
+        if (PType) {
+                BlkIOAttrib = BLK_IO_SEL_PARTITIONED_GPT;
+                BlkIOAttrib |= BLK_IO_SEL_MEDIA_TYPE_NON_REMOVABLE;
+                BlkIOAttrib |= BLK_IO_SEL_MATCH_PARTITION_TYPE_GUID;
+
+                HandleFilter.RootDeviceType = NULL;
+                HandleFilter.PartitionType = PType;
+                HandleFilter.VolumeName = NULL;
+
+                MaxHandles = ARRAY_SIZE (InfoList);
+
+                Status = GetBlkIOHandles (BlkIOAttrib, &HandleFilter,
+                           InfoList, &MaxHandles);
+                if (Status == EFI_SUCCESS) {
+                        if (MaxHandles == 0) {
+                            DEBUG ((EFI_D_INFO,
+                            "Partition Not found: %s\n", Partition));
+                            Result = AVB_IO_RESULT_ERROR_NO_SUCH_PARTITION;
+                            goto out;
+                }
+
+                if (MaxHandles != 1) {
+                /* Unable to deterministically load from single partition */
+                             DEBUG ((EFI_D_INFO,
+                             "multiple partitions found: %s\n", Partition));
+                              Result = AVB_IO_RESULT_ERROR_IO;
+                              goto out;
+                     }
+                 } else {
+                      DEBUG ((EFI_D_INFO,
+                       "GetBlkIOHandles failed with error: %d\n", Status));
+                        Result = AVB_IO_RESULT_ERROR_IO;
+                        goto out;
+               }
+          } else {
+                Result = GetHandleInfo (Partition, InfoList);
+                if (Result != AVB_IO_RESULT_OK) {
+                      DEBUG ((EFI_D_ERROR,
+                       "AvbGetSizeOfPartition: GetHandleInfo failed"));
+                      goto out;
+           }
+       }
 
 	BlockIo = InfoList[0].BlkIo;
 	PartitionSize = (BlockIo->Media->LastBlock + 1) * BlockIo->Media->BlockSize;
diff --git a/QcomModulePkg/QcomModulePkg.dec b/QcomModulePkg/QcomModulePkg.dec
index a67cf46..e8fddfa 100644
--- a/QcomModulePkg/QcomModulePkg.dec
+++ b/QcomModulePkg/QcomModulePkg.dec
@@ -81,6 +81,8 @@
   gEfiBootImgPartitionGuid =            { 0x20117f86, 0xe985, 0x4357, { 0xb9, 0xee, 0x37, 0x4b, 0xc1, 0xd8, 0x48, 0x7d } }
   gEfiRecoveryImgPartitionGuid =        { 0x9D72D4E4, 0x9958, 0x42DA, { 0xAC, 0x26, 0xBE, 0xA7, 0xA9, 0x0B, 0x04, 0x34 } }
   gEfiSystemPartitionGuid =             { 0x97D7B011, 0x54DA, 0x4835, { 0xB3, 0xC4, 0x91, 0x7A, 0xD6, 0xE7, 0x3D, 0x74 } }
+  gEfiVbmetaPartitionGuid =             { 0x4b7a15d6, 0x322c, 0x42ac, { 0x81, 0x10, 0x88, 0xb7, 0xda, 0x0c, 0x5d, 0x77 } }
+  gEfiDtboPartitionGuid =               { 0x24d0d418, 0xd31d, 0x4d8d, { 0xac, 0x2c, 0x4d, 0x43, 0x05, 0x18, 0x84, 0x50 } }
   gEfiDipPartitionGuid =                { 0x4114B077, 0x005D, 0x4E12, { 0xAC, 0x8C, 0xB4, 0x93, 0xBD, 0xA6, 0x84, 0xFB } }
   gEfiMdtpPartitionGuid =               { 0x3878408A, 0xE263, 0x4B67, { 0xB8, 0x78, 0x63, 0x40, 0xB3, 0x5B, 0x11, 0xE3 } }
   gBlockIoRefreshGuid =                 { 0xb1eb3d10, 0x9d67, 0x40ca, { 0x95, 0x59, 0xf1, 0x48, 0x8b, 0x1b, 0x2d, 0xdb } }