From 3690d30a6e3cdc4f159bae74a88d6e201e3d69f6 Mon Sep 17 00:00:00 2001
From: Jiaxin Wu <jiaxin.wu@intel.com>
Date: Wed, 26 Jun 2024 16:04:50 +0800
Subject: [PATCH] UefiCpuPkg/PiSmmCpuDxeSmm: Check logging PF address for MM

This patch is to make sure only logging PF address for MM
can run into the SmmProfilePFHandler.

Signed-off-by: Jiaxin Wu <jiaxin.wu@intel.com>
Cc: Ray Ni <ray.ni@intel.com>
Cc: Rahul Kumar <rahul1.kumar@intel.com>
Cc: Gerd Hoffmann <kraxel@redhat.com>
Cc: Star Zeng <star.zeng@intel.com>
Cc: Dun Tan <dun.tan@intel.com>
Cc: Hongbin1 Zhang <hongbin1.zhang@intel.com>
Cc: Wei6 Xu <wei6.xu@intel.com>
Cc: Yuanhao Xie <yuanhao.xie@intel.com>
---
 UefiCpuPkg/PiSmmCpuDxeSmm/NonMmramMapDxeSmm.c | 18 +++++++++++
 .../PiSmmCpuDxeSmm/NonMmramMapStandaloneMm.c  | 32 +++++++++++++++++++
 UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuCommon.h    | 13 ++++++++
 UefiCpuPkg/PiSmmCpuDxeSmm/X64/PageTbl.c       |  7 ++++
 4 files changed, 70 insertions(+)

diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/NonMmramMapDxeSmm.c b/UefiCpuPkg/PiSmmCpuDxeSmm/NonMmramMapDxeSmm.c
index 8dbcd8d950..27f159b98c 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/NonMmramMapDxeSmm.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/NonMmramMapDxeSmm.c
@@ -460,6 +460,24 @@ GetSmmProfileData (
   return Base;
 }
 
+/**
+  Return if the Address is the NonMmram logging Address.
+
+  @param[in] Address the address to be checked
+
+  @return TRUE  The address is the NonMmram logging Address.
+  @return FALSE The address is not the NonMmram logging Address.
+**/
+BOOLEAN
+IsNonMmramLoggingAddress (
+  IN UINT64  Address
+  )
+{
+  ASSERT (FALSE);
+
+  return TRUE;
+}
+
 /**
   Return if the Address is forbidden as SMM communication buffer.
 
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/NonMmramMapStandaloneMm.c b/UefiCpuPkg/PiSmmCpuDxeSmm/NonMmramMapStandaloneMm.c
index 8c5286b744..70e94a6b76 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/NonMmramMapStandaloneMm.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/NonMmramMapStandaloneMm.c
@@ -46,6 +46,38 @@ GetSmmProfileData (
   return SmmProfileDataHob.MemoryAllocation->AllocDescriptor.MemoryBaseAddress;
 }
 
+/**
+  Return if the Address is the NonMmram logging Address.
+
+  @param[in] Address the address to be checked
+
+  @return TRUE  The address is the NonMmram logging Address.
+  @return FALSE The address is not the NonMmram logging Address.
+**/
+BOOLEAN
+IsNonMmramLoggingAddress (
+  IN UINT64  Address
+  )
+{
+  EFI_PEI_HOB_POINTERS  Hob;
+
+  Hob.Raw = GetFirstHob (EFI_HOB_TYPE_RESOURCE_DESCRIPTOR);
+  while (Hob.Raw != NULL) {
+    if ((Address >= Hob.ResourceDescriptor->PhysicalStart) && (Address < Hob.ResourceDescriptor->PhysicalStart + Hob.ResourceDescriptor->ResourceLength)) {
+      if ((Hob.ResourceDescriptor->ResourceAttribute & MM_RESOURCE_ATTRIBUTE_LOGGING) != 0) {
+        return TRUE;
+      }
+
+      return FALSE;
+    }
+
+    Hob.Raw = GET_NEXT_HOB (Hob);
+    Hob.Raw = GetNextHob (EFI_HOB_TYPE_RESOURCE_DESCRIPTOR, Hob.Raw);
+  }
+
+  return FALSE;
+}
+
 /**
   Return if the Address is forbidden as SMM communication buffer.
 
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuCommon.h b/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuCommon.h
index 44bf3f3aba..73ace4efa4 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuCommon.h
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuCommon.h
@@ -979,6 +979,19 @@ GetSmmProfileData (
   IN OUT  UINT64  *Size
   );
 
+/**
+  Return if the Address is the NonMmram logging Address.
+
+  @param[in] Address the address to be checked
+
+  @return TRUE  The address is the NonMmram logging Address.
+  @return FALSE The address is not the NonMmram logging Address.
+**/
+BOOLEAN
+IsNonMmramLoggingAddress (
+  IN UINT64  Address
+  );
+
 /**
   Return if the Address is forbidden as SMM communication buffer.
 
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/X64/PageTbl.c b/UefiCpuPkg/PiSmmCpuDxeSmm/X64/PageTbl.c
index a5253e8167..f56d2849de 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/X64/PageTbl.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/X64/PageTbl.c
@@ -821,6 +821,13 @@ SmiPFHandler (
   }
 
   if (FeaturePcdGet (PcdCpuSmmProfileEnable)) {
+    if (mIsStandaloneMm) {
+      //
+      // Only logging ranges shall run here in MM env.
+      //
+      ASSERT (IsNonMmramLoggingAddress (PFAddress));
+    }
+
     SmmProfilePFHandler (
       SystemContext.SystemContextX64->Rip,
       SystemContext.SystemContextX64->ExceptionData