UefiCpuPkg/MtrrTest: Add test cases for TME-MK enable case

When TME-MK is enabled, the MtrrLib should substract the TME-MK
reserved bits from the max PA returned from CPUID instruction.

The new test case guarantees such behavior in MtrrLib.

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/Library/MtrrLib/UnitTest/MtrrLibUnitTest.c

@@ -1,7 +1,7 @@
 /** @file
   Unit tests of the MtrrLib instance of the MtrrLib class
 
-  Copyright (c) 2020, Intel Corporation. All rights reserved.<BR>
+  Copyright (c) 2020 - 2023, Intel Corporation. All rights reserved.<BR>
   SPDX-License-Identifier: BSD-2-Clause-Patent
 
 **/
@@ -30,6 +30,8 @@ STATIC MTRR_LIB_SYSTEM_PARAMETER  mSystemParameters[] = {
   { 48, TRUE, TRUE, CacheWriteThrough,   12 },
   { 48, TRUE, TRUE, CacheWriteProtected, 12 },
   { 48, TRUE, TRUE, CacheWriteCombining, 12 },
+
+  { 48, TRUE, TRUE, CacheWriteBack,      12, 7}, // 7 bits for MKTME
 };
 
 UINT32  mFixedMtrrsIndex[] = {
@@ -219,7 +221,7 @@ UnitTestMtrrSetMemoryAttributesInMtrrSettings (
     &WcCount
     );
   GenerateValidAndConfigurableMtrrPairs (
-    SystemParameter->PhysicalAddressBits,
+    SystemParameter->PhysicalAddressBits - SystemParameter->MkTmeKeyidBits,
     RawMtrrRange,
     UcCount,
     WtCount,
@@ -232,7 +234,7 @@ UnitTestMtrrSetMemoryAttributesInMtrrSettings (
   ExpectedMemoryRangesCount = ARRAY_SIZE (ExpectedMemoryRanges);
   GetEffectiveMemoryRanges (
     SystemParameter->DefaultCacheType,
-    SystemParameter->PhysicalAddressBits,
+    SystemParameter->PhysicalAddressBits - SystemParameter->MkTmeKeyidBits,
     RawMtrrRange,
     ExpectedVariableMtrrUsage,
     ExpectedMemoryRanges,
@@ -278,7 +280,7 @@ UnitTestMtrrSetMemoryAttributesInMtrrSettings (
     ActualMemoryRangesCount = ARRAY_SIZE (ActualMemoryRanges);
     CollectTestResult (
       SystemParameter->DefaultCacheType,
-      SystemParameter->PhysicalAddressBits,
+      SystemParameter->PhysicalAddressBits - SystemParameter->MkTmeKeyidBits,
       SystemParameter->VariableMtrrCount,
       &LocalMtrrs,
       ActualMemoryRanges,
@@ -325,7 +327,7 @@ UnitTestInvalidMemoryLayouts (
   SystemParameter = (MTRR_LIB_SYSTEM_PARAMETER *)Context;
 
   RangeCount = Random32 (1, ARRAY_SIZE (Ranges));
-  MaxAddress = 1ull << SystemParameter->PhysicalAddressBits;
+  MaxAddress = 1ull << (SystemParameter->PhysicalAddressBits - SystemParameter->MkTmeKeyidBits);
 
   for (Index = 0; Index < RangeCount; Index++) {
     do {
@@ -967,7 +969,7 @@ UnitTestMtrrSetMemoryAttributeInMtrrSettings (
     &WcCount
     );
   GenerateValidAndConfigurableMtrrPairs (
-    SystemParameter->PhysicalAddressBits,
+    SystemParameter->PhysicalAddressBits - SystemParameter->MkTmeKeyidBits,
     RawMtrrRange,
     UcCount,
     WtCount,
@@ -980,7 +982,7 @@ UnitTestMtrrSetMemoryAttributeInMtrrSettings (
   ExpectedMemoryRangesCount = ARRAY_SIZE (ExpectedMemoryRanges);
   GetEffectiveMemoryRanges (
     SystemParameter->DefaultCacheType,
-    SystemParameter->PhysicalAddressBits,
+    SystemParameter->PhysicalAddressBits - SystemParameter->MkTmeKeyidBits,
     RawMtrrRange,
     ExpectedVariableMtrrUsage,
     ExpectedMemoryRanges,
@@ -1019,7 +1021,7 @@ UnitTestMtrrSetMemoryAttributeInMtrrSettings (
     ActualMemoryRangesCount = ARRAY_SIZE (ActualMemoryRanges);
     CollectTestResult (
       SystemParameter->DefaultCacheType,
-      SystemParameter->PhysicalAddressBits,
+      SystemParameter->PhysicalAddressBits - SystemParameter->MkTmeKeyidBits,
       SystemParameter->VariableMtrrCount,
       &LocalMtrrs,
       ActualMemoryRanges,

UefiCpuPkg/Library/MtrrLib/UnitTest/MtrrLibUnitTest.h

@@ -1,6 +1,6 @@
 /** @file
 
-  Copyright (c) 2020, Intel Corporation. All rights reserved.<BR>
+  Copyright (c) 2020 - 2023, Intel Corporation. All rights reserved.<BR>
   SPDX-License-Identifier: BSD-2-Clause-Patent
 
 **/
@@ -40,6 +40,7 @@ typedef struct {
   BOOLEAN                   FixedMtrrSupported;
   MTRR_MEMORY_CACHE_TYPE    DefaultCacheType;
   UINT32                    VariableMtrrCount;
+  UINT8                     MkTmeKeyidBits;
 } MTRR_LIB_SYSTEM_PARAMETER;
 
 extern UINT32   mFixedMtrrsIndex[];

UefiCpuPkg/Library/MtrrLib/UnitTest/Support.c

@@ -17,7 +17,9 @@ MSR_IA32_MTRR_PHYSBASE_REGISTER  mVariableMtrrsPhysBase[MTRR_NUMBER_OF_VARIABLE_
 MSR_IA32_MTRR_PHYSMASK_REGISTER              mVariableMtrrsPhysMask[MTRR_NUMBER_OF_VARIABLE_MTRR];
 MSR_IA32_MTRR_DEF_TYPE_REGISTER              mDefTypeMsr;
 MSR_IA32_MTRRCAP_REGISTER                    mMtrrCapMsr;
+MSR_IA32_TME_ACTIVATE_REGISTER               mTmeActivateMsr;
 CPUID_VERSION_INFO_EDX                       mCpuidVersionInfoEdx;
+CPUID_STRUCTURED_EXTENDED_FEATURE_FLAGS_ECX  mCpuidExtendedFeatureFlagsEcx;
 CPUID_VIR_PHY_ADDRESS_SIZE_EAX               mCpuidVirPhyAddressSizeEax;
 
 BOOLEAN       mRandomInput;
@@ -86,6 +88,94 @@ GenerateRandomNumbers (
   fclose (File);
 }
 
+/**
+  Retrieves CPUID information using an extended leaf identifier.
+
+  Executes the CPUID instruction with EAX set to the value specified by Index
+  and ECX set to the value specified by SubIndex. This function always returns
+  Index. This function is only available on IA-32 and x64.
+
+  If Eax is not NULL, then the value of EAX after CPUID is returned in Eax.
+  If Ebx is not NULL, then the value of EBX after CPUID is returned in Ebx.
+  If Ecx is not NULL, then the value of ECX after CPUID is returned in Ecx.
+  If Edx is not NULL, then the value of EDX after CPUID is returned in Edx.
+
+  @param  Index     The 32-bit value to load into EAX prior to invoking the
+                    CPUID instruction.
+  @param  SubIndex  The 32-bit value to load into ECX prior to invoking the
+                    CPUID instruction.
+  @param  Eax       The pointer to the 32-bit EAX value returned by the CPUID
+                    instruction. This is an optional parameter that may be
+                    NULL.
+  @param  Ebx       The pointer to the 32-bit EBX value returned by the CPUID
+                    instruction. This is an optional parameter that may be
+                    NULL.
+  @param  Ecx       The pointer to the 32-bit ECX value returned by the CPUID
+                    instruction. This is an optional parameter that may be
+                    NULL.
+  @param  Edx       The pointer to the 32-bit EDX value returned by the CPUID
+                    instruction. This is an optional parameter that may be
+                    NULL.
+
+  @return Index.
+
+**/
+UINT32
+EFIAPI
+UnitTestMtrrLibAsmCpuidEx (
+  IN      UINT32  Index,
+  IN      UINT32  SubIndex,
+  OUT     UINT32  *Eax   OPTIONAL,
+  OUT     UINT32  *Ebx   OPTIONAL,
+  OUT     UINT32  *Ecx   OPTIONAL,
+  OUT     UINT32  *Edx   OPTIONAL
+  )
+{
+  switch (Index) {
+    case CPUID_SIGNATURE:
+      if (Eax != NULL) {
+        *Eax = CPUID_STRUCTURED_EXTENDED_FEATURE_FLAGS;
+      }
+
+      return Index;
+      break;
+    case CPUID_VERSION_INFO:
+      if (Edx != NULL) {
+        *Edx = mCpuidVersionInfoEdx.Uint32;
+      }
+
+      return Index;
+      break;
+    case CPUID_STRUCTURED_EXTENDED_FEATURE_FLAGS:
+      if (Ecx != NULL) {
+        *Ecx = mCpuidExtendedFeatureFlagsEcx.Uint32;
+      }
+
+      return Index;
+      break;
+    case CPUID_EXTENDED_FUNCTION:
+      if (Eax != NULL) {
+        *Eax = CPUID_VIR_PHY_ADDRESS_SIZE;
+      }
+
+      return Index;
+      break;
+    case CPUID_VIR_PHY_ADDRESS_SIZE:
+      if (Eax != NULL) {
+        *Eax = mCpuidVirPhyAddressSizeEax.Uint32;
+      }
+
+      return Index;
+      break;
+  }
+
+  //
+  // Should never fall through to here
+  //
+  ASSERT (FALSE);
+  return Index;
+}
+
 /**
   Retrieves CPUID information.
 
@@ -121,42 +211,7 @@ UnitTestMtrrLibAsmCpuid (
   OUT     UINT32  *Edx   OPTIONAL
   )
 {
-  switch (Index) {
+  return UnitTestMtrrLibAsmCpuidEx (Index, 0, Eax, Ebx, Ecx, Edx);
-    case CPUID_SIGNATURE:
-      if (Eax != NULL) {
-        *Eax = CPUID_VERSION_INFO;
-      }
-
-      return Index;
-      break;
-    case CPUID_VERSION_INFO:
-      if (Edx != NULL) {
-        *Edx = mCpuidVersionInfoEdx.Uint32;
-      }
-
-      return Index;
-      break;
-    case CPUID_EXTENDED_FUNCTION:
-      if (Eax != NULL) {
-        *Eax = CPUID_VIR_PHY_ADDRESS_SIZE;
-      }
-
-      return Index;
-      break;
-    case CPUID_VIR_PHY_ADDRESS_SIZE:
-      if (Eax != NULL) {
-        *Eax = mCpuidVirPhyAddressSizeEax.Uint32;
-      }
-
-      return Index;
-      break;
-  }
-
-  //
-  // Should never fall through to here
-  //
-  ASSERT (FALSE);
-  return Index;
 }
 
 /**
@@ -207,6 +262,10 @@ UnitTestMtrrLibAsmReadMsr64 (
     return mMtrrCapMsr.Uint64;
   }
 
+  if (MsrIndex == MSR_IA32_TME_ACTIVATE) {
+    return mTmeActivateMsr.Uint64;
+  }
+
   //
   // Should never fall through to here
   //
@@ -325,9 +384,21 @@ InitializeMtrrRegs (
   // Hook BaseLib functions used by MtrrLib that require some emulation.
   //
   gUnitTestHostBaseLib.X86->AsmCpuid   = UnitTestMtrrLibAsmCpuid;
+  gUnitTestHostBaseLib.X86->AsmCpuidEx = UnitTestMtrrLibAsmCpuidEx;
+
   gUnitTestHostBaseLib.X86->AsmReadMsr64  = UnitTestMtrrLibAsmReadMsr64;
   gUnitTestHostBaseLib.X86->AsmWriteMsr64 = UnitTestMtrrLibAsmWriteMsr64;
 
+  if (SystemParameter->MkTmeKeyidBits != 0) {
+    mCpuidExtendedFeatureFlagsEcx.Bits.TME_EN = 1;
+    mTmeActivateMsr.Bits.TmeEnable            = 1;
+    mTmeActivateMsr.Bits.MkTmeKeyidBits       = SystemParameter->MkTmeKeyidBits;
+  } else {
+    mCpuidExtendedFeatureFlagsEcx.Bits.TME_EN = 0;
+    mTmeActivateMsr.Bits.TmeEnable            = 0;
+    mTmeActivateMsr.Bits.MkTmeKeyidBits       = 0;
+  }
+
   return UNIT_TEST_PASSED;
 }