[PATCH 2/2] MdeModulePkg ErstFmpDxe: Create ESRT in ReadyToBoot event

[Star Zeng <star.zeng@intel.com>]

Current code just creates ESRT entry in FMP notification
and installs ESRT configuration table in ReadyToBoot event.

The LastAttemptVersion and LastAttemptStatus in ESRT will
be out of date after system continues to boot without reset
after capsule update (reset is required or capsule update
is failed).

This patches updates the code to create ESRT based on all
FMP instances in ReadyToBoot event.

Cc: Michael D Kinney <michael.d.kinney@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.1
Signed-off-by: Star Zeng <star.zeng@intel.com>
---
 MdeModulePkg/Universal/EsrtFmpDxe/EsrtFmp.c | 239 ++++++++++++----------------
 1 file changed, 103 insertions(+), 136 deletions(-)

diff --git a/MdeModulePkg/Universal/EsrtFmpDxe/EsrtFmp.c b/MdeModulePkg/Universal/EsrtFmpDxe/EsrtFmp.c
index cc9dd63b53cd..219059760e97 100644
--- a/MdeModulePkg/Universal/EsrtFmpDxe/EsrtFmp.c
+++ b/MdeModulePkg/Universal/EsrtFmpDxe/EsrtFmp.c
@@ -55,47 +55,34 @@ PrintTable (
 //
 #define GROWTH_STEP  10
 
-//
-// Module globals.
-//
-EFI_EVENT                  mEsrtReadyToBootEvent;
-EFI_SYSTEM_RESOURCE_TABLE  *mTable = NULL;
-BOOLEAN                    mEsrtInstalled = FALSE;
-EFI_EVENT                  mFmpInstallEvent;
-VOID                       *mFmpInstallEventRegistration = NULL;
-
 /**
   Install EFI System Resource Table into the UEFI Configuration Table
 
+  @param[in] Table                  Pointer to the ESRT.
+
   @return  Status code.
 
 **/
 EFI_STATUS
 InstallEfiSystemResourceTableInUefiConfigurationTable (
-   VOID
+  IN EFI_SYSTEM_RESOURCE_TABLE      *Table
   )
 {
   EFI_STATUS Status;
 
   Status = EFI_SUCCESS;
-  if (!mEsrtInstalled) {
-    if (mTable == NULL) {
-      DEBUG ((DEBUG_ERROR, "EsrtFmpDxe: Can't install ESRT table because it is NULL. \n"));
-      Status = EFI_OUT_OF_RESOURCES;
-    } else if (mTable->FwResourceCount == 0) {
-      DEBUG ((DEBUG_ERROR, "EsrtFmpDxe: Can't install ESRT table because it has zero Entries. \n"));
-      Status = EFI_UNSUPPORTED;
+  if (Table->FwResourceCount == 0) {
+    DEBUG ((DEBUG_ERROR, "EsrtFmpDxe: Can't install ESRT table because it has zero Entries. \n"));
+    Status = EFI_UNSUPPORTED;
+  } else {
+    //
+    // Install the pointer into config table
+    //
+    Status = gBS->InstallConfigurationTable (&gEfiSystemResourceTableGuid, Table);
+    if (EFI_ERROR (Status)) {
+      DEBUG ((DEBUG_ERROR, "EsrtFmpDxe: Can't install ESRT table.  Status: %r. \n", Status));
     } else {
-      //
-      // Install the pointer into config table
-      //
-      Status = gBS->InstallConfigurationTable (&gEfiSystemResourceTableGuid, mTable);
-      if (EFI_ERROR (Status)) {
-        DEBUG ((DEBUG_ERROR, "EsrtFmpDxe: Can't install ESRT table.  Status: %r. \n", Status));
-      } else {
-        DEBUG ((DEBUG_INFO, "EsrtFmpDxe: Installed ESRT table. \n"));
-        mEsrtInstalled = TRUE;
-      }
+      DEBUG ((DEBUG_INFO, "EsrtFmpDxe: Installed ESRT table. \n"));
     }
   }
   return Status;
@@ -135,15 +122,18 @@ IsSystemFmp (
   given a FMP descriptor.  If the guid is already in the ESRT it
   will be ignored.  The ESRT will grow if it does not have enough room.
 
-  @param[in]  FmpImageInfoBuf    Pointer to the EFI_FIRMWARE_IMAGE_DESCRIPTOR.
-  @param[in]  FmpVersion         FMP Version.
+  @param[in, out] Table             On input, pointer to the pointer to the ESRT.
+                                    On output, same as input or pointer to the pointer
+                                    to new enlarged ESRT.
+  @param[in]      FmpImageInfoBuf   Pointer to the EFI_FIRMWARE_IMAGE_DESCRIPTOR.
+  @param[in]      FmpVersion        FMP Version.
 
   @return  Status code.
 
 **/
 EFI_STATUS
-EFIAPI
 CreateEsrtEntry (
+  IN OUT EFI_SYSTEM_RESOURCE_TABLE  **Table,
   IN EFI_FIRMWARE_IMAGE_DESCRIPTOR  *FmpImageInfoBuf,
   IN UINT32                         FmpVersion
   )
@@ -156,18 +146,11 @@ CreateEsrtEntry (
   Index = 0;
   Entry = NULL;
 
-  //
-  // Get our ESRT table.  This should never be null at this point
-  //
-  if (mTable == NULL) {
-    return EFI_DEVICE_ERROR;
-  }
-
-  Entry = (EFI_SYSTEM_RESOURCE_ENTRY *)(mTable + 1);
+  Entry = (EFI_SYSTEM_RESOURCE_ENTRY *)((*Table) + 1);
   //
   // Make sure Guid isn't already in the list
   //
-  for (Index = 0; Index < mTable->FwResourceCount; Index++) {
+  for (Index = 0; Index < (*Table)->FwResourceCount; Index++) {
     if (CompareGuid (&Entry->FwClass, &FmpImageInfoBuf->ImageTypeId)) {
       DEBUG ((DEBUG_ERROR, "EsrtFmpDxe: ESRT Entry already exists for FMP Instance with GUID %g\n", &Entry->FwClass));
       return EFI_INVALID_PARAMETER;
@@ -178,18 +161,8 @@ CreateEsrtEntry (
   //
   // Grow table if needed
   //
-  if (mTable->FwResourceCount >= mTable->FwResourceCountMax) {
-    //
-    // Can't grow table after installed.
-    // Only because didn't add support for this.
-    // Would need to re-install ESRT in system table if wanted to support
-    //
-    if (mEsrtInstalled) {
-      DEBUG ((DEBUG_ERROR, "EsrtFmpDxe: Failed to install entry because ESRT table needed to grow after table already installed. \n"));
-      return EFI_OUT_OF_RESOURCES;
-    }
-
-    NewSize  = ((mTable->FwResourceCountMax + GROWTH_STEP) * sizeof (EFI_SYSTEM_RESOURCE_ENTRY)) + sizeof (EFI_SYSTEM_RESOURCE_TABLE);
+  if ((*Table)->FwResourceCount >= (*Table)->FwResourceCountMax) {
+    NewSize  = (((*Table)->FwResourceCountMax + GROWTH_STEP) * sizeof (EFI_SYSTEM_RESOURCE_ENTRY)) + sizeof (EFI_SYSTEM_RESOURCE_TABLE);
     NewTable = AllocateZeroPool (NewSize);
     if (NewTable == NULL) {
       DEBUG ((DEBUG_ERROR, "EsrtFmpDxe: Failed to allocate memory larger table for ESRT. \n"));
@@ -200,8 +173,8 @@ CreateEsrtEntry (
     //
     CopyMem (
       NewTable,
-      mTable,
-      ((mTable->FwResourceCountMax) * sizeof (EFI_SYSTEM_RESOURCE_ENTRY)) + sizeof (EFI_SYSTEM_RESOURCE_TABLE)
+      (*Table),
+      (((*Table)->FwResourceCountMax) * sizeof (EFI_SYSTEM_RESOURCE_ENTRY)) + sizeof (EFI_SYSTEM_RESOURCE_TABLE)
       );
     //
     // Update max
@@ -210,25 +183,25 @@ CreateEsrtEntry (
     //
     // Free old table
     //
-    FreePool (mTable);
+    FreePool ((*Table));
     //
     // Reassign pointer to new table.
     //
-    mTable = NewTable;
+    (*Table) = NewTable;
   }
 
   //
   // ESRT table has enough room for the new entry so add new entry
   //
-  Entry = (EFI_SYSTEM_RESOURCE_ENTRY *)(((UINT8 *)mTable) + sizeof (EFI_SYSTEM_RESOURCE_TABLE));
+  Entry = (EFI_SYSTEM_RESOURCE_ENTRY *)(((UINT8 *)(*Table)) + sizeof (EFI_SYSTEM_RESOURCE_TABLE));
   //
   // Move to the location of new entry
   //
-  Entry = Entry + mTable->FwResourceCount;
+  Entry = Entry + (*Table)->FwResourceCount;
   //
   // Increment resource count
   //
-  mTable->FwResourceCount++;
+  (*Table)->FwResourceCount++;
 
   CopyGuid (&Entry->FwClass, &FmpImageInfoBuf->ImageTypeId);
 
@@ -264,23 +237,23 @@ CreateEsrtEntry (
 }
 
 /**
-  Notify function for every Firmware Management Protocol being installed.
-  Get the descriptors from FMP Instance and create ESRT entries (ESRE)
+  Function to create ESRT based on FMP Instances.
+  Create ESRT table, get the descriptors from FMP Instance and
+  create ESRT entries (ESRE).
 
-  @param[in]  Event    The Event that is being processed.
-  @param[in]  Context  The Event Context.
+  @return Pointer to the ESRT created.
 
 **/
-VOID
-EFIAPI
-FmpInstallProtocolNotify (
-  IN EFI_EVENT  Event,
-  IN VOID       *Context
+EFI_SYSTEM_RESOURCE_TABLE *
+CreateFmpBasedEsrt (
+  VOID
   )
 {
   EFI_STATUS                        Status;
-  EFI_HANDLE                        Handle;
-  UINTN                             BufferSize;
+  EFI_SYSTEM_RESOURCE_TABLE         *Table;
+  UINTN                             NoProtocols;
+  VOID                              **Buffer;
+  UINTN                             Index;
   EFI_FIRMWARE_MANAGEMENT_PROTOCOL  *Fmp;
   UINTN                             DescriptorSize;
   EFI_FIRMWARE_IMAGE_DESCRIPTOR     *FmpImageInfoBuf;
@@ -292,29 +265,43 @@ FmpInstallProtocolNotify (
   CHAR16                            *PackageVersionName;
 
   Status             = EFI_SUCCESS;
-  Handle             = 0;
-  BufferSize         = 0;
+  Table              = NULL;
+  NoProtocols        = 0;
+  Buffer             = NULL;
   PackageVersionName = NULL;
   FmpImageInfoBuf    = NULL;
   FmpImageInfoBufOrg = NULL;
   Fmp                = NULL;
 
-  DEBUG ((DEBUG_INFO, "FMP Installed Notify\n"));
-  while (TRUE) {
-    BufferSize = sizeof (EFI_HANDLE);
-    Status = gBS->LocateHandle (ByRegisterNotify, NULL, mFmpInstallEventRegistration, &BufferSize, &Handle);
-    if (EFI_ERROR (Status)) {
-      DEBUG ((DEBUG_WARN, "EsrtFmpDxe: Failed to Locate handle from notify value. Status: %r\n", Status));
-      return;
-    }
+  Status = EfiLocateProtocolBuffer (
+             &gEfiFirmwareManagementProtocolGuid,
+             &NoProtocols,
+             &Buffer
+             );
+  if (EFI_ERROR(Status) || (Buffer == NULL)) {
+    return NULL;
+  }
 
-    Status = gBS->HandleProtocol (Handle, &gEfiFirmwareManagementProtocolGuid, (VOID **)&Fmp);
-    if (EFI_ERROR (Status)) {
-      DEBUG ((DEBUG_ERROR, "EsrtFmpDxe: Failed to get FMP for a handle 0x%x\n", Handle));
-      continue;
-    }
-    ImageInfoSize = 0;
+  //
+  // Allocate Memory for table
+  //
+  Table = AllocateZeroPool (
+             (GROWTH_STEP * sizeof (EFI_SYSTEM_RESOURCE_ENTRY)) + sizeof (EFI_SYSTEM_RESOURCE_TABLE)
+             );
+  if (Table == NULL) {
+    DEBUG ((DEBUG_ERROR, "EsrtFmpDxe: Failed to allocate memory for ESRT.\n"));
+    gBS->FreePool (Buffer);
+    return NULL;
+  }
+
+  Table->FwResourceCount    = 0;
+  Table->FwResourceCountMax = GROWTH_STEP;
+  Table->FwResourceVersion  = EFI_SYSTEM_RESOURCE_TABLE_FIRMWARE_RESOURCE_VERSION; 
 
+  for (Index = 0; Index < NoProtocols; Index++) {
+    Fmp = (EFI_FIRMWARE_MANAGEMENT_PROTOCOL *) Buffer[Index];
+
+    ImageInfoSize = 0;
     Status = Fmp->GetImageInfo (
                     Fmp,                         // FMP Pointer
                     &ImageInfoSize,              // Buffer Size (in this case 0)
@@ -331,7 +318,6 @@ FmpInstallProtocolNotify (
       continue;
     }
 
-    FmpImageInfoBuf = NULL;
     FmpImageInfoBuf = AllocateZeroPool (ImageInfoSize);
     if (FmpImageInfoBuf == NULL) {
       DEBUG ((DEBUG_ERROR, "EsrtFmpDxe: Failed to get memory for descriptors.\n"));
@@ -352,7 +338,9 @@ FmpInstallProtocolNotify (
                     );
     if (EFI_ERROR (Status)) {
       DEBUG ((DEBUG_ERROR, "EsrtFmpDxe: Failure in GetImageInfo.  Status = %r\n", Status));
-      goto CleanUp;
+      FreePool (FmpImageInfoBufOrg);
+      FmpImageInfoBufOrg = NULL;
+      continue;
     }
 
     //
@@ -366,7 +354,7 @@ FmpInstallProtocolNotify (
         //
         // Create ESRT entry
         //
-        CreateEsrtEntry (FmpImageInfoBuf, FmpImageInfoDescriptorVer);
+        CreateEsrtEntry (&Table, FmpImageInfoBuf, FmpImageInfoDescriptorVer);
       }
       FmpImageInfoCount--;
       //
@@ -379,17 +367,12 @@ FmpInstallProtocolNotify (
       FreePool (PackageVersionName);
       PackageVersionName = NULL;
     }
-    if (FmpImageInfoBufOrg != NULL) {
-      FreePool (FmpImageInfoBufOrg);
-      FmpImageInfoBufOrg = NULL;
-    }
-  }
-
-CleanUp:
-  if (FmpImageInfoBufOrg != NULL) {
     FreePool (FmpImageInfoBufOrg);
+    FmpImageInfoBufOrg = NULL;
   }
-  return;
+
+  gBS->FreePool (Buffer);
+  return Table;
 }
 
 /**
@@ -407,14 +390,30 @@ EsrtReadyToBootEventNotify (
   IN VOID       *Context
   )
 {
-  InstallEfiSystemResourceTableInUefiConfigurationTable ();
+  EFI_STATUS                 Status;
+  EFI_SYSTEM_RESOURCE_TABLE  *Table;
+
+  Table = CreateFmpBasedEsrt ();
+  if (Table != NULL) {
+    //
+    // Print table on debug builds
+    //
+    DEBUG_CODE_BEGIN ();
+    PrintTable (Table);
+    DEBUG_CODE_END ();
+
+    Status = InstallEfiSystemResourceTableInUefiConfigurationTable (Table);
+    if (EFI_ERROR (Status)) {
+      FreePool (Table);
+    }
+  } else {
+    DEBUG ((DEBUG_ERROR, "EsrtFmpDxe: Can't install ESRT table because it is NULL. \n"));
+  }
 
   //
-  // Print table on debug builds
+  // Close the event to prevent it be signalled again.
   //
-  DEBUG_CODE_BEGIN ();
-  PrintTable (mTable);
-  DEBUG_CODE_END ();
+  gBS->CloseEvent (Event);
 }
 
 /**
@@ -435,39 +434,7 @@ EsrtFmpEntryPoint (
   )
 {
   EFI_STATUS  Status;
-
-  //
-  // Allocate Memory for table
-  //
-  mTable = AllocateZeroPool (
-             (GROWTH_STEP * sizeof (EFI_SYSTEM_RESOURCE_ENTRY)) + sizeof (EFI_SYSTEM_RESOURCE_TABLE)
-             );
-  ASSERT (mTable != NULL);
-  if (mTable == NULL) {
-    DEBUG ((DEBUG_ERROR, "EsrtFmpDxe: Failed to allocate memory for ESRT.\n"));
-    return EFI_OUT_OF_RESOURCES;
-  }
-
-  mTable->FwResourceCount    = 0;
-  mTable->FwResourceCountMax = GROWTH_STEP;
-  mTable->FwResourceVersion  = EFI_SYSTEM_RESOURCE_TABLE_FIRMWARE_RESOURCE_VERSION;
-
-  //
-  // Register notify function for all FMP installed
-  //
-  mFmpInstallEvent = EfiCreateProtocolNotifyEvent (
-                       &gEfiFirmwareManagementProtocolGuid,
-                       TPL_CALLBACK,
-                       FmpInstallProtocolNotify,
-                       NULL,
-                       &mFmpInstallEventRegistration
-                       );
-
-  ASSERT (mFmpInstallEvent != NULL);
-
-  if (mFmpInstallEvent == NULL) {
-    DEBUG ((DEBUG_ERROR, "EsrtFmpDxe: Failed to Create Protocol Notify Event for FMP.\n"));
-  }
+  EFI_EVENT   EsrtReadyToBootEvent;
 
   //
   // Register notify function to install ESRT on ReadyToBoot Event.
@@ -476,7 +443,7 @@ EsrtFmpEntryPoint (
              TPL_CALLBACK,
              EsrtReadyToBootEventNotify,
              NULL,
-             &mEsrtReadyToBootEvent
+             &EsrtReadyToBootEvent
              );
 
   ASSERT_EFI_ERROR (Status);
-- 
2.7.0.windows.1