[PATCH] MdeModulePkg/Core: fix guard page missing issue

[Jian J Wang <jian.j.wang@intel.com>]

This issue is a regression one caused by a patch at

    425d25699be83c35e12df8470b827d7fbcef3bce

That fix didn't take the 0 page to free into account, which still
needs to call UnsetGuardPage() even no memory needs to free.

The fix is just moving the calling of UnsetGuardPage() to the place
right after calling AdjustMemoryF().

Cc: Ruiyu Ni <ruiyu.ni@intel.com>
Cc: Star Zeng <star.zeng@intel.com>
Cc: Eric Dong <eric.dong@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.1
Signed-off-by: Jian J Wang <jian.j.wang@intel.com>
---
 MdeModulePkg/Core/Dxe/Mem/HeapGuard.c |  7 +++----
 MdeModulePkg/Core/Dxe/Mem/Pool.c      | 16 ++++++++--------
 2 files changed, 11 insertions(+), 12 deletions(-)

diff --git a/MdeModulePkg/Core/Dxe/Mem/HeapGuard.c b/MdeModulePkg/Core/Dxe/Mem/HeapGuard.c
index 92753c7269..392aeb8a02 100644
--- a/MdeModulePkg/Core/Dxe/Mem/HeapGuard.c
+++ b/MdeModulePkg/Core/Dxe/Mem/HeapGuard.c
@@ -1135,10 +1135,6 @@ CoreConvertPagesWithGuard (
     OldPages = NumberOfPages;
 
     AdjustMemoryF (&Start, &NumberOfPages);
-    if (NumberOfPages == 0) {
-      return EFI_SUCCESS;
-    }
-
     //
     // It's safe to unset Guard page inside memory lock because there should
     // be no memory allocation occurred in updating memory page attribute at
@@ -1147,6 +1143,9 @@ CoreConvertPagesWithGuard (
     // marking it usable (from non-present to present).
     //
     UnsetGuardForMemory (OldStart, OldPages);
+    if (NumberOfPages == 0) {
+      return EFI_SUCCESS;
+    }
   } else {
     AdjustMemoryA (&Start, &NumberOfPages);
   }
diff --git a/MdeModulePkg/Core/Dxe/Mem/Pool.c b/MdeModulePkg/Core/Dxe/Mem/Pool.c
index df9a1d28df..1ff2061f7f 100644
--- a/MdeModulePkg/Core/Dxe/Mem/Pool.c
+++ b/MdeModulePkg/Core/Dxe/Mem/Pool.c
@@ -642,15 +642,15 @@ CoreFreePoolPagesWithGuard (
   NoPagesGuarded = NoPages;
 
   AdjustMemoryF (&Memory, &NoPages);
+  //
+  // It's safe to unset Guard page inside memory lock because there should
+  // be no memory allocation occurred in updating memory page attribute at
+  // this point. And unsetting Guard page before free will prevent Guard
+  // page just freed back to pool from being allocated right away before
+  // marking it usable (from non-present to present).
+  //
+  UnsetGuardForMemory (MemoryGuarded, NoPagesGuarded);
   if (NoPages > 0) {
-    //
-    // It's safe to unset Guard page inside memory lock because there should
-    // be no memory allocation occurred in updating memory page attribute at
-    // this point. And unsetting Guard page before free will prevent Guard
-    // page just freed back to pool from being allocated right away before
-    // marking it usable (from non-present to present).
-    //
-    UnsetGuardForMemory (MemoryGuarded, NoPagesGuarded);
     CoreFreePoolPagesI (PoolType, Memory, NoPages);
   }
 }
-- 
2.14.1.windows.1