[PATCH 2/6] UefiCpuPkg/PiSmmCpu: Check for untested memory in GCD

[Hao Wu <hao.a.wu@intel.com>]

From: Jiewen Yao <jiewen.yao@intel.com>

It treats GCD untested memory as invalid SMM
communication buffer.

Cc: Star Zeng <star.zeng@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.1
Signed-off-by: Jiewen Yao <jiewen.yao@intel.com>
---
 UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c | 144 ++++++++++++++++----
 1 file changed, 120 insertions(+), 24 deletions(-)

diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c b/UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c
index 85e06b2e6a..b2ace6334e 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c
@@ -13,6 +13,13 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 
 #include "PiSmmCpuDxeSmm.h"
 
+//
+// attributes for reserved memory before it is promoted to system memory
+//
+#define EFI_MEMORY_PRESENT      0x0100000000000000ULL
+#define EFI_MEMORY_INITIALIZED  0x0200000000000000ULL
+#define EFI_MEMORY_TESTED       0x0400000000000000ULL
+
 #define NEXT_MEMORY_DESCRIPTOR(MemoryDescriptor, Size) \
   ((EFI_MEMORY_DESCRIPTOR *)((UINT8 *)(MemoryDescriptor) + (Size)))
 
@@ -23,6 +30,9 @@ EFI_MEMORY_DESCRIPTOR *mUefiMemoryMap;
 UINTN                 mUefiMemoryMapSize;
 UINTN                 mUefiDescriptorSize;
 
+EFI_GCD_MEMORY_SPACE_DESCRIPTOR   *mGcdMemSpace       = NULL;
+UINTN                             mGcdMemNumberOfDesc = 0;
+
 PAGE_ATTRIBUTE_TABLE mPageAttributeTable[] = {
   {Page4K,  SIZE_4KB, PAGING_4K_ADDRESS_MASK_64},
   {Page2M,  SIZE_2MB, PAGING_2M_ADDRESS_MASK_64},
@@ -1022,6 +1032,60 @@ MergeMemoryMapForNotPresentEntry (
   return ;
 }
 
+/**
+  This function caches the GCD memory map information.
+**/
+VOID
+GetGcdMemoryMap (
+  VOID
+  )
+{
+  UINTN                            NumberOfDescriptors;
+  EFI_GCD_MEMORY_SPACE_DESCRIPTOR  *MemSpaceMap;
+  EFI_STATUS                       Status;
+  UINTN                            Index;
+
+  Status = gDS->GetMemorySpaceMap (&NumberOfDescriptors, &MemSpaceMap);
+  if (EFI_ERROR (Status)) {
+    return ;
+  }
+
+  mGcdMemNumberOfDesc = 0;
+  for (Index = 0; Index < NumberOfDescriptors; Index++) {
+    if (MemSpaceMap[Index].GcdMemoryType == EfiGcdMemoryTypeReserved &&
+        (MemSpaceMap[Index].Capabilities & (EFI_MEMORY_PRESENT | EFI_MEMORY_INITIALIZED | EFI_MEMORY_TESTED)) ==
+          (EFI_MEMORY_PRESENT | EFI_MEMORY_INITIALIZED)
+          ) {
+      mGcdMemNumberOfDesc++;
+    }
+  }
+
+  mGcdMemSpace = AllocateZeroPool (mGcdMemNumberOfDesc * sizeof (EFI_GCD_MEMORY_SPACE_DESCRIPTOR));
+  ASSERT (mGcdMemSpace != NULL);
+  if (mGcdMemSpace == NULL) {
+    mGcdMemNumberOfDesc = 0;
+    gBS->FreePool (MemSpaceMap);
+    return ;
+  }
+
+  mGcdMemNumberOfDesc = 0;
+  for (Index = 0; Index < NumberOfDescriptors; Index++) {
+    if (MemSpaceMap[Index].GcdMemoryType == EfiGcdMemoryTypeReserved &&
+        (MemSpaceMap[Index].Capabilities & (EFI_MEMORY_PRESENT | EFI_MEMORY_INITIALIZED | EFI_MEMORY_TESTED)) ==
+          (EFI_MEMORY_PRESENT | EFI_MEMORY_INITIALIZED)
+          ) {
+      CopyMem (
+        &mGcdMemSpace[mGcdMemNumberOfDesc],
+        &MemSpaceMap[Index],
+        sizeof(EFI_GCD_MEMORY_SPACE_DESCRIPTOR)
+        );
+      mGcdMemNumberOfDesc++;
+    }
+  }
+
+  gBS->FreePool (MemSpaceMap);
+}
+
 /**
   This function caches the UEFI memory map information.
 **/
@@ -1081,6 +1145,11 @@ GetUefiMemoryMap (
   ASSERT (mUefiMemoryMap != NULL);
 
   gBS->FreePool (MemoryMap);
+
+  //
+  // Get additional information from GCD memory map.
+  //
+  GetGcdMemoryMap ();
 }
 
 /**
@@ -1102,33 +1171,52 @@ SetUefiMemMapAttributes (
 
   DEBUG ((DEBUG_INFO, "SetUefiMemMapAttributes\n"));
 
-  if (mUefiMemoryMap == NULL) {
-    DEBUG ((DEBUG_INFO, "UefiMemoryMap - NULL\n"));
-    return ;
+  if (mUefiMemoryMap != NULL) {
+    MemoryMapEntryCount = mUefiMemoryMapSize/mUefiDescriptorSize;
+    MemoryMap = mUefiMemoryMap;
+    for (Index = 0; Index < MemoryMapEntryCount; Index++) {
+      if (IsUefiPageNotPresent(MemoryMap)) {
+        Status = SmmSetMemoryAttributes (
+                   MemoryMap->PhysicalStart,
+                   EFI_PAGES_TO_SIZE((UINTN)MemoryMap->NumberOfPages),
+                   EFI_MEMORY_RP
+                   );
+        DEBUG ((
+          DEBUG_INFO,
+          "UefiMemory protection: 0x%lx - 0x%lx %r\n",
+          MemoryMap->PhysicalStart,
+          MemoryMap->PhysicalStart + (UINT64)EFI_PAGES_TO_SIZE((UINTN)MemoryMap->NumberOfPages),
+          Status
+          ));
+      }
+      MemoryMap = NEXT_MEMORY_DESCRIPTOR(MemoryMap, mUefiDescriptorSize);
+    }
   }
+  //
+  // Do not free mUefiMemoryMap, it will be checked in IsSmmCommBufferForbiddenAddress().
+  //
 
-  MemoryMapEntryCount = mUefiMemoryMapSize/mUefiDescriptorSize;
-  MemoryMap = mUefiMemoryMap;
-  for (Index = 0; Index < MemoryMapEntryCount; Index++) {
-    if (IsUefiPageNotPresent(MemoryMap)) {
+  //
+  // Set untested memory as not present.
+  //
+  if (mGcdMemSpace != NULL) {
+    for (Index = 0; Index < mGcdMemNumberOfDesc; Index++) {
       Status = SmmSetMemoryAttributes (
-                 MemoryMap->PhysicalStart,
-                 EFI_PAGES_TO_SIZE((UINTN)MemoryMap->NumberOfPages),
+                 mGcdMemSpace[Index].BaseAddress,
+                 mGcdMemSpace[Index].Length,
                  EFI_MEMORY_RP
                  );
       DEBUG ((
         DEBUG_INFO,
-        "UefiMemory protection: 0x%lx - 0x%lx %r\n",
-        MemoryMap->PhysicalStart,
-        MemoryMap->PhysicalStart + (UINT64)EFI_PAGES_TO_SIZE((UINTN)MemoryMap->NumberOfPages),
+        "GcdMemory protection: 0x%lx - 0x%lx %r\n",
+        mGcdMemSpace[Index].BaseAddress,
+        mGcdMemSpace[Index].BaseAddress + mGcdMemSpace[Index].Length,
         Status
         ));
     }
-    MemoryMap = NEXT_MEMORY_DESCRIPTOR(MemoryMap, mUefiDescriptorSize);
   }
-
   //
-  // Do free mUefiMemoryMap, it will be checked in IsSmmCommBufferForbiddenAddress().
+  // Do not free mGcdMemSpace, it will be checked in IsSmmCommBufferForbiddenAddress().
   //
 }
 
@@ -1149,21 +1237,29 @@ IsSmmCommBufferForbiddenAddress (
   UINTN                 MemoryMapEntryCount;
   UINTN                 Index;
 
-  if (mUefiMemoryMap == NULL) {
-    return FALSE;
+  if (mUefiMemoryMap != NULL) {
+    MemoryMap = mUefiMemoryMap;
+    MemoryMapEntryCount = mUefiMemoryMapSize/mUefiDescriptorSize;
+    for (Index = 0; Index < MemoryMapEntryCount; Index++) {
+      if (IsUefiPageNotPresent (MemoryMap)) {
+        if ((Address >= MemoryMap->PhysicalStart) &&
+            (Address < MemoryMap->PhysicalStart + EFI_PAGES_TO_SIZE((UINTN)MemoryMap->NumberOfPages)) ) {
+          return TRUE;
+        }
+      }
+      MemoryMap = NEXT_MEMORY_DESCRIPTOR(MemoryMap, mUefiDescriptorSize);
+    }
   }
 
-  MemoryMap = mUefiMemoryMap;
-  MemoryMapEntryCount = mUefiMemoryMapSize/mUefiDescriptorSize;
-  for (Index = 0; Index < MemoryMapEntryCount; Index++) {
-    if (IsUefiPageNotPresent (MemoryMap)) {
-      if ((Address >= MemoryMap->PhysicalStart) &&
-          (Address < MemoryMap->PhysicalStart + EFI_PAGES_TO_SIZE((UINTN)MemoryMap->NumberOfPages)) ) {
+  if (mGcdMemSpace != NULL) {
+    for (Index = 0; Index < mGcdMemNumberOfDesc; Index++) {
+      if ((Address >= mGcdMemSpace[Index].BaseAddress) &&
+          (Address < mGcdMemSpace[Index].BaseAddress + mGcdMemSpace[Index].Length) ) {
         return TRUE;
       }
     }
-    MemoryMap = NEXT_MEMORY_DESCRIPTOR(MemoryMap, mUefiDescriptorSize);
   }
+
   return FALSE;
 }
 
-- 
2.16.2.windows.1