MdeModulePkg DxeCore: Enhance MemoryAttributesTable installation

Current MemoryAttributesTable will be installed on ReadyToBoot event
at TPL_NOTIFY level, it maybe incorrect when PcdHiiOsRuntimeSupport
= TRUE as HiiDatabaseDxe will have runtime memory allocation for HII
OS runtime support on and after ReadyToBoot. The issue was exposed at
http://article.gmane.org/gmane.comp.bios.edk2.devel/10125.

To make sure the correctness of MemoryAttributesTable, this patch is
to enhance MemoryAttributesTable installation to install
MemoryAttributesTable on ReadyToBoot event at TPL_CALLBACK - 1 level
to make sure it is at the last of ReadyToBoot event, and also hook
runtime memory allocation after ReadyToBoot.

Cc: Jiewen Yao <jiewen.yao@intel.com>
Cc: Liming Gao <liming.gao@intel.com>
Cc: Feng Tian <feng.tian@intel.com>
Cc: Ard Biesheuvel <ard.biesheuvel@linaro.org>
Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Star Zeng <star.zeng@intel.com>
Reviewed-by: Jiewen Yao <jiewen.yao@intel.com>

MdeModulePkg/Core/Dxe/DxeMain.h

@@ -276,6 +276,8 @@ extern EFI_RUNTIME_SERVICES                     *gDxeCoreRT;
 extern EFI_DXE_SERVICES                         *gDxeCoreDS;
 extern EFI_HANDLE                               gDxeCoreImageHandle;
 
+extern BOOLEAN                                  gMemoryMapTerminated;
+
 extern EFI_DECOMPRESS_PROTOCOL                  gEfiDecompress;
 
 extern EFI_RUNTIME_ARCH_PROTOCOL                *gRuntime;
@@ -2858,6 +2860,16 @@ CoreInitializeMemoryAttributesTable (
   VOID
   );
 
+/**
+  Install MemoryAttributesTable on memory allocation.
+
+  @param[in] MemoryType EFI memory type.
+**/
+VOID
+InstallMemoryAttributesTableOnMemoryAllocation (
+  IN EFI_MEMORY_TYPE    MemoryType
+  );
+
 /**
   Insert image record.
 

MdeModulePkg/Core/Dxe/DxeMain/DxeMain.c

@@ -207,6 +207,7 @@ EFI_SYSTEM_TABLE      *gDxeCoreST = NULL;
 EFI_RUNTIME_SERVICES  *gDxeCoreRT = &mEfiRuntimeServicesTableTemplate;
 EFI_HANDLE            gDxeCoreImageHandle = NULL;
 
+BOOLEAN               gMemoryMapTerminated = FALSE;
 
 //
 // EFI Decompress Protocol
@@ -751,6 +752,8 @@ CoreExitBootServices (
     return Status;
   }
 
+  gMemoryMapTerminated = TRUE;
+
   //
   // Notify other drivers that we are exiting boot services.
   //

MdeModulePkg/Core/Dxe/Mem/Page.c

@@ -1336,6 +1336,7 @@ CoreAllocatePages (
   Status = CoreInternalAllocatePages (Type, MemoryType, NumberOfPages, Memory);
   if (!EFI_ERROR (Status)) {
     CoreUpdateProfile ((EFI_PHYSICAL_ADDRESS) (UINTN) RETURN_ADDRESS (0), MemoryProfileActionAllocatePages, MemoryType, EFI_PAGES_TO_SIZE (NumberOfPages), (VOID *) (UINTN) *Memory);
+    InstallMemoryAttributesTableOnMemoryAllocation (MemoryType);
   }
   return Status;
 }
@@ -1444,6 +1445,7 @@ CoreFreePages (
   Status = CoreInternalFreePages (Memory, NumberOfPages, &MemoryType);
   if (!EFI_ERROR (Status)) {
     CoreUpdateProfile ((EFI_PHYSICAL_ADDRESS) (UINTN) RETURN_ADDRESS (0), MemoryProfileActionFreePages, (EFI_MEMORY_TYPE) 0, EFI_PAGES_TO_SIZE (NumberOfPages), (VOID *) (UINTN) Memory);
+    InstallMemoryAttributesTableOnMemoryAllocation (MemoryType);
   }
   return Status;
 }

MdeModulePkg/Core/Dxe/Mem/Pool.c

@@ -277,6 +277,7 @@ CoreAllocatePool (
   Status = CoreInternalAllocatePool (PoolType, Size, Buffer);
   if (!EFI_ERROR (Status)) {
     CoreUpdateProfile ((EFI_PHYSICAL_ADDRESS) (UINTN) RETURN_ADDRESS (0), MemoryProfileActionAllocatePool, PoolType, Size, *Buffer);
+    InstallMemoryAttributesTableOnMemoryAllocation (PoolType);
   }
   return Status;
 }
@@ -505,6 +506,7 @@ CoreFreePool (
   Status = CoreInternalFreePool (Buffer, &PoolType);
   if (!EFI_ERROR (Status)) {
     CoreUpdateProfile ((EFI_PHYSICAL_ADDRESS) (UINTN) RETURN_ADDRESS (0), MemoryProfileActionFreePool, (EFI_MEMORY_TYPE) 0, 0, Buffer);
+    InstallMemoryAttributesTableOnMemoryAllocation (PoolType);
   }
   return Status;
 }

MdeModulePkg/Core/Dxe/Misc/MemoryAttributesTable.c

@@ -71,18 +71,16 @@ CoreGetMemoryMapPropertiesTable (
   );
 
 extern EFI_PROPERTIES_TABLE  mPropertiesTable;
+EFI_MEMORY_ATTRIBUTES_TABLE  *mMemoryAttributesTable = NULL;
+BOOLEAN                      mMemoryAttributesTableReadyToBoot = FALSE;
 
 /**
   Install MemoryAttributesTable.
 
-  @param[in]  Event     The Event this notify function registered to.
-  @param[in]  Context   Pointer to the context data registered to the Event.
 **/
 VOID
-EFIAPI
 InstallMemoryAttributesTable (
-  EFI_EVENT                               Event,
+  VOID
-  VOID                                    *Context
   )
 {
   UINTN                          MemoryMapSize;
@@ -97,12 +95,29 @@ InstallMemoryAttributesTable (
   EFI_MEMORY_ATTRIBUTES_TABLE    *MemoryAttributesTable;
   EFI_MEMORY_DESCRIPTOR          *MemoryAttributesEntry;
 
+  if (gMemoryMapTerminated) {
+    //
+    // Directly return after MemoryMap terminated.
+    //
+    return;
+  }
+
   if ((mPropertiesTable.MemoryProtectionAttribute & EFI_PROPERTIES_RUNTIME_MEMORY_PROTECTION_NON_EXECUTABLE_PE_DATA) == 0) {
     DEBUG ((EFI_D_VERBOSE, "MemoryProtectionAttribute NON_EXECUTABLE_PE_DATA is not set, "));
     DEBUG ((EFI_D_VERBOSE, "because Runtime Driver Section Alignment is not %dK.\n", EFI_ACPI_RUNTIME_PAGE_ALLOCATION_ALIGNMENT >> 10));
     return ;
   }
 
+  if (mMemoryAttributesTable == NULL) {
+    //
+    // InstallConfigurationTable here to occupy one entry for MemoryAttributesTable
+    // before GetMemoryMap below, as InstallConfigurationTable may allocate runtime
+    // memory for the new entry.
+    //
+    Status = gBS->InstallConfigurationTable (&gEfiMemoryAttributesTableGuid, (VOID *) (UINTN) MAX_ADDRESS);
+    ASSERT_EFI_ERROR (Status);
+  }
+
   MemoryMapSize = 0;
   MemoryMap = NULL;
   Status = CoreGetMemoryMapPropertiesTable (
@@ -177,8 +192,52 @@ InstallMemoryAttributesTable (
   MemoryMap = MemoryMapStart;
   FreePool (MemoryMap);
 
+  //
+  // Update configuratoin table for MemoryAttributesTable.
+  //
   Status = gBS->InstallConfigurationTable (&gEfiMemoryAttributesTableGuid, MemoryAttributesTable);
   ASSERT_EFI_ERROR (Status);
+
+  if (mMemoryAttributesTable != NULL) {
+    FreePool (mMemoryAttributesTable);
+  }
+  mMemoryAttributesTable = MemoryAttributesTable; 
+}
+
+/**
+  Install MemoryAttributesTable on memory allocation.
+
+  @param[in] MemoryType EFI memory type.
+**/
+VOID
+InstallMemoryAttributesTableOnMemoryAllocation (
+  IN EFI_MEMORY_TYPE    MemoryType
+  )
+{
+  //
+  // Install MemoryAttributesTable after ReadyToBoot on runtime memory allocation.
+  //
+  if (mMemoryAttributesTableReadyToBoot &&
+      ((MemoryType == EfiRuntimeServicesCode) || (MemoryType == EfiRuntimeServicesData))) {
+    InstallMemoryAttributesTable ();
+  }
+}
+
+/**
+  Install MemoryAttributesTable on ReadyToBoot.
+
+  @param[in] Event      The Event this notify function registered to.
+  @param[in] Context    Pointer to the context data registered to the Event.
+**/
+VOID
+EFIAPI
+InstallMemoryAttributesTableOnReadyToBoot (
+  IN EFI_EVENT          Event,
+  IN VOID               *Context
+  )
+{
+  InstallMemoryAttributesTable ();
+  mMemoryAttributesTableReadyToBoot = TRUE; 
 }
 
 /**
@@ -194,17 +253,16 @@ CoreInitializeMemoryAttributesTable (
   EFI_EVENT   ReadyToBootEvent;
 
   //
-  // Construct the table at ReadyToBoot, because this should be
+  // Construct the table at ReadyToBoot.
-  // last point to allocate RuntimeCode/RuntimeData.
   //
-  Status = gBS->CreateEventEx (
+  Status = CoreCreateEventInternal (
-                  EVT_NOTIFY_SIGNAL,
+             EVT_NOTIFY_SIGNAL,
-                  TPL_NOTIFY,
+             TPL_CALLBACK - 1,
-                  InstallMemoryAttributesTable,
+             InstallMemoryAttributesTableOnReadyToBoot,
-                  NULL,
+             NULL,
-                  &gEfiEventReadyToBootGuid,
+             &gEfiEventReadyToBootGuid,
-                  &ReadyToBootEvent
+             &ReadyToBootEvent
-                  );
+             );
   ASSERT_EFI_ERROR (Status);
   return ;
 }