2 files changed, 125 insertions, 3 deletions
diff --git a/MdeModulePkg/Library/PiSmmCoreMemoryAllocationLib/MemoryAllocationLib.c b/MdeModulePkg/Library/PiSmmCoreMemoryAllocationLib/MemoryAllocationLib.c
index 7df6a1a73b..7435844fb9 100644
--- a/MdeModulePkg/Library/PiSmmCoreMemoryAllocationLib/MemoryAllocationLib.c
+++ b/MdeModulePkg/Library/PiSmmCoreMemoryAllocationLib/MemoryAllocationLib.c
@@ -14,11 +14,93 @@
 
 #include <PiSmm.h>
 
+#include <Protocol/SmmAccess2.h>
 #include <Library/MemoryAllocationLib.h>
+#include <Library/UefiBootServicesTableLib.h>
 #include <Library/BaseMemoryLib.h>
 #include <Library/DebugLib.h>
 #include "PiSmmCoreMemoryAllocationServices.h"
 
+EFI_SMRAM_DESCRIPTOR  *mSmramRanges    = NULL;
+UINTN                 mSmramRangeCount = 0;
+
+/**
+  This function gets and caches SMRAM ranges that are present in the system.
+    
+  It will ASSERT() if SMM Access2 Protocol doesn't exist.
+  It will ASSERT() if SMRAM ranges can't be got.
+  It will ASSERT() if Resource can't be allocated for cache SMRAM range. 
+
+**/
+VOID
+EFIAPI
+GetSmramRanges (
+  VOID
+  )
+{
+  EFI_STATUS                Status;
+  EFI_SMM_ACCESS2_PROTOCOL  *SmmAccess;
+  UINTN                     Size;
+
+  //
+  // Locate SMM Access2 Protocol
+  //
+  Status = gBS->LocateProtocol (
+                  &gEfiSmmAccess2ProtocolGuid, 
+                  NULL, 
+                  (VOID **)&SmmAccess
+                  );
+  ASSERT_EFI_ERROR (Status);
+
+  //
+  // Get SMRAM range information
+  //
+  Size = 0;
+  Status = SmmAccess->GetCapabilities (SmmAccess, &Size, NULL);
+  ASSERT (Status == EFI_BUFFER_TOO_SMALL);
+
+  mSmramRanges = (EFI_SMRAM_DESCRIPTOR *) AllocatePool (Size);
+  ASSERT (mSmramRanges != NULL);
+
+  Status = SmmAccess->GetCapabilities (SmmAccess, &Size, mSmramRanges);
+  ASSERT_EFI_ERROR (Status);
+
+  mSmramRangeCount = Size / sizeof (EFI_SMRAM_DESCRIPTOR);
+}
+
+/**
+  Check whether the start address of buffer is within any of the SMRAM ranges.
+
+  @param[in]  Buffer   The pointer to the buffer to be checked.
+
+  @retval     TURE     The buffer is in SMRAM ranges.
+  @retval     FALSE    The buffer is out of SMRAM ranges.
+**/
+BOOLEAN
+EFIAPI
+BufferInSmram (
+  IN VOID *Buffer
+  )
+{
+  UINTN  Index;
+
+  if (mSmramRanges == NULL) {
+    //
+    // SMRAM ranges is not got. Try to get them all.
+    //
+    GetSmramRanges();
+  }
+
+  for (Index = 0; Index < mSmramRangeCount; Index ++) {
+    if (((EFI_PHYSICAL_ADDRESS) (UINTN) Buffer >= mSmramRanges[Index].CpuStart) && 
+        ((EFI_PHYSICAL_ADDRESS) (UINTN) Buffer < (mSmramRanges[Index].CpuStart + mSmramRanges[Index].PhysicalSize))) {
+      return TRUE;
+    }
+  }
+
+  return FALSE;
+}
+
 /**
   Allocates one or more 4KB pages of a certain memory type.
 
@@ -145,7 +227,19 @@ FreePages (
   EFI_STATUS  Status;
 
   ASSERT (Pages != 0);
-  Status = SmmFreePages ((EFI_PHYSICAL_ADDRESS) (UINTN) Buffer, Pages);
+  if (BufferInSmram (Buffer)) {
+    //
+    // When Buffer is in SMRAM range, it should be allocated by SmmAllocatePages() service.
+    // So, SmmFreePages() service is used to free it.
+    //
+    Status = SmmFreePages ((EFI_PHYSICAL_ADDRESS) (UINTN) Buffer, Pages);
+  } else {
+    //
+    // When Buffer is out of SMRAM range, it should be allocated by gBS->AllocatePages() service.
+    // So, gBS->FreePages() service is used to free it.
+    //
+    Status = gBS->FreePages ((EFI_PHYSICAL_ADDRESS) (UINTN) Buffer, Pages);
+  }
   ASSERT_EFI_ERROR (Status);
 }
 
@@ -342,7 +436,19 @@ FreeAlignedPages (
   EFI_STATUS  Status;
 
   ASSERT (Pages != 0);
-  Status = SmmFreePages ((EFI_PHYSICAL_ADDRESS) (UINTN) Buffer, Pages);
+  if (BufferInSmram (Buffer)) {
+    //
+    // When Buffer is in SMRAM range, it should be allocated by SmmAllocatePages() service.
+    // So, SmmFreePages() service is used to free it.
+    //
+    Status = SmmFreePages ((EFI_PHYSICAL_ADDRESS) (UINTN) Buffer, Pages);
+  } else {
+    //
+    // When Buffer is out of SMRAM range, it should be allocated by gBS->AllocatePages() service.
+    // So, gBS->FreePages() service is used to free it.
+    //
+    Status = gBS->FreePages ((EFI_PHYSICAL_ADDRESS) (UINTN) Buffer, Pages);
+  }
   ASSERT_EFI_ERROR (Status);
 }
 
@@ -808,7 +914,19 @@ FreePool (
 {
   EFI_STATUS    Status;
 
-  Status = SmmFreePool (Buffer);
+  if (BufferInSmram (Buffer)) {
+    //
+    // When Buffer is in SMRAM range, it should be allocated by SmmAllocatePool() service.
+    // So, SmmFreePool() service is used to free it.
+    //
+    Status = SmmFreePool (Buffer);
+  } else {
+    //
+    // When Buffer is out of SMRAM range, it should be allocated by gBS->AllocatePool() service.
+    // So, gBS->FreePool() service is used to free it.
+    //
+    Status = gBS->FreePool (Buffer);
+  }
   ASSERT_EFI_ERROR (Status);
 }
 
diff --git a/MdeModulePkg/Library/PiSmmCoreMemoryAllocationLib/PiSmmCoreMemoryAllocationLib.inf b/MdeModulePkg/Library/PiSmmCoreMemoryAllocationLib/PiSmmCoreMemoryAllocationLib.inf
index fb51095bad..8f3e0fc116 100644
--- a/MdeModulePkg/Library/PiSmmCoreMemoryAllocationLib/PiSmmCoreMemoryAllocationLib.inf
+++ b/MdeModulePkg/Library/PiSmmCoreMemoryAllocationLib/PiSmmCoreMemoryAllocationLib.inf
@@ -40,3 +40,7 @@
 [LibraryClasses]
   DebugLib
   BaseMemoryLib
+  UefiBootServicesTableLib
+
+[Protocols]
+  gEfiSmmAccess2ProtocolGuid    ## CONSUMES