[PATCH] UefiCpuPkg/PiSmmCpu: Remove hardcode 48 address size limitation

["Ni, Ray" <ray.ni@intel.com>]

5-level paging can be enabled on CPU which supports up to 52 physical
address size. But when the feature was enabled, the 48 address size
limit was not removed and the 5-level paging testing didn't access
address >= 2^48. So the issue wasn't detected until recently an
address >= 2^48 is accessed.

Change-Id: Iaedc73be318d4b4122071efc3ba6e967a4b58fc3
Signed-off-by: Ray Ni <ray.ni@intel.com>
Cc: Eric Dong <eric.dong@intel.com>
Cc: Laszlo Ersek <lersek@redhat.com>
Cc: Rahul Kumar <rahul1.kumar@intel.com>
---
 UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c | 26 ++++++++++++++++++--------
 1 file changed, 18 insertions(+), 8 deletions(-)

diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c b/UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c
index fd6583f9d1..89143810b6 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c
@@ -1887,11 +1887,13 @@ InitializeMpServiceData (
   IN UINTN       ShadowStackSize
   )
 {
-  UINT32                    Cr3;
-  UINTN                     Index;
-  UINT8                     *GdtTssTables;
-  UINTN                     GdtTableStepSize;
-  CPUID_VERSION_INFO_EDX    RegEdx;
+  UINT32                          Cr3;
+  UINTN                           Index;
+  UINT8                           *GdtTssTables;
+  UINTN                           GdtTableStepSize;
+  CPUID_VERSION_INFO_EDX          RegEdx;
+  UINT32                          MaxExtendedFunction;
+  CPUID_VIR_PHY_ADDRESS_SIZE_EAX  VirPhyAddressSize;
 
   //
   // Determine if this CPU supports machine check
@@ -1918,9 +1920,17 @@ InitializeMpServiceData (
   // Initialize physical address mask
   // NOTE: Physical memory above virtual address limit is not supported !!!
   //
-  AsmCpuid (0x80000008, (UINT32*)&Index, NULL, NULL, NULL);
-  gPhyMask = LShiftU64 (1, (UINT8)Index) - 1;
-  gPhyMask &= (1ull << 48) - EFI_PAGE_SIZE;
+  AsmCpuid (CPUID_EXTENDED_FUNCTION, &MaxExtendedFunction, NULL, NULL, NULL);
+  if (MaxExtendedFunction >= CPUID_VIR_PHY_ADDRESS_SIZE) {
+    AsmCpuid (CPUID_VIR_PHY_ADDRESS_SIZE, &VirPhyAddressSize.Uint32, NULL, NULL, NULL);
+  } else {
+    VirPhyAddressSize.Bits.PhysicalAddressBits = 36;
+  }
+  gPhyMask  = LShiftU64 (1, VirPhyAddressSize.Bits.PhysicalAddressBits) - 1;
+  //
+  // Clear the low 12 bits
+  //
+  gPhyMask &= 0xfffffffffffff000ULL;
 
   //
   // Create page tables
-- 
2.31.1.windows.1