]> git.proxmox.com Git - mirror_edk2.git/blobdiff - MdeModulePkg/Library/DxeCapsuleLibFmp/DxeCapsuleLib.c
MdeModulePkg/CapsuleLib: Fix runtime issue
[mirror_edk2.git] / MdeModulePkg / Library / DxeCapsuleLibFmp / DxeCapsuleLib.c
index af08886d02b9baa87970254acdaf91c2f3dff0d7..9ed0be3b653c9d59fde48ede0184750d8a69d29c 100644 (file)
@@ -10,7 +10,7 @@
   ValidateFmpCapsule(), DisplayCapsuleImage(), ConvertBmpToGopBlt() will\r
   receive untrusted input and do basic validation.\r
 \r
-  Copyright (c) 2016, Intel Corporation. All rights reserved.<BR>\r
+  Copyright (c) 2016 - 2017, Intel Corporation. All rights reserved.<BR>\r
   This program and the accompanying materials\r
   are licensed and made available under the terms and conditions of the BSD License\r
   which accompanies this distribution.  The full text of the license may be found at\r
 #include <Protocol/FirmwareManagement.h>\r
 #include <Protocol/DevicePath.h>\r
 \r
-BOOLEAN            mAreAllImagesProcessed;\r
-\r
-EFI_SYSTEM_RESOURCE_TABLE *mEsrtTable              = NULL;\r
-BOOLEAN                   mIsVirtualAddrConverted  = FALSE;\r
-BOOLEAN                   mDxeCapsuleLibEndOfDxe   = FALSE;\r
+BOOLEAN                   mDxeCapsuleLibEndOfDxe       = FALSE;\r
+EFI_EVENT                 mDxeCapsuleLibEndOfDxeEvent  = NULL;\r
 \r
 /**\r
   Initialize capsule related variables.\r
@@ -61,23 +58,6 @@ InitCapsuleVariable (
   VOID\r
   );\r
 \r
-/**\r
-  Check if this FMP capsule is processed.\r
-\r
-  @param[in] CapsuleHeader  The capsule image header\r
-  @param[in] PayloadIndex   FMP payload index\r
-  @param[in] ImageHeader    FMP image header\r
-\r
-  @retval TRUE  This FMP capsule is processed.\r
-  @retval FALSE This FMP capsule is not processed.\r
-**/\r
-BOOLEAN\r
-IsFmpCapsuleProcessed (\r
-  IN EFI_CAPSULE_HEADER                            *CapsuleHeader,\r
-  IN UINTN                                         PayloadIndex,\r
-  IN EFI_FIRMWARE_MANAGEMENT_CAPSULE_IMAGE_HEADER  *ImageHeader\r
-  );\r
-\r
 /**\r
   Record capsule status variable.\r
 \r
@@ -100,6 +80,7 @@ RecordCapsuleStatusVariable (
   @param[in] CapsuleStatus  The capsule process stauts\r
   @param[in] PayloadIndex   FMP payload index\r
   @param[in] ImageHeader    FMP image header\r
+  @param[in] FmpDevicePath  DevicePath associated with the FMP producer\r
 \r
   @retval EFI_SUCCESS          The capsule status variable is recorded.\r
   @retval EFI_OUT_OF_RESOURCES No resource to record the capsule status variable.\r
@@ -109,7 +90,8 @@ RecordFmpCapsuleStatusVariable (
   IN EFI_CAPSULE_HEADER                            *CapsuleHeader,\r
   IN EFI_STATUS                                    CapsuleStatus,\r
   IN UINTN                                         PayloadIndex,\r
-  IN EFI_FIRMWARE_MANAGEMENT_CAPSULE_IMAGE_HEADER  *ImageHeader\r
+  IN EFI_FIRMWARE_MANAGEMENT_CAPSULE_IMAGE_HEADER  *ImageHeader,\r
+  IN EFI_DEVICE_PATH_PROTOCOL                      *FmpDevicePath OPTIONAL\r
   );\r
 \r
 /**\r
@@ -564,7 +546,7 @@ ConvertBmpToGopBlt (
 \r
     }\r
 \r
-    ImageIndex = (UINTN) (Image - ImageHeader);\r
+    ImageIndex = (UINTN) Image - (UINTN) ImageHeader;\r
     if ((ImageIndex % 4) != 0) {\r
       //\r
       // Bmp Image starts each row on a 32-bit boundary!\r
@@ -602,7 +584,7 @@ DisplayCapsuleImage (
   EFI_GRAPHICS_OUTPUT_PROTOCOL  *GraphicsOutput;\r
 \r
   ImagePayload = (DISPLAY_DISPLAY_PAYLOAD *)(CapsuleHeader + 1);\r
-  PayloadSize = (UINTN)(CapsuleHeader->CapsuleImageSize - sizeof(EFI_CAPSULE_HEADER));\r
+  PayloadSize = CapsuleHeader->CapsuleImageSize - sizeof(EFI_CAPSULE_HEADER);\r
 \r
   if (ImagePayload->Version != 1) {\r
     return EFI_UNSUPPORTED;\r
@@ -750,7 +732,7 @@ DumpFmpCapsule (
   for (Index = 0; Index < FmpCapsuleHeader->EmbeddedDriverCount; Index++) {\r
     DEBUG((DEBUG_VERBOSE, "  ItemOffsetList[%d]      - 0x%lx\n", Index, ItemOffsetList[Index]));\r
   }\r
-  for (; Index < (UINTN)(FmpCapsuleHeader->EmbeddedDriverCount + FmpCapsuleHeader->PayloadItemCount); Index++) {\r
+  for (; Index < (UINT32)FmpCapsuleHeader->EmbeddedDriverCount + FmpCapsuleHeader->PayloadItemCount; Index++) {\r
     DEBUG((DEBUG_VERBOSE, "  ItemOffsetList[%d]      - 0x%lx\n", Index, ItemOffsetList[Index]));\r
     ImageHeader = (EFI_FIRMWARE_MANAGEMENT_CAPSULE_IMAGE_HEADER *)((UINT8 *)FmpCapsuleHeader + ItemOffsetList[Index]);\r
 \r
@@ -766,6 +748,520 @@ DumpFmpCapsule (
   }\r
 }\r
 \r
+/**\r
+  Dump all FMP information.\r
+**/\r
+VOID\r
+DumpAllFmpInfo (\r
+  VOID\r
+  )\r
+{\r
+  EFI_STATUS                                    Status;\r
+  EFI_HANDLE                                    *HandleBuffer;\r
+  UINTN                                         NumberOfHandles;\r
+  EFI_FIRMWARE_MANAGEMENT_PROTOCOL              *Fmp;\r
+  UINTN                                         Index;\r
+  UINTN                                         ImageInfoSize;\r
+  EFI_FIRMWARE_IMAGE_DESCRIPTOR                 *FmpImageInfoBuf;\r
+  UINT32                                        FmpImageInfoDescriptorVer;\r
+  UINT8                                         FmpImageInfoCount;\r
+  UINTN                                         DescriptorSize;\r
+  UINT32                                        PackageVersion;\r
+  CHAR16                                        *PackageVersionName;\r
+\r
+  Status = gBS->LocateHandleBuffer (\r
+                  ByProtocol,\r
+                  &gEfiFirmwareManagementProtocolGuid,\r
+                  NULL,\r
+                  &NumberOfHandles,\r
+                  &HandleBuffer\r
+                  );\r
+  if (EFI_ERROR(Status)) {\r
+    return ;\r
+  }\r
+\r
+  for (Index = 0; Index < NumberOfHandles; Index++) {\r
+    Status = gBS->HandleProtocol(\r
+                    HandleBuffer[Index],\r
+                    &gEfiFirmwareManagementProtocolGuid,\r
+                    (VOID **)&Fmp\r
+                    );\r
+    if (EFI_ERROR(Status)) {\r
+      continue;\r
+    }\r
+\r
+    ImageInfoSize = 0;\r
+    Status = Fmp->GetImageInfo (\r
+                    Fmp,\r
+                    &ImageInfoSize,\r
+                    NULL,\r
+                    NULL,\r
+                    NULL,\r
+                    NULL,\r
+                    NULL,\r
+                    NULL\r
+                    );\r
+    if (Status != EFI_BUFFER_TOO_SMALL) {\r
+      continue;\r
+    }\r
+\r
+    FmpImageInfoBuf = AllocateZeroPool (ImageInfoSize);\r
+    if (FmpImageInfoBuf == NULL) {\r
+      continue;\r
+    }\r
+\r
+    PackageVersionName = NULL;\r
+    Status = Fmp->GetImageInfo (\r
+                    Fmp,\r
+                    &ImageInfoSize,               // ImageInfoSize\r
+                    FmpImageInfoBuf,              // ImageInfo\r
+                    &FmpImageInfoDescriptorVer,   // DescriptorVersion\r
+                    &FmpImageInfoCount,           // DescriptorCount\r
+                    &DescriptorSize,              // DescriptorSize\r
+                    &PackageVersion,              // PackageVersion\r
+                    &PackageVersionName           // PackageVersionName\r
+                    );\r
+    if (EFI_ERROR(Status)) {\r
+      FreePool(FmpImageInfoBuf);\r
+      continue;\r
+    }\r
+\r
+    DEBUG((DEBUG_INFO, "FMP (%d) ImageInfo:\n", Index));\r
+    DumpFmpImageInfo(\r
+      ImageInfoSize,               // ImageInfoSize\r
+      FmpImageInfoBuf,             // ImageInfo\r
+      FmpImageInfoDescriptorVer,   // DescriptorVersion\r
+      FmpImageInfoCount,           // DescriptorCount\r
+      DescriptorSize,              // DescriptorSize\r
+      PackageVersion,              // PackageVersion\r
+      PackageVersionName           // PackageVersionName\r
+      );\r
+\r
+    if (PackageVersionName != NULL) {\r
+      FreePool(PackageVersionName);\r
+    }\r
+\r
+    FreePool(FmpImageInfoBuf);\r
+  }\r
+\r
+  return ;\r
+}\r
+\r
+/**\r
+  Get FMP handle by ImageTypeId and HardwareInstance.\r
+\r
+  @param[in]     UpdateImageTypeId       Used to identify device firmware targeted by this update.\r
+  @param[in]     UpdateHardwareInstance  The HardwareInstance to target with this update.\r
+  @param[in,out] NoHandles               The number of handles returned in Buffer.\r
+  @param[out]    Buffer[out]             A pointer to the buffer to return the requested array of handles.\r
+\r
+  @retval EFI_SUCCESS            The array of handles was returned in Buffer, and the number of\r
+                                 handles in Buffer was returned in NoHandles.\r
+  @retval EFI_NOT_FOUND          No handles match the search.\r
+  @retval EFI_OUT_OF_RESOURCES   There is not enough pool memory to store the matching results.\r
+**/\r
+EFI_STATUS\r
+GetFmpHandleBufferByType (\r
+  IN     EFI_GUID                     *UpdateImageTypeId,\r
+  IN     UINT64                       UpdateHardwareInstance,\r
+  IN OUT UINTN                        *NoHandles,\r
+  OUT    EFI_HANDLE                   **Buffer\r
+  )\r
+{\r
+  EFI_STATUS                                    Status;\r
+  EFI_HANDLE                                    *HandleBuffer;\r
+  UINTN                                         NumberOfHandles;\r
+  EFI_HANDLE                                    *MatchedHandleBuffer;\r
+  UINTN                                         MatchedNumberOfHandles;\r
+  EFI_FIRMWARE_MANAGEMENT_PROTOCOL              *Fmp;\r
+  UINTN                                         Index;\r
+  UINTN                                         ImageInfoSize;\r
+  EFI_FIRMWARE_IMAGE_DESCRIPTOR                 *FmpImageInfoBuf;\r
+  UINT32                                        FmpImageInfoDescriptorVer;\r
+  UINT8                                         FmpImageInfoCount;\r
+  UINTN                                         DescriptorSize;\r
+  UINT32                                        PackageVersion;\r
+  CHAR16                                        *PackageVersionName;\r
+  UINTN                                         Index2;\r
+  EFI_FIRMWARE_IMAGE_DESCRIPTOR                 *TempFmpImageInfo;\r
+\r
+  *NoHandles = 0;\r
+  *Buffer = NULL;\r
+\r
+  Status = gBS->LocateHandleBuffer (\r
+                  ByProtocol,\r
+                  &gEfiFirmwareManagementProtocolGuid,\r
+                  NULL,\r
+                  &NumberOfHandles,\r
+                  &HandleBuffer\r
+                  );\r
+  if (EFI_ERROR(Status)) {\r
+    return Status;\r
+  }\r
+\r
+  MatchedNumberOfHandles = 0;\r
+  MatchedHandleBuffer = AllocateZeroPool (sizeof(EFI_HANDLE) * NumberOfHandles);\r
+  if (MatchedHandleBuffer == NULL) {\r
+    FreePool (HandleBuffer);\r
+    return EFI_OUT_OF_RESOURCES;\r
+  }\r
+\r
+  for (Index = 0; Index < NumberOfHandles; Index++) {\r
+    Status = gBS->HandleProtocol(\r
+                    HandleBuffer[Index],\r
+                    &gEfiFirmwareManagementProtocolGuid,\r
+                    (VOID **)&Fmp\r
+                    );\r
+    if (EFI_ERROR(Status)) {\r
+      continue;\r
+    }\r
+\r
+    ImageInfoSize = 0;\r
+    Status = Fmp->GetImageInfo (\r
+                    Fmp,\r
+                    &ImageInfoSize,\r
+                    NULL,\r
+                    NULL,\r
+                    NULL,\r
+                    NULL,\r
+                    NULL,\r
+                    NULL\r
+                    );\r
+    if (Status != EFI_BUFFER_TOO_SMALL) {\r
+      continue;\r
+    }\r
+\r
+    FmpImageInfoBuf = AllocateZeroPool (ImageInfoSize);\r
+    if (FmpImageInfoBuf == NULL) {\r
+      continue;\r
+    }\r
+\r
+    PackageVersionName = NULL;\r
+    Status = Fmp->GetImageInfo (\r
+                    Fmp,\r
+                    &ImageInfoSize,               // ImageInfoSize\r
+                    FmpImageInfoBuf,              // ImageInfo\r
+                    &FmpImageInfoDescriptorVer,   // DescriptorVersion\r
+                    &FmpImageInfoCount,           // DescriptorCount\r
+                    &DescriptorSize,              // DescriptorSize\r
+                    &PackageVersion,              // PackageVersion\r
+                    &PackageVersionName           // PackageVersionName\r
+                    );\r
+    if (EFI_ERROR(Status)) {\r
+      FreePool(FmpImageInfoBuf);\r
+      continue;\r
+    }\r
+\r
+    if (PackageVersionName != NULL) {\r
+      FreePool(PackageVersionName);\r
+    }\r
+\r
+    TempFmpImageInfo = FmpImageInfoBuf;\r
+    for (Index2 = 0; Index2 < FmpImageInfoCount; Index2++) {\r
+      //\r
+      // Check if this FMP instance matches\r
+      //\r
+      if (CompareGuid(UpdateImageTypeId, &TempFmpImageInfo->ImageTypeId)) {\r
+        if ((UpdateHardwareInstance == 0) ||\r
+            ((FmpImageInfoDescriptorVer >= EFI_FIRMWARE_IMAGE_DESCRIPTOR_VERSION) &&\r
+             (UpdateHardwareInstance == TempFmpImageInfo->HardwareInstance))) {\r
+          MatchedHandleBuffer[MatchedNumberOfHandles] = HandleBuffer[Index];\r
+          MatchedNumberOfHandles++;\r
+          break;\r
+        }\r
+      }\r
+      TempFmpImageInfo = (EFI_FIRMWARE_IMAGE_DESCRIPTOR *)((UINT8 *)TempFmpImageInfo + DescriptorSize);\r
+    }\r
+    FreePool(FmpImageInfoBuf);\r
+  }\r
+\r
+  if (MatchedNumberOfHandles == 0) {\r
+    return EFI_NOT_FOUND;\r
+  }\r
+\r
+  *NoHandles = MatchedNumberOfHandles;\r
+  *Buffer = MatchedHandleBuffer;\r
+\r
+  return EFI_SUCCESS;\r
+}\r
+\r
+/**\r
+  Return FmpImageInfoDescriptorVer by an FMP handle.\r
+\r
+  @param[in]  Handle   A FMP handle.\r
+\r
+  @return FmpImageInfoDescriptorVer associated with the FMP.\r
+**/\r
+UINT32\r
+GetFmpImageInfoDescriptorVer (\r
+  IN EFI_HANDLE                                   Handle\r
+  )\r
+{\r
+  EFI_STATUS                                    Status;\r
+  EFI_FIRMWARE_MANAGEMENT_PROTOCOL              *Fmp;\r
+  UINTN                                         ImageInfoSize;\r
+  EFI_FIRMWARE_IMAGE_DESCRIPTOR                 *FmpImageInfoBuf;\r
+  UINT32                                        FmpImageInfoDescriptorVer;\r
+  UINT8                                         FmpImageInfoCount;\r
+  UINTN                                         DescriptorSize;\r
+  UINT32                                        PackageVersion;\r
+  CHAR16                                        *PackageVersionName;\r
+\r
+  Status = gBS->HandleProtocol(\r
+                  Handle,\r
+                  &gEfiFirmwareManagementProtocolGuid,\r
+                  (VOID **)&Fmp\r
+                  );\r
+  if (EFI_ERROR(Status)) {\r
+    return 0;\r
+  }\r
+\r
+  ImageInfoSize = 0;\r
+  Status = Fmp->GetImageInfo (\r
+                  Fmp,\r
+                  &ImageInfoSize,\r
+                  NULL,\r
+                  NULL,\r
+                  NULL,\r
+                  NULL,\r
+                  NULL,\r
+                  NULL\r
+                  );\r
+  if (Status != EFI_BUFFER_TOO_SMALL) {\r
+    return 0;\r
+  }\r
+\r
+  FmpImageInfoBuf = AllocateZeroPool (ImageInfoSize);\r
+  if (FmpImageInfoBuf == NULL) {\r
+    return 0;\r
+  }\r
+\r
+  PackageVersionName = NULL;\r
+  Status = Fmp->GetImageInfo (\r
+                  Fmp,\r
+                  &ImageInfoSize,               // ImageInfoSize\r
+                  FmpImageInfoBuf,              // ImageInfo\r
+                  &FmpImageInfoDescriptorVer,   // DescriptorVersion\r
+                  &FmpImageInfoCount,           // DescriptorCount\r
+                  &DescriptorSize,              // DescriptorSize\r
+                  &PackageVersion,              // PackageVersion\r
+                  &PackageVersionName           // PackageVersionName\r
+                  );\r
+  if (EFI_ERROR(Status)) {\r
+    FreePool(FmpImageInfoBuf);\r
+    return 0;\r
+  }\r
+  return FmpImageInfoDescriptorVer;\r
+}\r
+\r
+/**\r
+  Set FMP image data.\r
+\r
+  @param[in]  Handle        A FMP handle.\r
+  @param[in]  ImageHeader   The payload image header.\r
+  @param[in]  PayloadIndex  The index of the payload.\r
+\r
+  @return The status of FMP->SetImage.\r
+**/\r
+EFI_STATUS\r
+SetFmpImageData (\r
+  IN EFI_HANDLE                                   Handle,\r
+  IN EFI_FIRMWARE_MANAGEMENT_CAPSULE_IMAGE_HEADER *ImageHeader,\r
+  IN UINTN                                        PayloadIndex\r
+  )\r
+{\r
+  EFI_STATUS                                    Status;\r
+  EFI_FIRMWARE_MANAGEMENT_PROTOCOL              *Fmp;\r
+  UINT8                                         *Image;\r
+  VOID                                          *VendorCode;\r
+  CHAR16                                        *AbortReason;\r
+\r
+  Status = gBS->HandleProtocol(\r
+                  Handle,\r
+                  &gEfiFirmwareManagementProtocolGuid,\r
+                  (VOID **)&Fmp\r
+                  );\r
+  if (EFI_ERROR(Status)) {\r
+    return Status;\r
+  }\r
+\r
+  if (ImageHeader->Version >= EFI_FIRMWARE_MANAGEMENT_CAPSULE_IMAGE_HEADER_INIT_VERSION) {\r
+    Image = (UINT8 *)(ImageHeader + 1);\r
+  } else {\r
+    //\r
+    // If the EFI_FIRMWARE_MANAGEMENT_CAPSULE_IMAGE_HEADER is version 1,\r
+    // Header should exclude UpdateHardwareInstance field\r
+    //\r
+    Image = (UINT8 *)ImageHeader + OFFSET_OF(EFI_FIRMWARE_MANAGEMENT_CAPSULE_IMAGE_HEADER, UpdateHardwareInstance);\r
+  }\r
+\r
+  if (ImageHeader->UpdateVendorCodeSize == 0) {\r
+    VendorCode = NULL;\r
+  } else {\r
+    VendorCode = Image + ImageHeader->UpdateImageSize;\r
+  }\r
+  AbortReason = NULL;\r
+  DEBUG((DEBUG_INFO, "Fmp->SetImage ...\n"));\r
+  DEBUG((DEBUG_INFO, "ImageTypeId - %g, ", &ImageHeader->UpdateImageTypeId));\r
+  DEBUG((DEBUG_INFO, "PayloadIndex - 0x%x, ", PayloadIndex));\r
+  DEBUG((DEBUG_INFO, "ImageIndex - 0x%x ", ImageHeader->UpdateImageIndex));\r
+  if (ImageHeader->Version >= EFI_FIRMWARE_MANAGEMENT_CAPSULE_IMAGE_HEADER_INIT_VERSION) {\r
+    DEBUG((DEBUG_INFO, "(UpdateHardwareInstance - 0x%x)", ImageHeader->UpdateHardwareInstance));\r
+  }\r
+  DEBUG((DEBUG_INFO, "\n"));\r
+  Status = Fmp->SetImage(\r
+                  Fmp,\r
+                  ImageHeader->UpdateImageIndex,          // ImageIndex\r
+                  Image,                                  // Image\r
+                  ImageHeader->UpdateImageSize,           // ImageSize\r
+                  VendorCode,                             // VendorCode\r
+                  Update_Image_Progress,                  // Progress\r
+                  &AbortReason                            // AbortReason\r
+                  );\r
+  DEBUG((DEBUG_INFO, "Fmp->SetImage - %r\n", Status));\r
+  if (AbortReason != NULL) {\r
+    DEBUG ((DEBUG_ERROR, "%s\n", AbortReason));\r
+    FreePool(AbortReason);\r
+  }\r
+\r
+  return Status;\r
+}\r
+\r
+/**\r
+  Start a UEFI image in the FMP payload.\r
+\r
+  @param[in]  ImageBuffer   A pointer to the memory location containing a copy of the image to be loaded..\r
+  @param[in]  ImageSize     The size in bytes of ImageBuffer.\r
+\r
+  @return The status of gBS->LoadImage and gBS->StartImage.\r
+**/\r
+EFI_STATUS\r
+StartFmpImage (\r
+  IN VOID   *ImageBuffer,\r
+  IN UINTN  ImageSize\r
+  )\r
+{\r
+  MEMMAP_DEVICE_PATH                            MemMapNode;\r
+  EFI_STATUS                                    Status;\r
+  EFI_HANDLE                                    ImageHandle;\r
+  EFI_DEVICE_PATH_PROTOCOL                      *DriverDevicePath;\r
+  UINTN                                         ExitDataSize;\r
+\r
+  SetDevicePathNodeLength (&MemMapNode.Header, sizeof (MemMapNode));\r
+  MemMapNode.Header.Type     = HARDWARE_DEVICE_PATH;\r
+  MemMapNode.Header.SubType  = HW_MEMMAP_DP;\r
+  MemMapNode.MemoryType      = EfiBootServicesCode;\r
+  MemMapNode.StartingAddress = (EFI_PHYSICAL_ADDRESS)(UINTN)ImageBuffer;\r
+  MemMapNode.EndingAddress   = (EFI_PHYSICAL_ADDRESS)(UINTN)((UINT8 *)ImageBuffer + ImageSize - 1);\r
+\r
+  DriverDevicePath = AppendDevicePathNode (NULL, &MemMapNode.Header);\r
+  if (DriverDevicePath == NULL) {\r
+    return EFI_OUT_OF_RESOURCES;\r
+  }\r
+\r
+  DEBUG((DEBUG_INFO, "FmpCapsule: LoadImage ...\n"));\r
+  Status = gBS->LoadImage(\r
+                  FALSE,\r
+                  gImageHandle,\r
+                  DriverDevicePath,\r
+                  ImageBuffer,\r
+                  ImageSize,\r
+                  &ImageHandle\r
+                  );\r
+  DEBUG((DEBUG_INFO, "FmpCapsule: LoadImage - %r\n", Status));\r
+  if (EFI_ERROR(Status)) {\r
+    FreePool(DriverDevicePath);\r
+    return Status;\r
+  }\r
+\r
+  DEBUG((DEBUG_INFO, "FmpCapsule: StartImage ...\n"));\r
+  Status = gBS->StartImage(\r
+                  ImageHandle,\r
+                  &ExitDataSize,\r
+                  NULL\r
+                  );\r
+  DEBUG((DEBUG_INFO, "FmpCapsule: StartImage - %r\n", Status));\r
+  if (EFI_ERROR(Status)) {\r
+    DEBUG ((DEBUG_ERROR, "Driver Return Status = %r\n", Status));\r
+  }\r
+\r
+  FreePool(DriverDevicePath);\r
+  return Status;\r
+}\r
+\r
+/**\r
+  Record FMP capsule status.\r
+\r
+  @param[in]  Handle        A FMP handle.\r
+  @param[in] CapsuleHeader  The capsule image header\r
+  @param[in] CapsuleStatus  The capsule process stauts\r
+  @param[in] PayloadIndex   FMP payload index\r
+  @param[in] ImageHeader    FMP image header\r
+**/\r
+VOID\r
+RecordFmpCapsuleStatus (\r
+  IN EFI_HANDLE                                    Handle,  OPTIONAL\r
+  IN EFI_CAPSULE_HEADER                            *CapsuleHeader,\r
+  IN EFI_STATUS                                    CapsuleStatus,\r
+  IN UINTN                                         PayloadIndex,\r
+  IN EFI_FIRMWARE_MANAGEMENT_CAPSULE_IMAGE_HEADER  *ImageHeader\r
+  )\r
+{\r
+  EFI_STATUS                                    Status;\r
+  EFI_DEVICE_PATH_PROTOCOL                      *FmpDevicePath;\r
+  UINT32                                        FmpImageInfoDescriptorVer;\r
+  EFI_STATUS                                    StatusEsrt;\r
+  ESRT_MANAGEMENT_PROTOCOL                      *EsrtProtocol;\r
+  EFI_SYSTEM_RESOURCE_ENTRY                     EsrtEntry;\r
+\r
+  FmpDevicePath = NULL;\r
+  if (Handle != NULL) {\r
+    gBS->HandleProtocol(\r
+           Handle,\r
+           &gEfiDevicePathProtocolGuid,\r
+           (VOID **)&FmpDevicePath\r
+           );\r
+  }\r
+\r
+  RecordFmpCapsuleStatusVariable (\r
+    CapsuleHeader,\r
+    CapsuleStatus,\r
+    PayloadIndex,\r
+    ImageHeader,\r
+    FmpDevicePath\r
+    );\r
+\r
+  //\r
+  // Update corresponding ESRT entry LastAttemp Status\r
+  //\r
+  Status = gBS->LocateProtocol(&gEsrtManagementProtocolGuid, NULL, (VOID **)&EsrtProtocol);\r
+  if (EFI_ERROR (Status)) {\r
+    return ;\r
+  }\r
+\r
+  if (Handle == NULL) {\r
+    return ;\r
+  }\r
+\r
+  //\r
+  // Update EsrtEntry For V1, V2 FMP instance.\r
+  // V3 FMP ESRT cache will be synced up through EsrtSyncFmp interface\r
+  //\r
+  FmpImageInfoDescriptorVer = GetFmpImageInfoDescriptorVer (Handle);\r
+  if (FmpImageInfoDescriptorVer < EFI_FIRMWARE_IMAGE_DESCRIPTOR_VERSION) {\r
+    StatusEsrt = EsrtProtocol->GetEsrtEntry(&ImageHeader->UpdateImageTypeId, &EsrtEntry);\r
+    if (!EFI_ERROR(StatusEsrt)){\r
+      if (!EFI_ERROR(CapsuleStatus)) {\r
+        EsrtEntry.LastAttemptStatus = LAST_ATTEMPT_STATUS_SUCCESS;\r
+      } else {\r
+        EsrtEntry.LastAttemptStatus = LAST_ATTEMPT_STATUS_ERROR_UNSUCCESSFUL;\r
+      }\r
+      EsrtEntry.LastAttemptVersion = 0;\r
+      EsrtProtocol->UpdateEsrtEntry(&EsrtEntry);\r
+    }\r
+  }\r
+}\r
+\r
 /**\r
   Process Firmware management protocol data capsule.\r
 \r
@@ -777,64 +1273,38 @@ DumpFmpCapsule (
   This function need support nested FMP capsule.\r
 \r
   @param[in]   CapsuleHeader         Points to a capsule header.\r
-  @param[out]  AreAllImagesProcessed If all the FMP images in the capsule are processed.\r
 \r
   @retval EFI_SUCESS            Process Capsule Image successfully.\r
   @retval EFI_UNSUPPORTED       Capsule image is not supported by the firmware.\r
   @retval EFI_VOLUME_CORRUPTED  FV volume in the capsule is corrupted.\r
   @retval EFI_OUT_OF_RESOURCES  Not enough memory.\r
+  @retval EFI_NOT_READY         No FMP protocol to handle this FMP capsule.\r
 **/\r
 EFI_STATUS\r
 ProcessFmpCapsuleImage (\r
-  IN EFI_CAPSULE_HEADER  *CapsuleHeader,\r
-  OUT BOOLEAN            *AreAllImagesProcessed\r
+  IN EFI_CAPSULE_HEADER  *CapsuleHeader\r
   )\r
 {\r
   EFI_STATUS                                    Status;\r
-  EFI_STATUS                                    StatusEsrt;\r
-  EFI_STATUS                                    StatusRet;\r
   EFI_FIRMWARE_MANAGEMENT_CAPSULE_HEADER        *FmpCapsuleHeader;\r
   EFI_FIRMWARE_MANAGEMENT_CAPSULE_IMAGE_HEADER  *ImageHeader;\r
-  UINT8                                         *Image;\r
-  EFI_HANDLE                                    ImageHandle;\r
   UINT64                                        *ItemOffsetList;\r
   UINT32                                        ItemNum;\r
   UINTN                                         Index;\r
-  UINTN                                         ExitDataSize;\r
   EFI_HANDLE                                    *HandleBuffer;\r
-  EFI_FIRMWARE_MANAGEMENT_PROTOCOL              *Fmp;\r
   UINTN                                         NumberOfHandles;\r
-  UINTN                                         DescriptorSize;\r
-  UINT8                                         FmpImageInfoCount;\r
-  UINT32                                        FmpImageInfoDescriptorVer;\r
-  UINTN                                         ImageInfoSize;\r
-  UINT32                                        PackageVersion;\r
-  CHAR16                                        *PackageVersionName;\r
-  CHAR16                                        *AbortReason;\r
-  EFI_FIRMWARE_IMAGE_DESCRIPTOR                 *FmpImageInfoBuf;\r
-  EFI_FIRMWARE_IMAGE_DESCRIPTOR                 *TempFmpImageInfo;\r
   UINTN                                         DriverLen;\r
-  UINTN                                         Index1;\r
+  UINT64                                        UpdateHardwareInstance;\r
   UINTN                                         Index2;\r
-  MEMMAP_DEVICE_PATH                            MemMapNode;\r
-  EFI_DEVICE_PATH_PROTOCOL                      *DriverDevicePath;\r
-  ESRT_MANAGEMENT_PROTOCOL                      *EsrtProtocol;\r
-  EFI_SYSTEM_RESOURCE_ENTRY                     EsrtEntry;\r
-  VOID                                          *VendorCode;\r
+  BOOLEAN                                       NotReady;\r
+  BOOLEAN                                       Abort;\r
 \r
   if (!IsFmpCapsuleGuid(&CapsuleHeader->CapsuleGuid)) {\r
-    return ProcessFmpCapsuleImage ((EFI_CAPSULE_HEADER *)((UINTN)CapsuleHeader + CapsuleHeader->HeaderSize), AreAllImagesProcessed);\r
+    return ProcessFmpCapsuleImage ((EFI_CAPSULE_HEADER *)((UINTN)CapsuleHeader + CapsuleHeader->HeaderSize));\r
   }\r
 \r
-  ASSERT(AreAllImagesProcessed != NULL);\r
-\r
-  Status           = EFI_SUCCESS;\r
-  StatusRet        = EFI_NOT_FOUND;\r
-  HandleBuffer     = NULL;\r
-  ExitDataSize     = 0;\r
-  DriverDevicePath = NULL;\r
-  EsrtProtocol     = NULL;\r
-  *AreAllImagesProcessed = FALSE;\r
+  NotReady = FALSE;\r
+  Abort = FALSE;\r
 \r
   DumpFmpCapsule(CapsuleHeader);\r
 \r
@@ -851,35 +1321,15 @@ ProcessFmpCapsuleImage (
   // capsule in which driver count and payload count are both zero is not processed.\r
   //\r
   if (ItemNum == 0) {\r
-    *AreAllImagesProcessed = TRUE;\r
     return EFI_SUCCESS;\r
   }\r
 \r
-  //\r
-  // Update corresponding ESRT entry LastAttemp Status\r
-  //\r
-  Status = gBS->LocateProtocol(&gEsrtManagementProtocolGuid, NULL, (VOID **)&EsrtProtocol);\r
-  if (EFI_ERROR (Status)) {\r
-    EsrtProtocol = NULL;\r
-  }\r
-\r
   //\r
   // 1. Try to load & start all the drivers within capsule\r
   //\r
-  SetDevicePathNodeLength (&MemMapNode.Header, sizeof (MemMapNode));\r
-  MemMapNode.Header.Type     = HARDWARE_DEVICE_PATH;\r
-  MemMapNode.Header.SubType  = HW_MEMMAP_DP;\r
-  MemMapNode.MemoryType      = EfiBootServicesCode;\r
-  MemMapNode.StartingAddress = (EFI_PHYSICAL_ADDRESS)(UINTN)CapsuleHeader;\r
-  MemMapNode.EndingAddress   = (EFI_PHYSICAL_ADDRESS)(UINTN)((UINT8 *)CapsuleHeader + CapsuleHeader->CapsuleImageSize - 1);\r
-\r
-  DriverDevicePath = AppendDevicePathNode (NULL, &MemMapNode.Header);\r
-  if (DriverDevicePath == NULL) {\r
-    return EFI_OUT_OF_RESOURCES;\r
-  }\r
-\r
   for (Index = 0; Index < FmpCapsuleHeader->EmbeddedDriverCount; Index++) {\r
-    if (FmpCapsuleHeader->PayloadItemCount == 0 && Index == (UINTN)FmpCapsuleHeader->EmbeddedDriverCount - 1) {\r
+    if ((FmpCapsuleHeader->PayloadItemCount == 0) &&\r
+        (Index == (UINTN)FmpCapsuleHeader->EmbeddedDriverCount - 1)) {\r
       //\r
       // When driver is last element in the ItemOffsetList array, the driver size is calculated by reference CapsuleImageSize in EFI_CAPSULE_HEADER\r
       //\r
@@ -888,32 +1338,13 @@ ProcessFmpCapsuleImage (
       DriverLen = (UINTN)ItemOffsetList[Index + 1] - (UINTN)ItemOffsetList[Index];\r
     }\r
 \r
-    DEBUG((DEBUG_INFO, "FmpCapsule: LoadImage ...\n"));\r
-    Status = gBS->LoadImage(\r
-                    FALSE,\r
-                    gImageHandle,\r
-                    DriverDevicePath,\r
-                    (UINT8 *)FmpCapsuleHeader + ItemOffsetList[Index],\r
-                    DriverLen,\r
-                    &ImageHandle\r
-                    );\r
-    DEBUG((DEBUG_INFO, "FmpCapsule: LoadImage - %r\n", Status));\r
-    if (EFI_ERROR(Status)) {\r
-      StatusRet = Status;\r
-      goto EXIT;\r
-    }\r
-\r
-    DEBUG((DEBUG_INFO, "FmpCapsule: StartImage ...\n"));\r
-    Status = gBS->StartImage(\r
-                    ImageHandle,\r
-                    &ExitDataSize,\r
-                    NULL\r
-                    );\r
-    DEBUG((DEBUG_INFO, "FmpCapsule: StartImage - %r\n", Status));\r
+    Status = StartFmpImage (\r
+               (UINT8 *)FmpCapsuleHeader + ItemOffsetList[Index],\r
+               DriverLen\r
+               );\r
     if (EFI_ERROR(Status)) {\r
       DEBUG ((DEBUG_ERROR, "Driver Return Status = %r\n", Status));\r
-      StatusRet = Status;\r
-      goto EXIT;\r
+      return Status;\r
     }\r
   }\r
 \r
@@ -922,199 +1353,80 @@ ProcessFmpCapsuleImage (
   //\r
   DEBUG((DEBUG_INFO, "FmpCapsule: route payload to right FMP instance ...\n"));\r
 \r
-  Status = gBS->LocateHandleBuffer (\r
-                  ByProtocol,\r
-                  &gEfiFirmwareManagementProtocolGuid,\r
-                  NULL,\r
-                  &NumberOfHandles,\r
-                  &HandleBuffer\r
-                  );\r
-\r
-  if (!EFI_ERROR(Status)) {\r
-    for(Index1 = 0; Index1 < NumberOfHandles; Index1++) {\r
-      Status = gBS->HandleProtocol(\r
-                      HandleBuffer[Index1],\r
-                      &gEfiFirmwareManagementProtocolGuid,\r
-                      (VOID **)&Fmp\r
-                      );\r
-      if (EFI_ERROR(Status)) {\r
-        continue;\r
-      }\r
+  DumpAllFmpInfo ();\r
 \r
-      ImageInfoSize = 0;\r
-      Status = Fmp->GetImageInfo (\r
-                      Fmp,\r
-                      &ImageInfoSize,\r
-                      NULL,\r
-                      NULL,\r
-                      NULL,\r
-                      NULL,\r
-                      NULL,\r
-                      NULL\r
-                      );\r
-      if (Status != EFI_BUFFER_TOO_SMALL) {\r
-        continue;\r
-      }\r
+  //\r
+  // Check all the payload entry in capsule payload list\r
+  //\r
+  for (Index = FmpCapsuleHeader->EmbeddedDriverCount; Index < ItemNum; Index++) {\r
+    ImageHeader  = (EFI_FIRMWARE_MANAGEMENT_CAPSULE_IMAGE_HEADER *)((UINT8 *)FmpCapsuleHeader + ItemOffsetList[Index]);\r
 \r
-      FmpImageInfoBuf = NULL;\r
-      FmpImageInfoBuf = AllocateZeroPool (ImageInfoSize);\r
-      if (FmpImageInfoBuf == NULL) {\r
-        StatusRet = EFI_OUT_OF_RESOURCES;\r
-        goto EXIT;\r
-      }\r
+    UpdateHardwareInstance = 0;\r
+    if (ImageHeader->Version >= EFI_FIRMWARE_MANAGEMENT_CAPSULE_IMAGE_HEADER_INIT_VERSION) {\r
+      UpdateHardwareInstance = ImageHeader->UpdateHardwareInstance;\r
+    }\r
 \r
-      PackageVersionName = NULL;\r
-      Status = Fmp->GetImageInfo (\r
-                      Fmp,\r
-                      &ImageInfoSize,               // ImageInfoSize\r
-                      FmpImageInfoBuf,              // ImageInfo\r
-                      &FmpImageInfoDescriptorVer,   // DescriptorVersion\r
-                      &FmpImageInfoCount,           // DescriptorCount\r
-                      &DescriptorSize,              // DescriptorSize\r
-                      &PackageVersion,              // PackageVersion\r
-                      &PackageVersionName           // PackageVersionName\r
-                      );\r
+    Status = GetFmpHandleBufferByType (\r
+               &ImageHeader->UpdateImageTypeId,\r
+               UpdateHardwareInstance,\r
+               &NumberOfHandles,\r
+               &HandleBuffer\r
+               );\r
+    if (EFI_ERROR(Status)) {\r
+      NotReady = TRUE;\r
+      RecordFmpCapsuleStatus (\r
+        NULL,\r
+        CapsuleHeader,\r
+        EFI_NOT_READY,\r
+        Index - FmpCapsuleHeader->EmbeddedDriverCount,\r
+        ImageHeader\r
+        );\r
+      continue;\r
+    }\r
 \r
-      //\r
-      // If FMP GetInformation interface failed, skip this resource\r
-      //\r
-      if (EFI_ERROR(Status)) {\r
-        FreePool(FmpImageInfoBuf);\r
+    for (Index2 = 0; Index2 < NumberOfHandles; Index2++) {\r
+      if (Abort) {\r
+        RecordFmpCapsuleStatus (\r
+          HandleBuffer[Index2],\r
+          CapsuleHeader,\r
+          EFI_ABORTED,\r
+          Index - FmpCapsuleHeader->EmbeddedDriverCount,\r
+          ImageHeader\r
+          );\r
         continue;\r
       }\r
 \r
-      DEBUG((DEBUG_INFO, "FMP (%d) ImageInfo:\n", Index));\r
-      DumpFmpImageInfo(\r
-        ImageInfoSize,               // ImageInfoSize\r
-        FmpImageInfoBuf,             // ImageInfo\r
-        FmpImageInfoDescriptorVer,   // DescriptorVersion\r
-        FmpImageInfoCount,           // DescriptorCount\r
-        DescriptorSize,              // DescriptorSize\r
-        PackageVersion,              // PackageVersion\r
-        PackageVersionName           // PackageVersionName\r
-        );\r
-\r
-      if (PackageVersionName != NULL) {\r
-        FreePool(PackageVersionName);\r
+      Status = SetFmpImageData (\r
+                 HandleBuffer[Index2],\r
+                 ImageHeader,\r
+                 Index - FmpCapsuleHeader->EmbeddedDriverCount\r
+                 );\r
+      if (Status != EFI_SUCCESS) {\r
+        Abort = TRUE;\r
       }\r
 \r
-      TempFmpImageInfo = FmpImageInfoBuf;\r
-      for (Index2 = 0; Index2 < FmpImageInfoCount; Index2++) {\r
-        //\r
-        // Check all the payload entry in capsule payload list\r
-        //\r
-        for (Index = FmpCapsuleHeader->EmbeddedDriverCount; Index < ItemNum; Index++) {\r
-          ImageHeader  = (EFI_FIRMWARE_MANAGEMENT_CAPSULE_IMAGE_HEADER *)((UINT8 *)FmpCapsuleHeader + ItemOffsetList[Index]);\r
-\r
-          if (IsFmpCapsuleProcessed(CapsuleHeader, Index - FmpCapsuleHeader->EmbeddedDriverCount, ImageHeader)) {\r
-            DEBUG((DEBUG_INFO, "FMP Capsule already processed (%g):", CapsuleHeader));\r
-            DEBUG((DEBUG_INFO, "ImageTypeId - %g, ", &ImageHeader->UpdateImageTypeId));\r
-            DEBUG((DEBUG_INFO, "PayloadIndex - 0x%x, ImageIndex - 0x%x\n", Index - FmpCapsuleHeader->EmbeddedDriverCount, ImageHeader->UpdateImageIndex));\r
-            continue;\r
-          }\r
-\r
-          if (CompareGuid(&ImageHeader->UpdateImageTypeId, &TempFmpImageInfo->ImageTypeId)) {\r
-            AbortReason = NULL;\r
-            if (ImageHeader->Version >= EFI_FIRMWARE_MANAGEMENT_CAPSULE_IMAGE_HEADER_INIT_VERSION) {\r
-              if(ImageHeader->UpdateHardwareInstance != 0){\r
-                //\r
-                // FMP Version is >=2 & UpdateHardwareInstance Skip 2 case\r
-                //  1. FMP Image info Version < 3\r
-                //  2. HardwareInstance doesn't match\r
-                //\r
-                if (FmpImageInfoDescriptorVer < EFI_FIRMWARE_IMAGE_DESCRIPTOR_VERSION ||\r
-                   ImageHeader->UpdateHardwareInstance != TempFmpImageInfo->HardwareInstance) {\r
-                  continue;\r
-                }\r
-              }\r
-              Image = (UINT8 *)(ImageHeader + 1);\r
-            } else {\r
-              //\r
-              // If the EFI_FIRMWARE_MANAGEMENT_CAPSULE_IMAGE_HEADER is version 1, only match ImageTypeId.\r
-              // Header should exclude UpdateHardwareInstance field\r
-              //\r
-              Image = (UINT8 *)ImageHeader + OFFSET_OF(EFI_FIRMWARE_MANAGEMENT_CAPSULE_IMAGE_HEADER, UpdateHardwareInstance);\r
-            }\r
-\r
-            if (ImageHeader->UpdateVendorCodeSize == 0) {\r
-              VendorCode = NULL;\r
-            } else {\r
-              VendorCode = Image + ImageHeader->UpdateImageSize;\r
-            }\r
-            DEBUG((DEBUG_INFO, "Fmp->SetImage ...\n"));\r
-            Status = Fmp->SetImage(\r
-                            Fmp,\r
-                            ImageHeader->UpdateImageIndex,          // ImageIndex\r
-                            Image,                                  // Image\r
-                            ImageHeader->UpdateImageSize,           // ImageSize\r
-                            VendorCode,                                   // VendorCode\r
-                            Update_Image_Progress,                  // Progress\r
-                            &AbortReason                            // AbortReason\r
-                            );\r
-            DEBUG((DEBUG_INFO, "Fmp->SetImage - %r\n", Status));\r
-            if (AbortReason != NULL) {\r
-              DEBUG ((DEBUG_ERROR, "%s\n", AbortReason));\r
-              FreePool(AbortReason);\r
-            }\r
-            RecordFmpCapsuleStatusVariable(\r
-              CapsuleHeader,                                 // CapsuleGuid\r
-              Status,                                        // CapsuleStatus\r
-              Index - FmpCapsuleHeader->EmbeddedDriverCount, // PayloadIndex\r
-              ImageHeader                                    // ImageHeader\r
-              );\r
-            if (StatusRet != EFI_SUCCESS) {\r
-              StatusRet = Status;\r
-            }\r
-            //\r
-            // Update EsrtEntry For V1, V2 FMP instance. V3 FMP ESRT cache will be synced up through EsrtSyncFmp interface\r
-            //\r
-            if (FmpImageInfoDescriptorVer < EFI_FIRMWARE_IMAGE_DESCRIPTOR_VERSION && EsrtProtocol != NULL) {\r
-               StatusEsrt = EsrtProtocol->GetEsrtEntry(&TempFmpImageInfo->ImageTypeId, &EsrtEntry);\r
-               if (!EFI_ERROR(StatusEsrt)){\r
-                 if (!EFI_ERROR(Status)) {\r
-                   EsrtEntry.LastAttemptStatus = LAST_ATTEMPT_STATUS_SUCCESS;\r
-                 } else {\r
-                   EsrtEntry.LastAttemptStatus = LAST_ATTEMPT_STATUS_ERROR_UNSUCCESSFUL;\r
-                 }\r
-                 EsrtEntry.LastAttemptVersion = 0;\r
-                 EsrtProtocol->UpdateEsrtEntry(&EsrtEntry);\r
-               }\r
-             }\r
-          }\r
-        }\r
-        //\r
-        // Use DescriptorSize to move ImageInfo Pointer to stay compatible with different ImageInfo version\r
-        //\r
-        TempFmpImageInfo = (EFI_FIRMWARE_IMAGE_DESCRIPTOR *)((UINT8 *)TempFmpImageInfo + DescriptorSize);\r
-      }\r
-      FreePool(FmpImageInfoBuf);\r
+      RecordFmpCapsuleStatus (\r
+        HandleBuffer[Index2],\r
+        CapsuleHeader,\r
+        Status,\r
+        Index - FmpCapsuleHeader->EmbeddedDriverCount,\r
+        ImageHeader\r
+        );\r
     }\r
-  }\r
-\r
-  //\r
-  // final check for AreAllImagesProcessed\r
-  //\r
-  *AreAllImagesProcessed = TRUE;\r
-  for (Index = FmpCapsuleHeader->EmbeddedDriverCount; Index < ItemNum; Index++) {\r
-    ImageHeader = (EFI_FIRMWARE_MANAGEMENT_CAPSULE_IMAGE_HEADER *)((UINT8 *)FmpCapsuleHeader + ItemOffsetList[Index]);\r
-\r
-    if (!IsFmpCapsuleProcessed(CapsuleHeader, Index - FmpCapsuleHeader->EmbeddedDriverCount, ImageHeader)) {\r
-      *AreAllImagesProcessed = FALSE;\r
-      break;\r
+    if (HandleBuffer != NULL) {\r
+      FreePool(HandleBuffer);\r
     }\r
   }\r
 \r
-EXIT:\r
-\r
-  if (HandleBuffer != NULL) {\r
-    FreePool(HandleBuffer);\r
+  if (NotReady) {\r
+    return EFI_NOT_READY;\r
   }\r
 \r
-  if (DriverDevicePath != NULL) {\r
-    FreePool(DriverDevicePath);\r
-  }\r
-\r
-  return StatusRet;\r
+  //\r
+  // always return SUCCESS to indicate this capsule is processed.\r
+  // The status of SetImage is recorded in capsule result variable.\r
+  //\r
+  return EFI_SUCCESS;\r
 }\r
 \r
 /**\r
@@ -1141,33 +1453,44 @@ IsNestedFmpCapsule (
   EFI_SYSTEM_RESOURCE_ENTRY  Entry;\r
 \r
   EsrtGuidFound = FALSE;\r
-\r
-  //\r
-  // Check ESRT protocol\r
-  //\r
-  Status = gBS->LocateProtocol(&gEsrtManagementProtocolGuid, NULL, (VOID **)&EsrtProtocol);\r
-  if (!EFI_ERROR(Status)) {\r
-    Status = EsrtProtocol->GetEsrtEntry(&CapsuleHeader->CapsuleGuid, &Entry);\r
-    if (!EFI_ERROR(Status)) {\r
-      EsrtGuidFound = TRUE;\r
-    }\r
-  }\r
-\r
-  //\r
-  // Check ESRT configuration table\r
-  //\r
-  if (!EsrtGuidFound) {\r
-    Status = EfiGetSystemConfigurationTable(&gEfiSystemResourceTableGuid, (VOID **)&Esrt);\r
-    if (!EFI_ERROR(Status)) {\r
-      ASSERT (Esrt != NULL);\r
-      EsrtEntry = (VOID *)(Esrt + 1);\r
-      for (Index = 0; Index < Esrt->FwResourceCount; Index++, EsrtEntry++) {\r
+  if (mIsVirtualAddrConverted) {\r
+    if(mEsrtTable != NULL) {\r
+      EsrtEntry = (EFI_SYSTEM_RESOURCE_ENTRY *)(mEsrtTable + 1);\r
+      for (Index = 0; Index < mEsrtTable->FwResourceCount ; Index++, EsrtEntry++) {\r
         if (CompareGuid(&EsrtEntry->FwClass, &CapsuleHeader->CapsuleGuid)) {\r
           EsrtGuidFound = TRUE;\r
           break;\r
         }\r
       }\r
     }\r
+  } else {\r
+    //\r
+    // Check ESRT protocol\r
+    //\r
+    Status = gBS->LocateProtocol(&gEsrtManagementProtocolGuid, NULL, (VOID **)&EsrtProtocol);\r
+    if (!EFI_ERROR(Status)) {\r
+      Status = EsrtProtocol->GetEsrtEntry(&CapsuleHeader->CapsuleGuid, &Entry);\r
+      if (!EFI_ERROR(Status)) {\r
+        EsrtGuidFound = TRUE;\r
+      }\r
+    }\r
+\r
+    //\r
+    // Check ESRT configuration table\r
+    //\r
+    if (!EsrtGuidFound) {\r
+      Status = EfiGetSystemConfigurationTable(&gEfiSystemResourceTableGuid, (VOID **)&Esrt);\r
+      if (!EFI_ERROR(Status)) {\r
+        ASSERT (Esrt != NULL);\r
+        EsrtEntry = (VOID *)(Esrt + 1);\r
+        for (Index = 0; Index < Esrt->FwResourceCount; Index++, EsrtEntry++) {\r
+          if (CompareGuid(&EsrtEntry->FwClass, &CapsuleHeader->CapsuleGuid)) {\r
+            EsrtGuidFound = TRUE;\r
+            break;\r
+          }\r
+        }\r
+      }\r
+    }\r
   }\r
   if (!EsrtGuidFound) {\r
     return FALSE;\r
@@ -1267,9 +1590,9 @@ ProcessCapsuleImage (
   )\r
 {\r
   EFI_STATUS                   Status;\r
-  BOOLEAN                      AreAllImagesProcessed;\r
 \r
   if (SupportCapsuleImage (CapsuleHeader) != EFI_SUCCESS) {\r
+    RecordCapsuleStatusVariable(CapsuleHeader, EFI_UNSUPPORTED);\r
     return EFI_UNSUPPORTED;\r
   }\r
 \r
@@ -1292,6 +1615,7 @@ ProcessCapsuleImage (
     Status = ValidateFmpCapsule(CapsuleHeader, NULL);\r
     DEBUG((DEBUG_INFO, "ValidateFmpCapsule - %r\n", Status));\r
     if (EFI_ERROR(Status)) {\r
+      RecordCapsuleStatusVariable(CapsuleHeader, Status);\r
       return Status;\r
     }\r
 \r
@@ -1299,13 +1623,9 @@ ProcessCapsuleImage (
     // Press EFI FMP Capsule\r
     //\r
     DEBUG((DEBUG_INFO, "ProcessFmpCapsuleImage ...\n"));\r
-    Status = ProcessFmpCapsuleImage(CapsuleHeader, &AreAllImagesProcessed);\r
+    Status = ProcessFmpCapsuleImage(CapsuleHeader);\r
     DEBUG((DEBUG_INFO, "ProcessFmpCapsuleImage - %r\n", Status));\r
 \r
-    if (!AreAllImagesProcessed) {\r
-      mAreAllImagesProcessed = FALSE;\r
-    }\r
-\r
     return Status;\r
   }\r
 \r
@@ -1344,7 +1664,6 @@ DxeCapsuleLibConstructor (
   IN EFI_SYSTEM_TABLE   *SystemTable\r
   )\r
 {\r
-  EFI_EVENT     EndOfDxeEvent;\r
   EFI_STATUS    Status;\r
 \r
   Status = gBS->CreateEventEx (\r
@@ -1353,7 +1672,7 @@ DxeCapsuleLibConstructor (
                   DxeCapsuleLibEndOfDxe,\r
                   NULL,\r
                   &gEfiEndOfDxeEventGroupGuid,\r
-                  &EndOfDxeEvent\r
+                  &mDxeCapsuleLibEndOfDxeEvent\r
                   );\r
   ASSERT_EFI_ERROR (Status);\r
 \r
@@ -1361,3 +1680,29 @@ DxeCapsuleLibConstructor (
 \r
   return EFI_SUCCESS;\r
 }\r
+\r
+/**\r
+  The destructor function closes the End of DXE event.\r
+\r
+  @param  ImageHandle   The firmware allocated handle for the EFI image.\r
+  @param  SystemTable   A pointer to the EFI System Table.\r
+\r
+  @retval EFI_SUCCESS   The destructor completed successfully.\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+DxeCapsuleLibDestructor (\r
+  IN EFI_HANDLE         ImageHandle,\r
+  IN EFI_SYSTEM_TABLE   *SystemTable\r
+  )\r
+{\r
+  EFI_STATUS    Status;\r
+\r
+  //\r
+  // Close the End of DXE event.\r
+  //\r
+  Status = gBS->CloseEvent (mDxeCapsuleLibEndOfDxeEvent);\r
+  ASSERT_EFI_ERROR (Status);\r
+\r
+  return EFI_SUCCESS;\r
+}\r