MdeModulePkg/PiSmmCore: Fix issues in Heap Guard

[mirror_edk2.git] / MdeModulePkg / Core / PiSmmCore / HeapGuard.c

diff --git a/MdeModulePkg/Core/PiSmmCore/HeapGuard.c b/MdeModulePkg/Core/PiSmmCore/HeapGuard.c
index 1d5fb8cdb5a0409cc34b618d3363ea3f7b183196..04fa1747a1e15ea9397ee85762744d951b163f8a 100644 (file)
--- a/MdeModulePkg/Core/PiSmmCore/HeapGuard.c
+++ b/MdeModulePkg/Core/PiSmmCore/HeapGuard.c
@@ -776,6 +776,7 @@ UnsetGuardForMemory (
   )\r
 {\r
   EFI_PHYSICAL_ADDRESS  GuardPage;\r
+  UINT64                GuardBitmap;\r
 \r
   if (NumberOfPages == 0) {\r
     return;\r
@@ -784,16 +785,29 @@ UnsetGuardForMemory (
   //\r
   // Head Guard must be one page before, if any.\r
   //\r
+  //          MSB-> 1     0 <-LSB\r
+  //          -------------------\r
+  //  Head Guard -> 0     1 -> Don't free Head Guard  (shared Guard)\r
+  //  Head Guard -> 0     0 -> Free Head Guard either (not shared Guard)\r
+  //                1     X -> Don't free first page  (need a new Guard)\r
+  //                           (it'll be turned into a Guard page later)\r
+  //          -------------------\r
+  //      Start -> -1    -2\r
+  //\r
   GuardPage = Memory - EFI_PAGES_TO_SIZE (1);\r
-  if (IsHeadGuard (GuardPage)) {\r
-    if (!IsMemoryGuarded (GuardPage - EFI_PAGES_TO_SIZE (1))) {\r
+  GuardBitmap = GetGuardedMemoryBits (Memory - EFI_PAGES_TO_SIZE (2), 2);\r
+  if ((GuardBitmap & BIT1) == 0) {\r
+    //\r
+    // Head Guard exists.\r
+    //\r
+    if ((GuardBitmap & BIT0) == 0) {\r
       //\r
       // If the head Guard is not a tail Guard of adjacent memory block,\r
       // unset it.\r
       //\r
       UnsetGuardPage (GuardPage);\r
     }\r
-  } else if (IsMemoryGuarded (GuardPage)) {\r
+  } else {\r
     //\r
     // Pages before memory to free are still in Guard. It's a partial free\r
     // case. Turn first page of memory block to free into a new Guard.\r
@@ -804,16 +818,29 @@ UnsetGuardForMemory (
   //\r
   // Tail Guard must be the page after this memory block to free, if any.\r
   //\r
+  //   MSB-> 1     0 <-LSB\r
+  //  --------------------\r
+  //         1     0 <- Tail Guard -> Don't free Tail Guard  (shared Guard)\r
+  //         0     0 <- Tail Guard -> Free Tail Guard either (not shared Guard)\r
+  //         X     1               -> Don't free last page   (need a new Guard)\r
+  //                                 (it'll be turned into a Guard page later)\r
+  //  --------------------\r
+  //        +1    +0 <- End\r
+  //\r
   GuardPage = Memory + EFI_PAGES_TO_SIZE (NumberOfPages);\r
-  if (IsTailGuard (GuardPage)) {\r
-    if (!IsMemoryGuarded (GuardPage + EFI_PAGES_TO_SIZE (1))) {\r
+  GuardBitmap = GetGuardedMemoryBits (GuardPage, 2);\r
+  if ((GuardBitmap & BIT0) == 0) {\r
+    //\r
+    // Tail Guard exists.\r
+    //\r
+    if ((GuardBitmap & BIT1) == 0) {\r
       //\r
       // If the tail Guard is not a head Guard of adjacent memory block,\r
       // free it; otherwise, keep it.\r
       //\r
       UnsetGuardPage (GuardPage);\r
     }\r
-  } else if (IsMemoryGuarded (GuardPage)) {\r
+  } else {\r
     //\r
     // Pages after memory to free are still in Guard. It's a partial free\r
     // case. We need to keep one page to be a head Guard.\r
@@ -903,6 +930,7 @@ AdjustMemoryF (
   EFI_PHYSICAL_ADDRESS  Start;\r
   EFI_PHYSICAL_ADDRESS  MemoryToTest;\r
   UINTN                 PagesToFree;\r
+  UINT64                GuardBitmap;\r
 \r
   if (Memory == NULL || NumberOfPages == NULL || *NumberOfPages == 0) {\r
     return;\r
@@ -914,9 +942,22 @@ AdjustMemoryF (
   //\r
   // Head Guard must be one page before, if any.\r
   //\r
-  MemoryToTest = Start - EFI_PAGES_TO_SIZE (1);\r
-  if (IsHeadGuard (MemoryToTest)) {\r
-    if (!IsMemoryGuarded (MemoryToTest - EFI_PAGES_TO_SIZE (1))) {\r
+  //          MSB-> 1     0 <-LSB\r
+  //          -------------------\r
+  //  Head Guard -> 0     1 -> Don't free Head Guard  (shared Guard)\r
+  //  Head Guard -> 0     0 -> Free Head Guard either (not shared Guard)\r
+  //                1     X -> Don't free first page  (need a new Guard)\r
+  //                           (it'll be turned into a Guard page later)\r
+  //          -------------------\r
+  //      Start -> -1    -2\r
+  //\r
+  MemoryToTest = Start - EFI_PAGES_TO_SIZE (2);\r
+  GuardBitmap = GetGuardedMemoryBits (MemoryToTest, 2);\r
+  if ((GuardBitmap & BIT1) == 0) {\r
+    //\r
+    // Head Guard exists.\r
+    //\r
+    if ((GuardBitmap & BIT0) == 0) {\r
       //\r
       // If the head Guard is not a tail Guard of adjacent memory block,\r
       // free it; otherwise, keep it.\r
@@ -924,10 +965,10 @@ AdjustMemoryF (
       Start       -= EFI_PAGES_TO_SIZE (1);\r
       PagesToFree += 1;\r
     }\r
-  } else if (IsMemoryGuarded (MemoryToTest)) {\r
+  } else {\r
     //\r
-    // Pages before memory to free are still in Guard. It's a partial free\r
-    // case. We need to keep one page to be a tail Guard.\r
+    // No Head Guard, and pages before memory to free are still in Guard. It's a\r
+    // partial free case. We need to keep one page to be a tail Guard.\r
     //\r
     Start       += EFI_PAGES_TO_SIZE (1);\r
     PagesToFree -= 1;\r
@@ -936,19 +977,32 @@ AdjustMemoryF (
   //\r
   // Tail Guard must be the page after this memory block to free, if any.\r
   //\r
+  //   MSB-> 1     0 <-LSB\r
+  //  --------------------\r
+  //         1     0 <- Tail Guard -> Don't free Tail Guard  (shared Guard)\r
+  //         0     0 <- Tail Guard -> Free Tail Guard either (not shared Guard)\r
+  //         X     1               -> Don't free last page   (need a new Guard)\r
+  //                                 (it'll be turned into a Guard page later)\r
+  //  --------------------\r
+  //        +1    +0 <- End\r
+  //\r
   MemoryToTest = Start + EFI_PAGES_TO_SIZE (PagesToFree);\r
-  if (IsTailGuard (MemoryToTest)) {\r
-    if (!IsMemoryGuarded (MemoryToTest + EFI_PAGES_TO_SIZE (1))) {\r
+  GuardBitmap = GetGuardedMemoryBits (MemoryToTest, 2);\r
+  if ((GuardBitmap & BIT0) == 0) {\r
+    //\r
+    // Tail Guard exists.\r
+    //\r
+    if ((GuardBitmap & BIT1) == 0) {\r
       //\r
       // If the tail Guard is not a head Guard of adjacent memory block,\r
       // free it; otherwise, keep it.\r
       //\r
       PagesToFree += 1;\r
     }\r
-  } else if (IsMemoryGuarded (MemoryToTest)) {\r
+  } else if (PagesToFree > 0) {\r
     //\r
-    // Pages after memory to free are still in Guard. It's a partial free\r
-    // case. We need to keep one page to be a head Guard.\r
+    // No Tail Guard, and pages after memory to free are still in Guard. It's a\r
+    // partial free case. We need to keep one page to be a head Guard.\r
     //\r
     PagesToFree -= 1;\r
   }\r
@@ -1146,6 +1200,9 @@ SmmInternalFreePagesExWithGuard (
 \r
   AdjustMemoryF (&MemoryToFree, &PagesToFree);\r
   UnsetGuardForMemory (Memory, NumberOfPages);\r
+  if (PagesToFree == 0) {\r
+    return EFI_SUCCESS;\r
+  }\r
 \r
   return SmmInternalFreePagesEx (MemoryToFree, PagesToFree, AddRegion);\r
 }\r