From f1674e665c3de23135009f5dd645b8af1f7d0379 Mon Sep 17 00:00:00 2001
From: Vivian Nowka-Keane <vnowkakeane@linux.microsoft.com>
Date: Tue, 22 Oct 2024 14:20:41 -0700
Subject: [PATCH] UefiCpuPkg: SmmProfile: Use public Architectural MSRs from
 MdePkg

Replaced local Msr defines with inclusion of Register/Amd/Msr.h.

Signed-off-by: Vivian Nowka-Keane <vnowkakeane@linux.microsoft.com>
---
 UefiCpuPkg/PiSmmCpuDxeSmm/SmmProfile.c        | 44 +++++++++++++------
 .../PiSmmCpuDxeSmm/SmmProfileInternal.h       | 15 ++-----
 2 files changed, 35 insertions(+), 24 deletions(-)

diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/SmmProfile.c b/UefiCpuPkg/PiSmmCpuDxeSmm/SmmProfile.c
index ee64d6b6d0bc..4862cf075fed 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/SmmProfile.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/SmmProfile.c
@@ -131,7 +131,13 @@ DisableBTS (
   VOID
   )
 {
-  AsmMsrAnd64 (MSR_DEBUG_CTL, ~((UINT64)(MSR_DEBUG_CTL_BTS | MSR_DEBUG_CTL_TR)));
+  MSR_IA32_DEBUGCTL_REGISTER  DebugCtl;
+
+  DebugCtl.Uint64   = AsmReadMsr64 (MSR_IA32_DEBUGCTL);
+  DebugCtl.Bits.BTS = 0;
+  DebugCtl.Bits.TR  = 0;
+
+  AsmWriteMsr64 (MSR_IA32_DEBUGCTL, DebugCtl.Uint64);
 }
 
 /**
@@ -143,7 +149,13 @@ EnableBTS (
   VOID
   )
 {
-  AsmMsrOr64 (MSR_DEBUG_CTL, (MSR_DEBUG_CTL_BTS | MSR_DEBUG_CTL_TR));
+  MSR_IA32_DEBUGCTL_REGISTER  DebugCtl;
+
+  DebugCtl.Uint64   = AsmReadMsr64 (MSR_IA32_DEBUGCTL);
+  DebugCtl.Bits.BTS = 1;
+  DebugCtl.Bits.TR  = 1;
+
+  AsmWriteMsr64 (MSR_IA32_DEBUGCTL, DebugCtl.Uint64);
 }
 
 /**
@@ -930,15 +942,15 @@ ActivateLBR (
   VOID
   )
 {
-  UINT64  DebugCtl;
+  MSR_IA32_DEBUGCTL_REGISTER  DebugCtl;
 
-  DebugCtl = AsmReadMsr64 (MSR_DEBUG_CTL);
-  if ((DebugCtl & MSR_DEBUG_CTL_LBR) != 0) {
+  DebugCtl.Uint64 = AsmReadMsr64 (MSR_IA32_DEBUGCTL);
+  if (DebugCtl.Bits.LBR) {
     return;
   }
 
-  DebugCtl |= MSR_DEBUG_CTL_LBR;
-  AsmWriteMsr64 (MSR_DEBUG_CTL, DebugCtl);
+  DebugCtl.Bits.LBR = 1;
+  AsmWriteMsr64 (MSR_IA32_DEBUGCTL, DebugCtl.Uint64);
 }
 
 /**
@@ -952,17 +964,23 @@ ActivateBTS (
   IN      UINTN  CpuIndex
   )
 {
-  UINT64  DebugCtl;
+  MSR_IA32_DEBUGCTL_REGISTER  DebugCtl;
 
-  DebugCtl = AsmReadMsr64 (MSR_DEBUG_CTL);
-  if ((DebugCtl & MSR_DEBUG_CTL_BTS) != 0) {
+  DebugCtl.Uint64 = AsmReadMsr64 (MSR_IA32_DEBUGCTL);
+  if ((DebugCtl.Bits.BTS)) {
     return;
   }
 
   AsmWriteMsr64 (MSR_DS_AREA, (UINT64)(UINTN)mMsrDsArea[CpuIndex]);
-  DebugCtl |= (UINT64)(MSR_DEBUG_CTL_BTS | MSR_DEBUG_CTL_TR);
-  DebugCtl &= ~((UINT64)MSR_DEBUG_CTL_BTINT);
-  AsmWriteMsr64 (MSR_DEBUG_CTL, DebugCtl);
+
+  //
+  // Enable BTS
+  //
+  DebugCtl.Bits.BTS = 1;
+  DebugCtl.Bits.TR  = 1;
+
+  DebugCtl.Bits.BTINT = 0;
+  AsmWriteMsr64 (MSR_IA32_DEBUGCTL, DebugCtl.Uint64);
 }
 
 /**
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/SmmProfileInternal.h b/UefiCpuPkg/PiSmmCpuDxeSmm/SmmProfileInternal.h
index e8f3f6492f80..df2343dc493c 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/SmmProfileInternal.h
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/SmmProfileInternal.h
@@ -39,20 +39,13 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 //
 // CPU generic definition
 //
-#define   MSR_EFER     0xc0000080
-#define   MSR_EFER_XD  0x800
+#define MSR_EFER_XD  0x800
 
-#define   CPUID1_EDX_BTS_AVAILABLE  0x200000
+#define CPUID1_EDX_BTS_AVAILABLE  0x200000
 
-#define   DR6_SINGLE_STEP  0x4000
-#define   RFLAG_TF         0x100
+#define DR6_SINGLE_STEP  0x4000
 
-#define MSR_DEBUG_CTL          0x1D9
-#define   MSR_DEBUG_CTL_LBR    0x1
-#define   MSR_DEBUG_CTL_TR     0x40
-#define   MSR_DEBUG_CTL_BTS    0x80
-#define   MSR_DEBUG_CTL_BTINT  0x100
-#define MSR_DS_AREA            0x600
+#define MSR_DS_AREA  0x600
 
 #define HEAP_GUARD_NONSTOP_MODE      \
         ((PcdGet8 (PcdHeapGuardPropertyMask) & (BIT6|BIT3|BIT2)) > BIT6)