MdeModulePkg/Core: allow HeapGuard even before CpuArchProtocol installed

[mirror_edk2.git] / MdeModulePkg / Core / Dxe / Gcd / Gcd.c

diff --git a/MdeModulePkg/Core/Dxe/Gcd/Gcd.c b/MdeModulePkg/Core/Dxe/Gcd/Gcd.c
index 8fbc3d282cdbe466a7fb355d4d7834ef6fc24891..77f4adb4bc0181a0900a1734698c4ba9ee432b4b 100644 (file)
--- a/MdeModulePkg/Core/Dxe/Gcd/Gcd.c
+++ b/MdeModulePkg/Core/Dxe/Gcd/Gcd.c
@@ -16,6 +16,7 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 \r
 #include "DxeMain.h"\r
 #include "Gcd.h"\r
+#include "Mem/HeapGuard.h"\r
 \r
 #define MINIMUM_INITIAL_MEMORY_SIZE 0x10000\r
 \r
@@ -391,12 +392,21 @@ CoreAllocateGcdMapEntry (
   IN OUT EFI_GCD_MAP_ENTRY  **BottomEntry\r
   )\r
 {\r
+  //\r
+  // Set to mOnGuarding to TRUE before memory allocation. This will make sure\r
+  // that the entry memory is not "guarded" by HeapGuard. Otherwise it might\r
+  // cause problem when it's freed (if HeapGuard is enabled).\r
+  //\r
+  mOnGuarding = TRUE;\r
   *TopEntry = AllocateZeroPool (sizeof (EFI_GCD_MAP_ENTRY));\r
+  mOnGuarding = FALSE;\r
   if (*TopEntry == NULL) {\r
     return EFI_OUT_OF_RESOURCES;\r
   }\r
 \r
+  mOnGuarding = TRUE;\r
   *BottomEntry = AllocateZeroPool (sizeof (EFI_GCD_MAP_ENTRY));\r
+  mOnGuarding = FALSE;\r
   if (*BottomEntry == NULL) {\r
     CoreFreePool (*TopEntry);\r
     return EFI_OUT_OF_RESOURCES;\r