[PATCH 5/6] UefiCpuPkg/PiSmmCpu: Check EFI_RUNTIME_RO in UEFI mem attrib table.

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

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

It treats the UEFI runtime page with EFI_MEMORY_RO attribute 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/PiSmmCpuDxeSmm.h         |  2 +
 UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.inf       |  1 +
 UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c | 75 +++++++++++++++++++-
 3 files changed, 75 insertions(+), 3 deletions(-)

diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h b/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h
index 1d016594e0..e3c7cff81c 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h
@@ -28,6 +28,7 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 #include <Protocol/SmmMemoryAttribute.h>
 
 #include <Guid/AcpiS3Context.h>
+#include <Guid/MemoryAttributesTable.h>
 #include <Guid/PiSmmMemoryAttributesTable.h>
 
 #include <Library/BaseLib.h>
@@ -45,6 +46,7 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 #include <Library/UefiBootServicesTableLib.h>
 #include <Library/UefiRuntimeServicesTableLib.h>
 #include <Library/DebugAgentLib.h>
+#include <Library/UefiLib.h>
 #include <Library/HobLib.h>
 #include <Library/LocalApicLib.h>
 #include <Library/UefiCpuLib.h>
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.inf b/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.inf
index 52d8c55075..a7fb7b0b14 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.inf
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.inf
@@ -117,6 +117,7 @@
   gEfiAcpi20TableGuid                      ## SOMETIMES_CONSUMES ## SystemTable
   gEfiAcpi10TableGuid                      ## SOMETIMES_CONSUMES ## SystemTable
   gEdkiiPiSmmMemoryAttributesTableGuid     ## CONSUMES ## SystemTable
+  gEfiMemoryAttributesTableGuid            ## CONSUMES ## SystemTable
 
 [FeaturePcd]
   gUefiCpuPkgTokenSpaceGuid.PcdCpuSmmDebug                         ## CONSUMES
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c b/UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c
index b2ace6334e..7d99f23426 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c
@@ -20,9 +20,6 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 #define EFI_MEMORY_INITIALIZED  0x0200000000000000ULL
 #define EFI_MEMORY_TESTED       0x0400000000000000ULL
 
-#define NEXT_MEMORY_DESCRIPTOR(MemoryDescriptor, Size) \
-  ((EFI_MEMORY_DESCRIPTOR *)((UINT8 *)(MemoryDescriptor) + (Size)))
-
 #define PREVIOUS_MEMORY_DESCRIPTOR(MemoryDescriptor, Size) \
   ((EFI_MEMORY_DESCRIPTOR *)((UINT8 *)(MemoryDescriptor) - (Size)))
 
@@ -33,6 +30,8 @@ UINTN                 mUefiDescriptorSize;
 EFI_GCD_MEMORY_SPACE_DESCRIPTOR   *mGcdMemSpace       = NULL;
 UINTN                             mGcdMemNumberOfDesc = 0;
 
+EFI_MEMORY_ATTRIBUTES_TABLE  *mUefiMemoryAttributesTable = NULL;
+
 PAGE_ATTRIBUTE_TABLE mPageAttributeTable[] = {
   {Page4K,  SIZE_4KB, PAGING_4K_ADDRESS_MASK_64},
   {Page2M,  SIZE_2MB, PAGING_2M_ADDRESS_MASK_64},
@@ -1086,6 +1085,26 @@ GetGcdMemoryMap (
   gBS->FreePool (MemSpaceMap);
 }
 
+/**
+  Get UEFI MemoryAttributesTable.
+**/
+VOID
+GetUefiMemoryAttributesTable (
+  VOID
+  )
+{
+  EFI_STATUS                   Status;
+  EFI_MEMORY_ATTRIBUTES_TABLE  *MemoryAttributesTable;
+  UINTN                        MemoryAttributesTableSize;
+
+  Status = EfiGetSystemConfigurationTable (&gEfiMemoryAttributesTableGuid, (VOID **)&MemoryAttributesTable);
+  if (!EFI_ERROR (Status)) {
+    MemoryAttributesTableSize = sizeof(EFI_MEMORY_ATTRIBUTES_TABLE) + MemoryAttributesTable->DescriptorSize * MemoryAttributesTable->NumberOfEntries;
+    mUefiMemoryAttributesTable = AllocateCopyPool (MemoryAttributesTableSize, MemoryAttributesTable);
+    ASSERT (mUefiMemoryAttributesTable != NULL);
+  }
+}
+
 /**
   This function caches the UEFI memory map information.
 **/
@@ -1150,6 +1169,11 @@ GetUefiMemoryMap (
   // Get additional information from GCD memory map.
   //
   GetGcdMemoryMap ();
+
+  //
+  // Get UEFI memory attributes table.
+  //
+  GetUefiMemoryAttributesTable ();
 }
 
 /**
@@ -1168,6 +1192,7 @@ SetUefiMemMapAttributes (
   EFI_MEMORY_DESCRIPTOR *MemoryMap;
   UINTN                 MemoryMapEntryCount;
   UINTN                 Index;
+  EFI_MEMORY_DESCRIPTOR *Entry;
 
   DEBUG ((DEBUG_INFO, "SetUefiMemMapAttributes\n"));
 
@@ -1218,6 +1243,35 @@ SetUefiMemMapAttributes (
   //
   // Do not free mGcdMemSpace, it will be checked in IsSmmCommBufferForbiddenAddress().
   //
+
+  //
+  // Set UEFI runtime memory with EFI_MEMORY_RO as not present.
+  //
+  if (mUefiMemoryAttributesTable != NULL) {
+    Entry = (EFI_MEMORY_DESCRIPTOR *)(mUefiMemoryAttributesTable + 1);
+    for (Index = 0; Index < mUefiMemoryAttributesTable->NumberOfEntries; Index++) {
+      if (Entry->Type == EfiRuntimeServicesCode || Entry->Type == EfiRuntimeServicesData) {
+        if ((Entry->Attribute & EFI_MEMORY_RO) != 0) {
+          Status = SmmSetMemoryAttributes (
+                     Entry->PhysicalStart,
+                     EFI_PAGES_TO_SIZE((UINTN)Entry->NumberOfPages),
+                     EFI_MEMORY_RP
+                     );
+          DEBUG ((
+            DEBUG_INFO,
+            "UefiMemoryAttribute protection: 0x%lx - 0x%lx %r\n",
+            Entry->PhysicalStart,
+            Entry->PhysicalStart + (UINT64)EFI_PAGES_TO_SIZE((UINTN)Entry->NumberOfPages),
+            Status
+            ));
+        }
+      }
+      Entry = NEXT_MEMORY_DESCRIPTOR (Entry, mUefiMemoryAttributesTable->DescriptorSize);
+    }
+  }
+  //
+  // Do not free mUefiMemoryAttributesTable, it will be checked in IsSmmCommBufferForbiddenAddress().
+  //
 }
 
 /**
@@ -1236,6 +1290,7 @@ IsSmmCommBufferForbiddenAddress (
   EFI_MEMORY_DESCRIPTOR *MemoryMap;
   UINTN                 MemoryMapEntryCount;
   UINTN                 Index;
+  EFI_MEMORY_DESCRIPTOR *Entry;
 
   if (mUefiMemoryMap != NULL) {
     MemoryMap = mUefiMemoryMap;
@@ -1260,6 +1315,20 @@ IsSmmCommBufferForbiddenAddress (
     }
   }
 
+  if (mUefiMemoryAttributesTable != NULL) {
+    Entry = (EFI_MEMORY_DESCRIPTOR *)(mUefiMemoryAttributesTable + 1);
+    for (Index = 0; Index < mUefiMemoryAttributesTable->NumberOfEntries; Index++) {
+      if (Entry->Type == EfiRuntimeServicesCode || Entry->Type == EfiRuntimeServicesData) {
+        if ((Entry->Attribute & EFI_MEMORY_RO) != 0) {
+          if ((Address >= Entry->PhysicalStart) &&
+              (Address < Entry->PhysicalStart + LShiftU64 (Entry->NumberOfPages, EFI_PAGE_SHIFT))) {
+            return TRUE;
+          }
+          Entry = NEXT_MEMORY_DESCRIPTOR (Entry, mUefiMemoryAttributesTable->DescriptorSize);
+        }
+      }
+    }
+  }
   return FALSE;
 }
 
-- 
2.16.2.windows.1