[PATCH] MdeModulePkg/DxeCapsuleLibFmp: Add runtime SetImage support

["Bob Morgan" <bobm@nvidia.com>]

Adds optional support for processing FMP capusle images after
ExitBootServices() if the ImageTypeIdGuid is mentioned in the new
PcdRuntimeFmpCapsuleImageTypeIdGuid list.

Cc: Jian J Wang <jian.j.wang@intel.com>
Cc: Liming Gao <gaoliming@byosoft.com.cn>
Cc: Guomin Jiang <guomin.jiang@intel.com>
Signed-off-by: Bob Morgan <bobm@nvidia.com>
---
 .../Library/DxeCapsuleLibFmp/DxeCapsuleLib.c  |  81 +++++++++---
 .../DxeCapsuleLibFmp/DxeCapsuleRuntime.c      | 119 ++++++++++++++++++
 .../DxeCapsuleLibFmp/DxeRuntimeCapsuleLib.inf |   4 +
 MdeModulePkg/MdeModulePkg.dec                 |   7 +-
 4 files changed, 192 insertions(+), 19 deletions(-)

diff --git a/MdeModulePkg/Library/DxeCapsuleLibFmp/DxeCapsuleLib.c b/MdeModulePkg/Library/DxeCapsuleLibFmp/DxeCapsuleLib.c
index 90942135d7..0000f91c6a 100644
--- a/MdeModulePkg/Library/DxeCapsuleLibFmp/DxeCapsuleLib.c
+++ b/MdeModulePkg/Library/DxeCapsuleLibFmp/DxeCapsuleLib.c
@@ -10,6 +10,7 @@
   ValidateFmpCapsule(), and DisplayCapsuleImage() receives untrusted input and
   performs basic validation.
 
+  Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.<BR>
   Copyright (c) 2016 - 2019, Intel Corporation. All rights reserved.<BR>
   SPDX-License-Identifier: BSD-2-Clause-Patent
 
@@ -41,6 +42,11 @@
 #include <Protocol/FirmwareManagementProgress.h>
 #include <Protocol/DevicePath.h>
 
+BOOLEAN (EFIAPI *mLibAtRuntimeFunction) (VOID)                  = NULL;
+EFI_FIRMWARE_MANAGEMENT_PROTOCOL    *mRuntimeFmp                = NULL;
+VOID                                **mRuntimeFmpProtocolArray  = NULL;
+EFI_GUID                            *mRuntimeFmpGuidArray       = NULL;
+
 EFI_SYSTEM_RESOURCE_TABLE *mEsrtTable                  = NULL;
 BOOLEAN                   mIsVirtualAddrConverted      = FALSE;
 
@@ -551,6 +557,11 @@ DumpAllFmpInfo (
   UINT32                                        PackageVersion;
   CHAR16                                        *PackageVersionName;
 
+  // Dump not supported at runtime.
+  if ((mLibAtRuntimeFunction != NULL) && mLibAtRuntimeFunction ()) {
+    return;
+  }
+
   Status = gBS->LocateHandleBuffer (
                   ByProtocol,
                   &gEfiFirmwareManagementProtocolGuid,
@@ -906,25 +917,35 @@ SetFmpImageData (
   CHAR16                                        *AbortReason;
   EFI_FIRMWARE_MANAGEMENT_UPDATE_IMAGE_PROGRESS ProgressCallback;
 
-  Status = gBS->HandleProtocol(
-                  Handle,
-                  &gEfiFirmwareManagementProtocolGuid,
-                  (VOID **)&Fmp
-                  );
-  if (EFI_ERROR(Status)) {
-    return Status;
-  }
+  // If not using optional runtime support, get FMP protocol for given Handle.
+  // Otherwise, use the one saved by ProcessFmpCapsuleImage().
+  if ((mLibAtRuntimeFunction == NULL) || !mLibAtRuntimeFunction ()) {
+    Status = gBS->HandleProtocol(
+                    Handle,
+                    &gEfiFirmwareManagementProtocolGuid,
+                    (VOID **)&Fmp
+                    );
+    if (EFI_ERROR(Status)) {
+      return Status;
+    }
 
-  //
-  // Lookup Firmware Management Progress Protocol before SetImage() is called
-  // This is an optional protocol that may not be present on Handle.
-  //
-  Status = gBS->HandleProtocol (
-                  Handle,
-                  &gEdkiiFirmwareManagementProgressProtocolGuid,
-                  (VOID **)&mFmpProgress
-                  );
-  if (EFI_ERROR (Status)) {
+    //
+    // Lookup Firmware Management Progress Protocol before SetImage() is called
+    // This is an optional protocol that may not be present on Handle.
+    //
+    Status = gBS->HandleProtocol (
+                    Handle,
+                    &gEdkiiFirmwareManagementProgressProtocolGuid,
+                    (VOID **)&mFmpProgress
+                    );
+    if (EFI_ERROR (Status)) {
+      mFmpProgress = NULL;
+    }
+  } else {
+    if (mRuntimeFmp == NULL) {
+      return EFI_UNSUPPORTED;
+    }
+    Fmp = mRuntimeFmp;
     mFmpProgress = NULL;
   }
 
@@ -1259,6 +1280,30 @@ ProcessFmpCapsuleImage (
       UpdateHardwareInstance = ImageHeader->UpdateHardwareInstance;
     }
 
+    // Optional runtime FMP SetImage processing sequence
+    if ((mLibAtRuntimeFunction != NULL) && mLibAtRuntimeFunction () &&
+        (mRuntimeFmpProtocolArray != NULL)) {
+      mRuntimeFmp = NULL;
+      Index2 = 0;
+      while (mRuntimeFmpProtocolArray[Index2] != NULL) {
+        if (CompareGuid (&ImageHeader->UpdateImageTypeId,
+                         &mRuntimeFmpGuidArray[Index2])) {
+          mRuntimeFmp = (EFI_FIRMWARE_MANAGEMENT_PROTOCOL *)
+            mRuntimeFmpProtocolArray[Index2];
+          break;
+        }
+        Index2++;
+      }
+
+      Status = SetFmpImageData (NULL,
+                                ImageHeader,
+                                Index - FmpCapsuleHeader->EmbeddedDriverCount);
+      if (EFI_ERROR (Status)) {
+        return Status;
+      }
+      continue;
+    }
+
     Status = GetFmpHandleBufferByType (
                &ImageHeader->UpdateImageTypeId,
                UpdateHardwareInstance,
diff --git a/MdeModulePkg/Library/DxeCapsuleLibFmp/DxeCapsuleRuntime.c b/MdeModulePkg/Library/DxeCapsuleLibFmp/DxeCapsuleRuntime.c
index f94044a409..6feb6dab79 100644
--- a/MdeModulePkg/Library/DxeCapsuleLibFmp/DxeCapsuleRuntime.c
+++ b/MdeModulePkg/Library/DxeCapsuleLibFmp/DxeCapsuleRuntime.c
@@ -1,6 +1,7 @@
 /** @file
   Capsule library runtime support.
 
+  Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.<BR>
   Copyright (c) 2016 - 2017, Intel Corporation. All rights reserved.<BR>
   SPDX-License-Identifier: BSD-2-Clause-Patent
 
@@ -19,7 +20,11 @@
 #include <Library/UefiBootServicesTableLib.h>
 #include <Library/UefiRuntimeServicesTableLib.h>
 #include <Library/MemoryAllocationLib.h>
+#include <Library/UefiRuntimeLib.h>
 
+extern BOOLEAN                   (EFIAPI *mLibAtRuntimeFunction) (VOID);
+extern VOID                      **mRuntimeFmpProtocolArray;
+extern EFI_GUID                  *mRuntimeFmpGuidArray;
 extern EFI_SYSTEM_RESOURCE_TABLE *mEsrtTable;
 extern BOOLEAN                   mIsVirtualAddrConverted;
 EFI_EVENT                 mDxeRuntimeCapsuleLibVirtualAddressChangeEvent  = NULL;
@@ -40,9 +45,121 @@ DxeCapsuleLibVirtualAddressChangeEvent (
   )
 {
   gRT->ConvertPointer (EFI_OPTIONAL_PTR, (VOID **)&mEsrtTable);
+
+  if (mRuntimeFmpProtocolArray != NULL) {
+    VOID **FmpArrayEntry;
+
+    FmpArrayEntry = mRuntimeFmpProtocolArray;
+    while (*FmpArrayEntry != NULL) {
+      EfiConvertPointer (0x0, (VOID **) FmpArrayEntry);
+      FmpArrayEntry++;
+    }
+    EfiConvertPointer (0x0, (VOID **) &mRuntimeFmpProtocolArray);
+  }
+  if (mRuntimeFmpGuidArray != NULL) {
+    EfiConvertPointer (0x0, (VOID **) &mRuntimeFmpGuidArray);
+  }
+  if (mLibAtRuntimeFunction != NULL ) {
+    EfiConvertPointer (0x0, (VOID **) &mLibAtRuntimeFunction);
+  }
+
   mIsVirtualAddrConverted = TRUE;
 }
 
+/**
+  Initialize optional runtime FMP arrays to support FMP SetImage processing
+  after ExitBootServices() is called.
+
+  The ImageTypeIdGuids of runtime-capable FMP protocol drivers are extracted
+  from the PcdRuntimeFmpCapsuleImageTypeIdGuid list and their protocol
+  structure pointers are saved in the mRuntimeFmpProtocolArray for use during
+  UpdateCapsule() processing. UpdateHardwareInstance is not supported.
+
+**/
+STATIC
+VOID
+EFIAPI
+InitializeRuntimeFmpArrays (
+  VOID
+  )
+{
+  EFI_GUID      *Guid;
+  UINTN         NumHandles;
+  EFI_HANDLE    *HandleBuffer;
+  EFI_STATUS    Status;
+  UINTN         Count;
+  UINTN         Index;
+  UINTN         FmpArrayIndex;
+
+  EFI_STATUS
+    GetFmpHandleBufferByType (
+      IN     EFI_GUID                     *UpdateImageTypeId,
+      IN     UINT64                       UpdateHardwareInstance,
+      OUT    UINTN                        *NoHandles, OPTIONAL
+      OUT    EFI_HANDLE                   **HandleBuf, OPTIONAL
+      OUT    BOOLEAN                      **ResetRequiredBuf OPTIONAL
+      );
+
+  Count = PcdGetSize (PcdRuntimeFmpCapsuleImageTypeIdGuid) / sizeof (GUID);
+  if (Count == 0) {
+    return;
+  }
+
+  // mRuntimeFmpProtocolArray is a NULL-terminated list of FMP protocol pointers
+  mRuntimeFmpProtocolArray = (VOID **)
+    AllocateRuntimeZeroPool ((Count + 1) * sizeof (VOID *));
+  if (mRuntimeFmpProtocolArray == NULL) {
+    DEBUG ((DEBUG_ERROR, "Error allocating mRuntimeFmpProtocolArray\n"));
+    return;
+  }
+  mRuntimeFmpGuidArray = (EFI_GUID *)
+    AllocateRuntimeZeroPool (Count * sizeof (EFI_GUID));
+  if (mRuntimeFmpGuidArray == NULL) {
+    DEBUG ((DEBUG_ERROR, "Error allocating mRuntimeFmpGuidArray"));
+    FreePool (mRuntimeFmpProtocolArray);
+    return;
+  }
+
+  // For each runtime ImageTypeIdGuid in the PCD, save its GUID and FMP protocol
+  FmpArrayIndex = 0;
+  Guid  = PcdGetPtr (PcdRuntimeFmpCapsuleImageTypeIdGuid);
+  for (Index = 0; Index < Count; Index++, Guid++) {
+    mRuntimeFmpGuidArray[FmpArrayIndex] = *Guid;
+    HandleBuffer = NULL;
+    Status = GetFmpHandleBufferByType (Guid,
+                                       0,
+                                       &NumHandles,
+                                       &HandleBuffer,
+                                       NULL);
+    if (EFI_ERROR (Status)) {
+      DEBUG ((DEBUG_ERROR,
+              "Error finding FMP handle for runtime ImageTypeIdGuid=%g: %r\n",
+              Guid, Status));
+      continue;
+    }
+
+    if (NumHandles > 1) {
+      DEBUG ((DEBUG_ERROR,
+              "FMP runtime ImageTypeIdGuid=%g returned %u handles, only 1 supported\n",
+              Guid, NumHandles));
+    }
+    Status = gBS->HandleProtocol (HandleBuffer[0],
+                                  &gEfiFirmwareManagementProtocolGuid,
+                                  &mRuntimeFmpProtocolArray[FmpArrayIndex]);
+    FreePool (HandleBuffer);
+    if (EFI_ERROR(Status)) {
+      DEBUG ((DEBUG_ERROR,
+              "Error getting FMP protocol for runtime ImageTypeIdGuid=%g: %r\n",
+              Guid, Status));
+      continue;
+    }
+
+    FmpArrayIndex++;
+  }
+
+  mLibAtRuntimeFunction = EfiAtRuntime;
+}
+
 /**
   Notify function for event group EFI_EVENT_GROUP_READY_TO_BOOT.
 
@@ -93,6 +210,8 @@ DxeCapsuleLibReadyToBootEventNotify (
     //
     mEsrtTable->FwResourceCountMax = mEsrtTable->FwResourceCount;
   }
+
+  InitializeRuntimeFmpArrays ();
 }
 
 /**
diff --git a/MdeModulePkg/Library/DxeCapsuleLibFmp/DxeRuntimeCapsuleLib.inf b/MdeModulePkg/Library/DxeCapsuleLibFmp/DxeRuntimeCapsuleLib.inf
index bf56f4623f..7b3f5e04f8 100644
--- a/MdeModulePkg/Library/DxeCapsuleLibFmp/DxeRuntimeCapsuleLib.inf
+++ b/MdeModulePkg/Library/DxeCapsuleLibFmp/DxeRuntimeCapsuleLib.inf
@@ -49,6 +49,7 @@
   PrintLib
   HobLib
   BmpSupportLib
+  PcdLib
 
 
 [Protocols]
@@ -70,5 +71,8 @@
   gEfiEventVirtualAddressChangeGuid       ## CONSUMES ## Event
   gEdkiiCapsuleOnDiskNameGuid             ## SOMETIMES_CONSUMES ## GUID
 
+[Pcd]
+  gEfiMdeModulePkgTokenSpaceGuid.PcdRuntimeFmpCapsuleImageTypeIdGuid
+
 [Depex]
   gEfiVariableWriteArchProtocolGuid
diff --git a/MdeModulePkg/MdeModulePkg.dec b/MdeModulePkg/MdeModulePkg.dec
index 133e04ee86..869aa892f7 100644
--- a/MdeModulePkg/MdeModulePkg.dec
+++ b/MdeModulePkg/MdeModulePkg.dec
@@ -3,7 +3,7 @@
 # It also provides the definitions(including PPIs/PROTOCOLs/GUIDs and library classes)
 # and libraries instances, which are used for those modules.
 #
-# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2019-2021, NVIDIA CORPORATION. All rights reserved.<BR>
 # Copyright (c) 2007 - 2021, Intel Corporation. All rights reserved.<BR>
 # Copyright (c) 2016, Linaro Ltd. All rights reserved.<BR>
 # (C) Copyright 2016 - 2019 Hewlett Packard Enterprise Development LP<BR>
@@ -2020,6 +2020,11 @@
   # @Prompt Capsule On Disk Temp Relocation file name in PEI phase
   gEfiMdeModulePkgTokenSpaceGuid.PcdCoDRelocationFileName|L"Cod.tmp"|VOID*|0x30001048
 
+  ## This PCD holds a list of GUIDs for the ImageTypeId to indicate the
+  #  FMP is runtime capable.
+  # @Prompt A list of runtime-capable FMP ImageTypeId GUIDs
+  gEfiMdeModulePkgTokenSpaceGuid.PcdRuntimeFmpCapsuleImageTypeIdGuid|{0x0}|VOID*|0x30001049
+
   ## This PCD hold a list GUIDs for the ImageTypeId to indicate the
   #  FMP capsule is a system FMP.
   # @Prompt A list of system FMP ImageTypeId GUIDs
-- 
2.17.1