MdeModulePkg/Core: Fix heap guard issues

[mirror_edk2.git] / MdeModulePkg / Core / Dxe / Mem / HeapGuard.c

diff --git a/MdeModulePkg/Core/Dxe/Mem/HeapGuard.c b/MdeModulePkg/Core/Dxe/Mem/HeapGuard.c
index bee229f4c8a73b3101b6c2c4e600706c73822766..0f035043e15a2ead446f9b8d0de89f08a7df2dff 100644 (file)
--- a/MdeModulePkg/Core/Dxe/Mem/HeapGuard.c
+++ b/MdeModulePkg/Core/Dxe/Mem/HeapGuard.c
@@ -583,8 +583,7 @@ SetGuardPage (
   mOnGuarding = TRUE;\r
   //\r
   // Note: This might overwrite other attributes needed by other features,\r
-  // such as memory protection (NX). Please make sure they are not enabled\r
-  // at the same time.\r
+  // such as NX memory protection.\r
   //\r
   gCpu->SetMemoryAttributes (gCpu, BaseAddress, EFI_PAGE_SIZE, EFI_MEMORY_RP);\r
   mOnGuarding = FALSE;\r
@@ -605,6 +604,18 @@ UnsetGuardPage (
   IN  EFI_PHYSICAL_ADDRESS      BaseAddress\r
   )\r
 {\r
+  UINT64          Attributes;\r
+\r
+  //\r
+  // Once the Guard page is unset, it will be freed back to memory pool. NX\r
+  // memory protection must be restored for this page if NX is enabled for free\r
+  // memory.\r
+  //\r
+  Attributes = 0;\r
+  if ((PcdGet64 (PcdDxeNxMemoryProtectionPolicy) & (1 << EfiConventionalMemory)) != 0) {\r
+    Attributes |= EFI_MEMORY_XP;\r
+  }\r
+\r
   //\r
   // Set flag to make sure allocating memory without GUARD for page table\r
   // operation; otherwise infinite loops could be caused.\r
@@ -615,7 +626,7 @@ UnsetGuardPage (
   // such as memory protection (NX). Please make sure they are not enabled\r
   // at the same time.\r
   //\r
-  gCpu->SetMemoryAttributes (gCpu, BaseAddress, EFI_PAGE_SIZE, 0);\r
+  gCpu->SetMemoryAttributes (gCpu, BaseAddress, EFI_PAGE_SIZE, Attributes);\r
   mOnGuarding = FALSE;\r
 }\r
 \r
@@ -869,6 +880,15 @@ AdjustMemoryS (
 {\r
   UINT64  Target;\r
 \r
+  //\r
+  // UEFI spec requires that allocated pool must be 8-byte aligned. If it's\r
+  // indicated to put the pool near the Tail Guard, we need extra bytes to\r
+  // make sure alignment of the returned pool address.\r
+  //\r
+  if ((PcdGet8 (PcdHeapGuardPropertyMask) & BIT7) == 0) {\r
+    SizeRequested = ALIGN_VALUE(SizeRequested, 8);\r
+  }\r
+\r
   Target = Start + Size - SizeRequested;\r
 \r
   //\r
@@ -1052,7 +1072,7 @@ AdjustPoolHeadA (
   IN UINTN                   Size\r
   )\r
 {\r
-  if ((PcdGet8 (PcdHeapGuardPropertyMask) & BIT7) != 0) {\r
+  if (Memory == 0 || (PcdGet8 (PcdHeapGuardPropertyMask) & BIT7) != 0) {\r
     //\r
     // Pool head is put near the head Guard\r
     //\r
@@ -1062,6 +1082,7 @@ AdjustPoolHeadA (
   //\r
   // Pool head is put near the tail Guard\r
   //\r
+  Size = ALIGN_VALUE (Size, 8);\r
   return (VOID *)(UINTN)(Memory + EFI_PAGES_TO_SIZE (NoPages) - Size);\r
 }\r
 \r
@@ -1077,7 +1098,7 @@ AdjustPoolHeadF (
   IN EFI_PHYSICAL_ADDRESS    Memory\r
   )\r
 {\r
-  if ((PcdGet8 (PcdHeapGuardPropertyMask) & BIT7) != 0) {\r
+  if (Memory == 0 || (PcdGet8 (PcdHeapGuardPropertyMask) & BIT7) != 0) {\r
     //\r
     // Pool head is put near the head Guard\r
     //\r