[PATCH] MdeModulePkg PiSmmCoreMemoryAllocLib: Fix a FreePool() assertion issue

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

When PiSmmCore links against PeiDxeDebugLibReportStatusCode, the code
flow below will cause a FreePool() assertion issue.

PiSmmCoreMemoryAllocationLibConstructor() ->
SmmInitializeMemoryServices() ->
DEBUG ((DEBUG_INFO, "SmmAddMemoryRegion\n")) in SmmAddMemoryRegion() ->
DebugPrint() -> REPORT_STATUS_CODE_EX() -> ReportStatusCodeEx() ->
AllocatePool()/FreePool(PiSmmCoreMemoryAllocLib) ->
ASSERT() at Head = CR (Buffer, POOL_HEAD, Data, POOL_HEAD_SIGNATURE)
  in CoreFreePoolI() of DxeCore Pool.c

It is because at the point of FreePool() in the code flow above,
mSmmCoreMemoryAllocLibSmramRanges/mSmmCoreMemoryAllocLibSmramRangeCount
are not been initialized yet, the FreePool() will be directed to
gBS->FreePool(), that is wrong.

This patch is to temporarily use BootServicesData to hold the
SmramRanges data before calling SmmInitializeMemoryServices().

Cc: Liming Gao <liming.gao@intel.com>
Cc: Jiewen Yao <jiewen.yao@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Star Zeng <star.zeng@intel.com>
---
 .../MemoryAllocationLib.c                          | 32 +++++++++++++++++++---
 1 file changed, 28 insertions(+), 4 deletions(-)

diff --git a/MdeModulePkg/Library/PiSmmCoreMemoryAllocationLib/MemoryAllocationLib.c b/MdeModulePkg/Library/PiSmmCoreMemoryAllocationLib/MemoryAllocationLib.c
index 96cb275cc9d7..4216a12d18f5 100644
--- a/MdeModulePkg/Library/PiSmmCoreMemoryAllocationLib/MemoryAllocationLib.c
+++ b/MdeModulePkg/Library/PiSmmCoreMemoryAllocationLib/MemoryAllocationLib.c
@@ -1068,20 +1068,44 @@ PiSmmCoreMemoryAllocationLibConstructor (
   IN EFI_SYSTEM_TABLE  *SystemTable
   )
 {
+  EFI_STATUS             Status;
   SMM_CORE_PRIVATE_DATA  *SmmCorePrivate;
   UINTN                  Size;
+  VOID                   *BootServicesData;
 
   SmmCorePrivate = (SMM_CORE_PRIVATE_DATA *)ImageHandle;
+
   //
-  // Initialize memory service using free SMRAM
+  // The FreePool()/FreePages() will need use SmramRanges data to know whether
+  // the buffer to free is in SMRAM range or not. And there may be FreePool()/
+  // FreePages() indrectly during calling SmmInitializeMemoryServices(), but
+  // no SMRAM could be allocated before calling SmmInitializeMemoryServices(),
+  // so temporarily use BootServicesData to hold the SmramRanges data.
   //
-  SmmInitializeMemoryServices (SmmCorePrivate->SmramRangeCount, SmmCorePrivate->SmramRanges);
-
   mSmmCoreMemoryAllocLibSmramRangeCount = SmmCorePrivate->SmramRangeCount;
   Size = mSmmCoreMemoryAllocLibSmramRangeCount * sizeof (EFI_SMRAM_DESCRIPTOR);
-  mSmmCoreMemoryAllocLibSmramRanges = (EFI_SMRAM_DESCRIPTOR *) AllocatePool (Size);
+  Status = gBS->AllocatePool (EfiBootServicesData, Size, (VOID **) &mSmmCoreMemoryAllocLibSmramRanges);
+  ASSERT_EFI_ERROR (Status);
   ASSERT (mSmmCoreMemoryAllocLibSmramRanges != NULL);
   CopyMem (mSmmCoreMemoryAllocLibSmramRanges, SmmCorePrivate->SmramRanges, Size);
 
+  //
+  // Initialize memory service using free SMRAM
+  //
+  SmmInitializeMemoryServices (SmmCorePrivate->SmramRangeCount, SmmCorePrivate->SmramRanges);
+
+  //
+  // Move the SmramRanges data from BootServicesData to SMRAM.
+  //
+  BootServicesData = mSmmCoreMemoryAllocLibSmramRanges;
+  mSmmCoreMemoryAllocLibSmramRanges = (EFI_SMRAM_DESCRIPTOR *) AllocateCopyPool (Size, (VOID *) BootServicesData);
+  ASSERT (mSmmCoreMemoryAllocLibSmramRanges != NULL);
+
+  //
+  // Free the temporarily used BootServicesData.
+  //
+  Status = gBS->FreePool (BootServicesData);
+  ASSERT_EFI_ERROR (Status);
+
   return EFI_SUCCESS;
 }
-- 
2.7.0.windows.1