[Patch V5 13/14] UefiCpuPkg: Refinement to smm runtime InitPaging() code

["duntan" <dun.tan@intel.com>]

This commit is code refinement to current smm runtime InitPaging()
page table update code. In InitPaging(), if PcdCpuSmmProfileEnable
is TRUE, use ConvertMemoryPageAttributes() API to map the range in
mProtectionMemRange to the attrbute recorded in the attribute field
of mProtectionMemRange, map the range outside mProtectionMemRange
as non-present. If PcdCpuSmmProfileEnable is FALSE, only need to
set the ranges not in mSmmCpuSmramRanges as NX.

Signed-off-by: Dun Tan <dun.tan@intel.com>
Cc: Eric Dong <eric.dong@intel.com>
Cc: Ray Ni <ray.ni@intel.com>
Cc: Rahul Kumar <rahul1.kumar@intel.com>
Cc: Gerd Hoffmann <kraxel@redhat.com>
---
 UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h |  37 +++++++++++++++++++++++++++++++++++++
 UefiCpuPkg/PiSmmCpuDxeSmm/SmmProfile.c     | 293 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
 2 files changed, 101 insertions(+), 229 deletions(-)

diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h b/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h
index 5399659bc0..12ad86028e 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h
@@ -725,6 +725,43 @@ SmmBlockingStartupThisAp (
   IN OUT  VOID              *ProcArguments OPTIONAL
   );
 
+/**
+  This function modifies the page attributes for the memory region specified by BaseAddress and
+  Length from their current attributes to the attributes specified by Attributes.
+
+  Caller should make sure BaseAddress and Length is at page boundary.
+
+  @param[in]   PageTableBase    The page table base.
+  @param[in]   BaseAddress      The physical address that is the start address of a memory region.
+  @param[in]   Length           The size in bytes of the memory region.
+  @param[in]   Attributes       The bit mask of attributes to modify for the memory region.
+  @param[in]   IsSet            TRUE means to set attributes. FALSE means to clear attributes.
+  @param[out]  IsModified       TRUE means page table modified. FALSE means page table not modified.
+
+  @retval RETURN_SUCCESS           The attributes were modified for the memory region.
+  @retval RETURN_ACCESS_DENIED     The attributes for the memory resource range specified by
+                                   BaseAddress and Length cannot be modified.
+  @retval RETURN_INVALID_PARAMETER Length is zero.
+                                   Attributes specified an illegal combination of attributes that
+                                   cannot be set together.
+  @retval RETURN_OUT_OF_RESOURCES  There are not enough system resources to modify the attributes of
+                                   the memory resource range.
+  @retval RETURN_UNSUPPORTED       The processor does not support one or more bytes of the memory
+                                   resource range specified by BaseAddress and Length.
+                                   The bit mask of attributes is not support for the memory resource
+                                   range specified by BaseAddress and Length.
+**/
+RETURN_STATUS
+ConvertMemoryPageAttributes (
+  IN  UINTN             PageTableBase,
+  IN  PAGING_MODE       PagingMode,
+  IN  PHYSICAL_ADDRESS  BaseAddress,
+  IN  UINT64            Length,
+  IN  UINT64            Attributes,
+  IN  BOOLEAN           IsSet,
+  OUT BOOLEAN           *IsModified   OPTIONAL
+  );
+
 /**
   This function sets the attributes for the memory region specified by BaseAddress and
   Length from their current attributes to the attributes specified by Attributes.
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/SmmProfile.c b/UefiCpuPkg/PiSmmCpuDxeSmm/SmmProfile.c
index 2a1f132b29..ab6b79ddf4 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/SmmProfile.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/SmmProfile.c
@@ -586,254 +586,89 @@ InitPaging (
   VOID
   )
 {
-  UINT64    Pml5Entry;
-  UINT64    Pml4Entry;
-  UINT64    *Pml5;
-  UINT64    *Pml4;
-  UINT64    *Pdpt;
-  UINT64    *Pd;
-  UINT64    *Pt;
-  UINTN     Address;
-  UINTN     Pml5Index;
-  UINTN     Pml4Index;
-  UINTN     PdptIndex;
-  UINTN     PdIndex;
-  UINTN     PtIndex;
-  UINTN     NumberOfPdptEntries;
-  UINTN     NumberOfPml4Entries;
-  UINTN     NumberOfPml5Entries;
-  UINTN     SizeOfMemorySpace;
-  BOOLEAN   Nx;
-  IA32_CR4  Cr4;
-  BOOLEAN   Enable5LevelPaging;
-  BOOLEAN   WpEnabled;
-  BOOLEAN   CetEnabled;
-
-  Cr4.UintN          = AsmReadCr4 ();
-  Enable5LevelPaging = (BOOLEAN)(Cr4.Bits.LA57 == 1);
-
-  if (sizeof (UINTN) == sizeof (UINT64)) {
-    if (!Enable5LevelPaging) {
-      Pml5Entry = (UINTN)mSmmProfileCr3 | IA32_PG_P;
-      Pml5      = &Pml5Entry;
-    } else {
-      Pml5 = (UINT64 *)(UINTN)mSmmProfileCr3;
-    }
-
-    SizeOfMemorySpace = HighBitSet64 (gPhyMask) + 1;
-    ASSERT (SizeOfMemorySpace <= 52);
-
-    //
-    // Calculate the table entries of PML5E, PML4E and PDPTE.
-    //
-    NumberOfPml5Entries = 1;
-    if (SizeOfMemorySpace > 48) {
-      if (Enable5LevelPaging) {
-        NumberOfPml5Entries = (UINTN)LShiftU64 (1, SizeOfMemorySpace - 48);
-      }
-
-      SizeOfMemorySpace = 48;
-    }
-
-    NumberOfPml4Entries = 1;
-    if (SizeOfMemorySpace > 39) {
-      NumberOfPml4Entries = (UINTN)LShiftU64 (1, SizeOfMemorySpace - 39);
-      SizeOfMemorySpace   = 39;
-    }
-
-    NumberOfPdptEntries = 1;
-    ASSERT (SizeOfMemorySpace > 30);
-    NumberOfPdptEntries = (UINTN)LShiftU64 (1, SizeOfMemorySpace - 30);
+  RETURN_STATUS  Status;
+  UINTN          Index;
+  UINTN          PageTable;
+  UINT64         Base;
+  UINT64         Length;
+  UINT64         Limit;
+  UINT64         PreviousAddress;
+  UINT64         MemoryAttrMask;
+  BOOLEAN        WpEnabled;
+  BOOLEAN        CetEnabled;
+
+  PageTable = AsmReadCr3 ();
+  if (sizeof (UINTN) == sizeof (UINT32)) {
+    Limit = BASE_4GB;
   } else {
-    Pml4Entry           = (UINTN)mSmmProfileCr3 | IA32_PG_P;
-    Pml4                = &Pml4Entry;
-    Pml5Entry           = (UINTN)Pml4 | IA32_PG_P;
-    Pml5                = &Pml5Entry;
-    NumberOfPml5Entries = 1;
-    NumberOfPml4Entries = 1;
-    NumberOfPdptEntries = 4;
+    Limit = (IsRestrictedMemoryAccess ()) ? LShiftU64 (1, mPhysicalAddressBits) : BASE_4GB;
   }
 
   DisableReadOnlyPageWriteProtect (&WpEnabled, &CetEnabled);
   //
-  // Go through page table and change 2MB-page into 4KB-page.
+  // [0, 4k] may be non-present.
   //
-  for (Pml5Index = 0; Pml5Index < NumberOfPml5Entries; Pml5Index++) {
-    if ((Pml5[Pml5Index] & IA32_PG_P) == 0) {
-      //
-      // If PML5 entry does not exist, skip it
-      //
-      continue;
-    }
+  PreviousAddress = ((PcdGet8 (PcdNullPointerDetectionPropertyMask) & BIT1) != 0) ? BASE_4KB : 0;
 
-    Pml4 = (UINT64 *)(UINTN)(Pml5[Pml5Index] & PHYSICAL_ADDRESS_MASK);
-    for (Pml4Index = 0; Pml4Index < NumberOfPml4Entries; Pml4Index++) {
-      if ((Pml4[Pml4Index] & IA32_PG_P) == 0) {
-        //
-        // If PML4 entry does not exist, skip it
-        //
-        continue;
+  DEBUG ((DEBUG_INFO, "Patch page table start ...\n"));
+  if (FeaturePcdGet (PcdCpuSmmProfileEnable)) {
+    for (Index = 0; Index < mProtectionMemRangeCount; Index++) {
+      MemoryAttrMask = 0;
+      if (mProtectionMemRange[Index].Nx == TRUE) {
+        MemoryAttrMask |= EFI_MEMORY_XP;
       }
 
-      Pdpt = (UINT64 *)(UINTN)(Pml4[Pml4Index] & ~mAddressEncMask & PHYSICAL_ADDRESS_MASK);
-      for (PdptIndex = 0; PdptIndex < NumberOfPdptEntries; PdptIndex++, Pdpt++) {
-        if ((*Pdpt & IA32_PG_P) == 0) {
-          //
-          // If PDPT entry does not exist, skip it
-          //
-          continue;
-        }
-
-        if ((*Pdpt & IA32_PG_PS) != 0) {
-          //
-          // This is 1G entry, skip it
-          //
-          continue;
-        }
-
-        Pd = (UINT64 *)(UINTN)(*Pdpt & ~mAddressEncMask & PHYSICAL_ADDRESS_MASK);
-        if (Pd == 0) {
-          continue;
-        }
-
-        for (PdIndex = 0; PdIndex < SIZE_4KB / sizeof (*Pd); PdIndex++, Pd++) {
-          if ((*Pd & IA32_PG_P) == 0) {
-            //
-            // If PD entry does not exist, skip it
-            //
-            continue;
-          }
-
-          Address = (UINTN)LShiftU64 (
-                             LShiftU64 (
-                               LShiftU64 ((Pml5Index << 9) + Pml4Index, 9) + PdptIndex,
-                               9
-                               ) + PdIndex,
-                             21
-                             );
-
-          //
-          // If it is 2M page, check IsAddressSplit()
-          //
-          if (((*Pd & IA32_PG_PS) != 0) && IsAddressSplit (Address)) {
-            //
-            // Based on current page table, create 4KB page table for split area.
-            //
-            ASSERT (Address == (*Pd & PHYSICAL_ADDRESS_MASK));
-
-            Pt = AllocatePageTableMemory (1);
-            ASSERT (Pt != NULL);
+      if (mProtectionMemRange[Index].Present == FALSE) {
+        MemoryAttrMask = EFI_MEMORY_RP;
+      }
 
-            // Split it
-            for (PtIndex = 0; PtIndex < SIZE_4KB / sizeof (*Pt); PtIndex++) {
-              Pt[PtIndex] = Address + ((PtIndex << 12) | mAddressEncMask | PAGE_ATTRIBUTE_BITS);
-            } // end for PT
+      Base   = mProtectionMemRange[Index].Range.Base;
+      Length = mProtectionMemRange[Index].Range.Top - Base;
+      if (MemoryAttrMask != 0) {
+        Status = ConvertMemoryPageAttributes (PageTable, mPagingMode, Base, Length, MemoryAttrMask, TRUE, NULL);
+        ASSERT_RETURN_ERROR (Status);
+      }
 
-            *Pd = (UINT64)(UINTN)Pt | mAddressEncMask | PAGE_ATTRIBUTE_BITS;
-          } // end if IsAddressSplit
-        } // end for PD
-      } // end for PDPT
-    } // end for PML4
-  } // end for PML5
+      if (Base > PreviousAddress) {
+        //
+        // Mark the ranges not in mProtectionMemRange as non-present.
+        //
+        MemoryAttrMask = EFI_MEMORY_RP;
+        Status         = ConvertMemoryPageAttributes (PageTable, mPagingMode, PreviousAddress, Base - PreviousAddress, MemoryAttrMask, TRUE, NULL);
+        ASSERT_RETURN_ERROR (Status);
+      }
 
-  //
-  // Go through page table and set several page table entries to absent or execute-disable.
-  //
-  DEBUG ((DEBUG_INFO, "Patch page table start ...\n"));
-  for (Pml5Index = 0; Pml5Index < NumberOfPml5Entries; Pml5Index++) {
-    if ((Pml5[Pml5Index] & IA32_PG_P) == 0) {
-      //
-      // If PML5 entry does not exist, skip it
-      //
-      continue;
+      PreviousAddress = Base + Length;
     }
 
-    Pml4 = (UINT64 *)(UINTN)(Pml5[Pml5Index] & PHYSICAL_ADDRESS_MASK);
-    for (Pml4Index = 0; Pml4Index < NumberOfPml4Entries; Pml4Index++) {
-      if ((Pml4[Pml4Index] & IA32_PG_P) == 0) {
+    //
+    // This assignment is for setting the last remaining range
+    //
+    MemoryAttrMask = EFI_MEMORY_RP;
+  } else {
+    MemoryAttrMask = EFI_MEMORY_XP;
+    for (Index = 0; Index < mSmmCpuSmramRangeCount; Index++) {
+      Base = mSmmCpuSmramRanges[Index].CpuStart;
+      if (Base > PreviousAddress) {
         //
-        // If PML4 entry does not exist, skip it
+        // Mark the ranges not in mSmmCpuSmramRanges as NX.
         //
-        continue;
+        Status = ConvertMemoryPageAttributes (PageTable, mPagingMode, PreviousAddress, Base - PreviousAddress, MemoryAttrMask, TRUE, NULL);
+        ASSERT_RETURN_ERROR (Status);
       }
 
-      Pdpt = (UINT64 *)(UINTN)(Pml4[Pml4Index] & ~mAddressEncMask & PHYSICAL_ADDRESS_MASK);
-      for (PdptIndex = 0; PdptIndex < NumberOfPdptEntries; PdptIndex++, Pdpt++) {
-        if ((*Pdpt & IA32_PG_P) == 0) {
-          //
-          // If PDPT entry does not exist, skip it
-          //
-          continue;
-        }
-
-        if ((*Pdpt & IA32_PG_PS) != 0) {
-          //
-          // This is 1G entry, set NX bit and skip it
-          //
-          if (mXdSupported) {
-            *Pdpt = *Pdpt | IA32_PG_NX;
-          }
-
-          continue;
-        }
-
-        Pd = (UINT64 *)(UINTN)(*Pdpt & ~mAddressEncMask & PHYSICAL_ADDRESS_MASK);
-        if (Pd == 0) {
-          continue;
-        }
+      PreviousAddress = mSmmCpuSmramRanges[Index].CpuStart + mSmmCpuSmramRanges[Index].PhysicalSize;
+    }
+  }
 
-        for (PdIndex = 0; PdIndex < SIZE_4KB / sizeof (*Pd); PdIndex++, Pd++) {
-          if ((*Pd & IA32_PG_P) == 0) {
-            //
-            // If PD entry does not exist, skip it
-            //
-            continue;
-          }
-
-          Address = (UINTN)LShiftU64 (
-                             LShiftU64 (
-                               LShiftU64 ((Pml5Index << 9) + Pml4Index, 9) + PdptIndex,
-                               9
-                               ) + PdIndex,
-                             21
-                             );
-
-          if ((*Pd & IA32_PG_PS) != 0) {
-            // 2MB page
-
-            if (!IsAddressValid (Address, &Nx)) {
-              //
-              // Patch to remove Present flag and RW flag
-              //
-              *Pd = *Pd & (INTN)(INT32)(~PAGE_ATTRIBUTE_BITS);
-            }
-
-            if (Nx && mXdSupported) {
-              *Pd = *Pd | IA32_PG_NX;
-            }
-          } else {
-            // 4KB page
-            Pt = (UINT64 *)(UINTN)(*Pd & ~mAddressEncMask & PHYSICAL_ADDRESS_MASK);
-            if (Pt == 0) {
-              continue;
-            }
-
-            for (PtIndex = 0; PtIndex < SIZE_4KB / sizeof (*Pt); PtIndex++, Pt++) {
-              if (!IsAddressValid (Address, &Nx)) {
-                *Pt = *Pt & (INTN)(INT32)(~PAGE_ATTRIBUTE_BITS);
-              }
-
-              if (Nx && mXdSupported) {
-                *Pt = *Pt | IA32_PG_NX;
-              }
-
-              Address += SIZE_4KB;
-            } // end for PT
-          } // end if PS
-        } // end for PD
-      } // end for PDPT
-    } // end for PML4
-  } // end for PML5
+  if (PreviousAddress < Limit) {
+    //
+    // Set the last remaining range to EFI_MEMORY_RP/EFI_MEMORY_XP.
+    // This path applies to both SmmProfile enable/disable case.
+    //
+    Status = ConvertMemoryPageAttributes (PageTable, mPagingMode, PreviousAddress, Limit - PreviousAddress, MemoryAttrMask, TRUE, NULL);
+    ASSERT_RETURN_ERROR (Status);
+  }
 
   EnableReadOnlyPageWriteProtect (WpEnabled, CetEnabled);
 
-- 
2.31.1.windows.1