1 files changed, 16 insertions, 1 deletions
diff --git a/UefiCpuPkg/CpuDxe/CpuDxe.c b/UefiCpuPkg/CpuDxe/CpuDxe.c
index 8ddebabd02..6ae2dcd1c7 100644
--- a/UefiCpuPkg/CpuDxe/CpuDxe.c
+++ b/UefiCpuPkg/CpuDxe/CpuDxe.c
@@ -25,6 +25,7 @@
 BOOLEAN                   InterruptState = FALSE;
 EFI_HANDLE                mCpuHandle = NULL;
 BOOLEAN                   mIsFlushingGCD;
+BOOLEAN                   mIsAllocatingPageTable = FALSE;
 UINT64                    mValidMtrrAddressMask;
 UINT64                    mValidMtrrBitsMask;
 UINT64                    mTimerPeriod = 0;
@@ -407,6 +408,20 @@ CpuSetMemoryAttributes (
     return EFI_SUCCESS;
   }
 
+  //
+  // During memory attributes updating, new pages may be allocated to setup
+  // smaller granularity of page table. Page allocation action might then cause
+  // another calling of CpuSetMemoryAttributes() recursively, due to memory
+  // protection policy configured (such as PcdDxeNxMemoryProtectionPolicy).
+  // Since this driver will always protect memory used as page table by itself,
+  // there's no need to apply protection policy requested from memory service.
+  // So it's safe to just return EFI_SUCCESS if this time of calling is caused
+  // by page table memory allocation.
+  //
+  if (mIsAllocatingPageTable) {
+    DEBUG((DEBUG_VERBOSE, "  Allocating page table memory\n"));
+    return EFI_SUCCESS;
+  }
 
   CacheAttributes = Attributes & CACHE_ATTRIBUTE_MASK;
   MemoryAttributes = Attributes & MEMORY_ATTRIBUTE_MASK;
@@ -487,7 +502,7 @@ CpuSetMemoryAttributes (
   //
   // Set memory attribute by page table
   //
-  return AssignMemoryPageAttributes (NULL, BaseAddress, Length, MemoryAttributes, AllocatePages);
+  return AssignMemoryPageAttributes (NULL, BaseAddress, Length, MemoryAttributes, NULL);
 }
 
 /**