From bb40c3836f896d231347a0a1c388ecc4af0b56c1 Mon Sep 17 00:00:00 2001
From: Ray Ni <ray.ni@intel.com>
Date: Mon, 27 Feb 2023 13:44:29 +0800
Subject: [PATCH] UefiCpuPkg/CpuDxe: Substract TME-MK KEY_ID_BITS from CPU max
 PA

CPUID enumeration of MAX_PA is unaffected by TME-MK activation and
will continue to report the maximum physical address bits available
for software to use, irrespective of the number of KeyID bits.

So, we need to check if TME is enabled and adjust the PA size
accordingly.

Signed-off-by: Ray Ni <ray.ni@intel.com>
Cc: Eric Dong <eric.dong@intel.com>
Cc: Rahul Kumar <rahul1.kumar@intel.com>
Cc: Gerd Hoffmann <kraxel@redhat.com>
Cc: Michael D Kinney <michael.d.kinney@intel.com>
Cc: Ahmad Anadani <ahmad.anadani@intel.com>
Acked-by: Gerd Hoffmann <kraxel@redhat.com>
Reviewed-by: Michael D Kinney <michael.d.kinney@intel.com>
---
 UefiCpuPkg/CpuDxe/CpuDxe.c | 24 ++++++++++++++++++++++--
 1 file changed, 22 insertions(+), 2 deletions(-)

diff --git a/UefiCpuPkg/CpuDxe/CpuDxe.c b/UefiCpuPkg/CpuDxe/CpuDxe.c
index 920976c576..3febd59d99 100644
--- a/UefiCpuPkg/CpuDxe/CpuDxe.c
+++ b/UefiCpuPkg/CpuDxe/CpuDxe.c
@@ -505,8 +505,11 @@ InitializeMtrrMask (
   VOID
   )
 {
-  UINT32                          MaxExtendedFunction;
-  CPUID_VIR_PHY_ADDRESS_SIZE_EAX  VirPhyAddressSize;
+  UINT32                                       MaxExtendedFunction;
+  CPUID_VIR_PHY_ADDRESS_SIZE_EAX               VirPhyAddressSize;
+  UINT32                                       MaxFunction;
+  CPUID_STRUCTURED_EXTENDED_FEATURE_FLAGS_ECX  ExtendedFeatureFlagsEcx;
+  MSR_IA32_TME_ACTIVATE_REGISTER               TmeActivate;
 
   AsmCpuid (CPUID_EXTENDED_FUNCTION, &MaxExtendedFunction, NULL, NULL, NULL);
 
@@ -516,6 +519,23 @@ InitializeMtrrMask (
     VirPhyAddressSize.Bits.PhysicalAddressBits = 36;
   }
 
+  //
+  // CPUID enumeration of MAX_PA is unaffected by TME-MK activation and will continue
+  // to report the maximum physical address bits available for software to use,
+  // irrespective of the number of KeyID bits.
+  // So, we need to check if TME is enabled and adjust the PA size accordingly.
+  //
+  AsmCpuid (CPUID_SIGNATURE, &MaxFunction, NULL, NULL, NULL);
+  if (MaxFunction >= CPUID_STRUCTURED_EXTENDED_FEATURE_FLAGS) {
+    AsmCpuidEx (CPUID_STRUCTURED_EXTENDED_FEATURE_FLAGS, 0, NULL, NULL, &ExtendedFeatureFlagsEcx.Uint32, NULL);
+    if (ExtendedFeatureFlagsEcx.Bits.TME_EN == 1) {
+      TmeActivate.Uint64 = AsmReadMsr64 (MSR_IA32_TME_ACTIVATE);
+      if (TmeActivate.Bits.TmeEnable == 1) {
+        VirPhyAddressSize.Bits.PhysicalAddressBits -= TmeActivate.Bits.MkTmeKeyidBits;
+      }
+    }
+  }
+
   mValidMtrrBitsMask    = LShiftU64 (1, VirPhyAddressSize.Bits.PhysicalAddressBits) - 1;
   mValidMtrrAddressMask = mValidMtrrBitsMask & 0xfffffffffffff000ULL;
 }