[PATCH 1/2] UefiCpuPkg/CpuDxe: allow accessing (DXE) page table in SMM mode

[Jian J Wang <jian.j.wang@intel.com>]

The SMM version of MemoryAllocationLib allows to free memory allocated
in DXE (before EndOfDxe). This is done by checking the memory range and
calling gBS services to do real operation if the memory to free is out
of SMRAM. This would cause problem if some memory related features, like
Heap Guard, have to update page table to change memory attributes.
Because page table in SMM mode is different from DXE mode, gBS memory
services cannot get the correct attributes of DXE memory from SMM page
table and then cause incorrect memory manipulations.

The solution in this patch is to store the DXE page table information
(e.g. value of CR0, CR3 registers, etc.) in a global variable of CpuDxe
driver. If CpuDxe detects it's in SMM mode, it will use this global
variable to access page table instead of current processor registers.

Change-Id: If810bb1828160b8bdd8cb616d86df2859c74971f
Cc: Eric Dong <eric.dong@intel.com>
Cc: Laszlo Ersek <lersek@redhat.com>
Cc: Jiewen Yao <jiewen.yao@intel.com>
Cc: Ruiyu Ni <ruiyu.ni@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.1
Signed-off-by: Jian J Wang <jian.j.wang@intel.com>
---
 UefiCpuPkg/CpuDxe/CpuDxe.c       |   2 +-
 UefiCpuPkg/CpuDxe/CpuDxe.inf     |   1 +
 UefiCpuPkg/CpuDxe/CpuPageTable.c | 108 ++++++++++++++++++++++++++-------------
 3 files changed, 75 insertions(+), 36 deletions(-)

diff --git a/UefiCpuPkg/CpuDxe/CpuDxe.c b/UefiCpuPkg/CpuDxe/CpuDxe.c
index 6ae2dcd1c7..1fd996fc3f 100644
--- a/UefiCpuPkg/CpuDxe/CpuDxe.c
+++ b/UefiCpuPkg/CpuDxe/CpuDxe.c
@@ -404,7 +404,7 @@ CpuSetMemoryAttributes (
   // to avoid unnecessary computing.
   //
   if (mIsFlushingGCD) {
-    DEBUG((DEBUG_INFO, "  Flushing GCD\n"));
+    DEBUG((DEBUG_GCD, "  Flushing GCD\n"));
     return EFI_SUCCESS;
   }
 
diff --git a/UefiCpuPkg/CpuDxe/CpuDxe.inf b/UefiCpuPkg/CpuDxe/CpuDxe.inf
index 3c938cee53..8c8773af90 100644
--- a/UefiCpuPkg/CpuDxe/CpuDxe.inf
+++ b/UefiCpuPkg/CpuDxe/CpuDxe.inf
@@ -66,6 +66,7 @@
 [Protocols]
   gEfiCpuArchProtocolGuid                       ## PRODUCES
   gEfiMpServiceProtocolGuid                     ## PRODUCES
+  gEfiSmmBase2ProtocolGuid
 
 [Guids]
   gIdleLoopEventGuid                            ## CONSUMES           ## Event
diff --git a/UefiCpuPkg/CpuDxe/CpuPageTable.c b/UefiCpuPkg/CpuDxe/CpuPageTable.c
index e2595b4d89..bf420d3792 100644
--- a/UefiCpuPkg/CpuDxe/CpuPageTable.c
+++ b/UefiCpuPkg/CpuDxe/CpuPageTable.c
@@ -23,6 +23,7 @@
 #include <Library/DebugLib.h>
 #include <Library/UefiBootServicesTableLib.h>
 #include <Protocol/MpService.h>
+#include <Protocol/SmmBase2.h>
 
 #include "CpuDxe.h"
 #include "CpuPageTable.h"
@@ -87,7 +88,33 @@ PAGE_ATTRIBUTE_TABLE mPageAttributeTable[] = {
   {Page1G,  SIZE_1GB, PAGING_1G_ADDRESS_MASK_64},
 };
 
-PAGE_TABLE_POOL   *mPageTablePool = NULL;
+PAGE_TABLE_POOL                   *mPageTablePool = NULL;
+PAGE_TABLE_LIB_PAGING_CONTEXT     mPagingContext;
+EFI_SMM_BASE2_PROTOCOL            *mSmmBase2 = NULL;
+
+BOOLEAN
+IsInSmm (
+  VOID
+  )
+{
+  EFI_STATUS              Status;
+  BOOLEAN                 InSmm;
+
+  InSmm = FALSE;
+  if (mSmmBase2 == NULL) {
+    Status = gBS->LocateProtocol (&gEfiSmmBase2ProtocolGuid, NULL,
+                                  (VOID **)&mSmmBase2);
+    if (EFI_ERROR (Status)) {
+      mSmmBase2 = NULL;
+    }
+  }
+
+  if (mSmmBase2 != NULL) {
+    mSmmBase2->InSmm (mSmmBase2, &InSmm);
+  }
+
+  return InSmm;
+}
 
 /**
   Return current paging context.
@@ -102,42 +129,45 @@ GetCurrentPagingContext (
   UINT32                         RegEax;
   UINT32                         RegEdx;
 
-  ZeroMem(PagingContext, sizeof(*PagingContext));
-  if (sizeof(UINTN) == sizeof(UINT64)) {
-    PagingContext->MachineType = IMAGE_FILE_MACHINE_X64;
-  } else {
-    PagingContext->MachineType = IMAGE_FILE_MACHINE_I386;
-  }
-  if ((AsmReadCr0 () & BIT31) != 0) {
-    PagingContext->ContextData.X64.PageTableBase = (AsmReadCr3 () & PAGING_4K_ADDRESS_MASK_64);
-  } else {
-    PagingContext->ContextData.X64.PageTableBase = 0;
-  }
+  if (!IsInSmm ()) {
+    if (sizeof(UINTN) == sizeof(UINT64)) {
+      mPagingContext.MachineType = IMAGE_FILE_MACHINE_X64;
+    } else {
+      mPagingContext.MachineType = IMAGE_FILE_MACHINE_I386;
+    }
+    if ((AsmReadCr0 () & BIT31) != 0) {
+      mPagingContext.ContextData.X64.PageTableBase = (AsmReadCr3 () & PAGING_4K_ADDRESS_MASK_64);
+    } else {
+      mPagingContext.ContextData.X64.PageTableBase = 0;
+    }
 
-  if ((AsmReadCr4 () & BIT4) != 0) {
-    PagingContext->ContextData.Ia32.Attributes |= PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_PSE;
-  }
-  if ((AsmReadCr4 () & BIT5) != 0) {
-    PagingContext->ContextData.Ia32.Attributes |= PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_PAE;
-  }
-  if ((AsmReadCr0 () & BIT16) != 0) {
-    PagingContext->ContextData.Ia32.Attributes |= PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_WP_ENABLE;
-  }
+    if ((AsmReadCr4 () & BIT4) != 0) {
+      mPagingContext.ContextData.Ia32.Attributes |= PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_PSE;
+    }
+    if ((AsmReadCr4 () & BIT5) != 0) {
+      mPagingContext.ContextData.Ia32.Attributes |= PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_PAE;
+    }
+    if ((AsmReadCr0 () & BIT16) != 0) {
+      mPagingContext.ContextData.Ia32.Attributes |= PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_WP_ENABLE;
+    }
 
-  AsmCpuid (0x80000000, &RegEax, NULL, NULL, NULL);
-  if (RegEax > 0x80000000) {
-    AsmCpuid (0x80000001, NULL, NULL, NULL, &RegEdx);
-    if ((RegEdx & BIT20) != 0) {
-      // XD supported
-      if ((AsmReadMsr64 (0xC0000080) & BIT11) != 0) {
-        // XD activated
-        PagingContext->ContextData.Ia32.Attributes |= PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_XD_ACTIVATED;
+    AsmCpuid (0x80000000, &RegEax, NULL, NULL, NULL);
+    if (RegEax > 0x80000000) {
+      AsmCpuid (0x80000001, NULL, NULL, NULL, &RegEdx);
+      if ((RegEdx & BIT20) != 0) {
+        // XD supported
+        if ((AsmReadMsr64 (0xC0000080) & BIT11) != 0) {
+          // XD activated
+          mPagingContext.ContextData.Ia32.Attributes |= PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_XD_ACTIVATED;
+        }
+      }
+      if ((RegEdx & BIT26) != 0) {
+        mPagingContext.ContextData.Ia32.Attributes |= PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_PAGE_1G_SUPPORT;
       }
-    }
-    if ((RegEdx & BIT26) != 0) {
-      PagingContext->ContextData.Ia32.Attributes |= PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_PAGE_1G_SUPPORT;
     }
   }
+
+  CopyMem (PagingContext, &mPagingContext, sizeof (mPagingContext));
 }
 
 /**
@@ -507,7 +537,10 @@ IsReadOnlyPageWriteProtected (
   VOID
   )
 {
-  return ((AsmReadCr0 () & BIT16) != 0);
+  if (!IsInSmm ()) {
+    return ((AsmReadCr0 () & BIT16) != 0);
+  }
+  return FALSE;
 }
 
 /**
@@ -518,7 +551,9 @@ DisableReadOnlyPageWriteProtect (
   VOID
   )
 {
-  AsmWriteCr0 (AsmReadCr0() & ~BIT16);
+  if (!IsInSmm ()) {
+    AsmWriteCr0 (AsmReadCr0 () & ~BIT16);
+  }
 }
 
 /**
@@ -529,7 +564,9 @@ EnableReadOnlyPageWriteProtect (
   VOID
   )
 {
-  AsmWriteCr0 (AsmReadCr0() | BIT16);
+  if (!IsInSmm ()) {
+    AsmWriteCr0 (AsmReadCr0 () | BIT16);
+  }
 }
 
 /**
@@ -1054,6 +1091,7 @@ InitializePageTableLib (
 {
   PAGE_TABLE_LIB_PAGING_CONTEXT     CurrentPagingContext;
 
+  ZeroMem (&mPagingContext, sizeof(mPagingContext));
   GetCurrentPagingContext (&CurrentPagingContext);
 
   //
-- 
2.16.2.windows.1