MdeModulePkg/Core: fix mem alloc issues in heap guard

[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 19245049c28b847c66665c5ea08a4a967632dd0b..ac043b5d9b264b2c3e27cff8727d1824d52aa304 100644 (file)
--- a/MdeModulePkg/Core/Dxe/Mem/HeapGuard.c
+++ b/MdeModulePkg/Core/Dxe/Mem/HeapGuard.c
@@ -225,8 +225,8 @@ FindGuardedMemoryMap (
   //\r
   // Adjust current map table depth according to the address to access\r
   //\r
-  while (mMapLevel < GUARDED_HEAP_MAP_TABLE_DEPTH\r
-         &&\r
+  while (AllocMapUnit &&\r
+         mMapLevel < GUARDED_HEAP_MAP_TABLE_DEPTH &&\r
          RShiftU64 (\r
            Address,\r
            mLevelShift[GUARDED_HEAP_MAP_TABLE_DEPTH - mMapLevel - 1]\r
@@ -904,6 +904,10 @@ AdjustMemoryS (
   }\r
 \r
   Target = Start + Size - SizeRequested;\r
+  ASSERT (Target >= Start);\r
+  if (Target == 0) {\r
+    return 0;\r
+  }\r
 \r
   if (!IsGuardPage (Start + Size)) {\r
     // No Guard at tail to share. One more page is needed.\r