From 12957e56d26d702ab7be913f6a7a9c14e4bac6a9 Mon Sep 17 00:00:00 2001
From: Jian J Wang <jian.j.wang@intel.com>
Date: Tue, 13 Mar 2018 16:29:51 +0800
Subject: [PATCH] MdeModulePkg/PiSmmCore: fix mem alloc issues in heap guard

This patch fixes the same issues in Heap Guard in DXE core, which is fixed
in another patch.

Cc: Star Zeng <star.zeng@intel.com>
Cc: Eric Dong <eric.dong@intel.com>
Cc: Jiewen Yao <jiewen.yao@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.1
Signed-off-by: Jian J Wang <jian.j.wang@intel.com>
Reviewed-by: Jiewen Yao <jiewen.yao@intel.com>
---
 MdeModulePkg/Core/PiSmmCore/HeapGuard.c | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/MdeModulePkg/Core/PiSmmCore/HeapGuard.c b/MdeModulePkg/Core/PiSmmCore/HeapGuard.c
index aa9c25d102..923af93de2 100644
--- a/MdeModulePkg/Core/PiSmmCore/HeapGuard.c
+++ b/MdeModulePkg/Core/PiSmmCore/HeapGuard.c
@@ -251,8 +251,8 @@ FindGuardedMemoryMap (
   //
   // Adjust current map table depth according to the address to access
   //
-  while (mMapLevel < GUARDED_HEAP_MAP_TABLE_DEPTH
-         &&
+  while (AllocMapUnit &&
+         mMapLevel < GUARDED_HEAP_MAP_TABLE_DEPTH &&
          RShiftU64 (
            Address,
            mLevelShift[GUARDED_HEAP_MAP_TABLE_DEPTH - mMapLevel - 1]
@@ -887,6 +887,10 @@ AdjustMemoryS (
   }
 
   Target = Start + Size - SizeRequested;
+  ASSERT (Target >= Start);
+  if (Target == 0) {
+    return 0;
+  }
 
   if (!IsGuardPage (Start + Size)) {
     // No Guard at tail to share. One more page is needed.
-- 
2.39.2