From d975c0de63dc32278aefa1f596e1da14ef492b9c Mon Sep 17 00:00:00 2001
From: kuqin <kuqin@microsoft.com>
Date: Fri, 19 May 2023 16:24:27 -0700
Subject: [PATCH] PiSmmCpuDxeSmm [3/8]: 202302 Override Integration

This change updated the override from mu_basecore upstream:
https://github.com/microsoft/mu_basecore/commit/13b97736c876919b9786055829caaa4fa46984b7

Signed-off-by: Kun Qin <kuqin@microsoft.com>
---
 MmSupervisorPkg/Core/Mem/SmmProfile.c | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/MmSupervisorPkg/Core/Mem/SmmProfile.c b/MmSupervisorPkg/Core/Mem/SmmProfile.c
index af0bf1a8..28934196 100644
--- a/MmSupervisorPkg/Core/Mem/SmmProfile.c
+++ b/MmSupervisorPkg/Core/Mem/SmmProfile.c
@@ -1,7 +1,7 @@
 /** @file
 Enable SMM profile.
 
-Copyright (c) 2012 - 2019, Intel Corporation. All rights reserved.<BR>
+Copyright (c) 2012 - 2023, Intel Corporation. All rights reserved.<BR>
 Copyright (c) 2017 - 2020, AMD Incorporated. All rights reserved.<BR>
 
 SPDX-License-Identifier: BSD-2-Clause-Patent
@@ -587,13 +587,18 @@ InitPaging (
     }
 
     SizeOfMemorySpace = HighBitSet64 (gPhyMask) + 1;
+    ASSERT (SizeOfMemorySpace <= 52);
+
     //
-    // Calculate the table entries of PML4E and PDPTE.
+    // Calculate the table entries of PML5E, PML4E and PDPTE.
     //
     NumberOfPml5Entries = 1;
     if (SizeOfMemorySpace > 48) {
-      NumberOfPml5Entries = (UINTN)LShiftU64 (1, SizeOfMemorySpace - 48);
-      SizeOfMemorySpace   = 48;
+      if (Enable5LevelPaging) {
+        NumberOfPml5Entries = (UINTN)LShiftU64 (1, SizeOfMemorySpace - 48);
+      }
+
+      SizeOfMemorySpace = 48;
     }
 
     NumberOfPml4Entries = 1;