From cee9d1b16bdcd7fe8bad6b39fe18af0a15941ca4 Mon Sep 17 00:00:00 2001
From: Kun Qin <kuqin@microsoft.com>
Date: Tue, 16 Jul 2024 15:50:01 -0700
Subject: [PATCH] MdeModulePkg: DxeCore: Fix Use-After-Free guard causing page
 fault

REF:https://bugzilla.tianocore.org/show_bug.cgi?id=2411

With Use-After-Free heap guard feature enabled, the DxeCore would blindly
attempt to "level-up" when the `GuardAllFreedPages` inspect a non-max
level table entry from the last loop. This could cause the next round of
inspection to dereference a potentially null pointer and as such causing
a page fault.

This change adds a null pointer check to prevent such case from happening.

Cc: Liming Gao <gaoliming@byosoft.com.cn>

Signed-off-by: Kun Qin <kun.qin@microsoft.com>
---
 MdeModulePkg/Core/Dxe/Mem/HeapGuard.c | 51 +++++++++++++++------------
 1 file changed, 28 insertions(+), 23 deletions(-)

diff --git a/MdeModulePkg/Core/Dxe/Mem/HeapGuard.c b/MdeModulePkg/Core/Dxe/Mem/HeapGuard.c
index 0c0ca61872..4071053036 100644
--- a/MdeModulePkg/Core/Dxe/Mem/HeapGuard.c
+++ b/MdeModulePkg/Core/Dxe/Mem/HeapGuard.c
@@ -1406,34 +1406,39 @@ GuardAllFreedPages (
       TableEntry = ((UINT64 *)(UINTN)(Tables[Level]))[Indices[Level]];
       Address    = Addresses[Level];
 
-      if (Level < GUARDED_HEAP_MAP_TABLE_DEPTH - 1) {
-        Level           += 1;
-        Tables[Level]    = TableEntry;
-        Addresses[Level] = Address;
-        Indices[Level]   = 0;
-
-        continue;
+      if (TableEntry == 0) {
+        GuardPageNumber = 0;
+        GuardPage       = (UINT64)-1;
       } else {
-        BitIndex = 1;
-        while (BitIndex != 0) {
-          if ((TableEntry & BitIndex) != 0) {
-            if (GuardPage == (UINT64)-1) {
-              GuardPage = Address;
+        if (Level < GUARDED_HEAP_MAP_TABLE_DEPTH - 1) {
+          Level           += 1;
+          Tables[Level]    = TableEntry;
+          Addresses[Level] = Address;
+          Indices[Level]   = 0;
+
+          continue;
+        } else {
+          BitIndex = 1;
+          while (BitIndex != 0) {
+            if ((TableEntry & BitIndex) != 0) {
+              if (GuardPage == (UINT64)-1) {
+                GuardPage = Address;
+              }
+
+              ++GuardPageNumber;
+            } else if (GuardPageNumber > 0) {
+              GuardFreedPages (GuardPage, GuardPageNumber);
+              GuardPageNumber = 0;
+              GuardPage       = (UINT64)-1;
             }
 
-            ++GuardPageNumber;
-          } else if (GuardPageNumber > 0) {
-            GuardFreedPages (GuardPage, GuardPageNumber);
-            GuardPageNumber = 0;
-            GuardPage       = (UINT64)-1;
-          }
+            if (TableEntry == 0) {
+              break;
+            }
 
-          if (TableEntry == 0) {
-            break;
+            Address += EFI_PAGES_TO_SIZE (1);
+            BitIndex = LShiftU64 (BitIndex, 1);
           }
-
-          Address += EFI_PAGES_TO_SIZE (1);
-          BitIndex = LShiftU64 (BitIndex, 1);
         }
       }
     }