From 7de570cb70c61ec653b01e6f8c4af985ca32284e Mon Sep 17 00:00:00 2001
From: Taylor Beebe <taylor.d.beebe@gmail.com>
Date: Fri, 26 Apr 2024 15:37:28 -0700
Subject: [PATCH] Return if Allocation Fails When Manipulating Guarded Memory
 Map

Description

This PR updates a set of ASSERTs checking if memory was successfully
allocated to branch if the allocation failed in addition to ASSERTing.

- [x] Impacts functionality?
  - **Functionality** - Does the change ultimately impact how firmware functions?
  - Examples: Add a new library, publish a new PPI, update an algorithm, ...
- [ ] Impacts security?
  - **Security** - Does the change have a direct security impact on an application,
    flow, or firmware?
  - Examples: Crypto algorithm change, buffer overflow fix, parameter
    validation improvement, ...
- [ ] Breaking change?
  - **Breaking change** - Will anyone consuming this change experience a break
    in build or boot behavior?
  - Examples: Add a new library class, move a module to a different repo, call
    a function in a new library class in a pre-existing module, ...
- [ ] Includes tests?
  - **Tests** - Does the change include any explicit test code?
  - Examples: Unit tests, integration tests, robot tests, ...
- [ ] Includes documentation?
  - **Documentation** - Does the change contain explicit documentation additions
    outside direct code modifications (and comments)?
  - Examples: Update readme file, add feature readme file, link to documentation
    on an a separate Web page, ...

How This Was Tested

Tested by booting Q35 to shell

Integration Instructions

N/A
---
 MdeModulePkg/Core/Dxe/Mem/HeapGuard.c | 20 ++++++++++++++++----
 1 file changed, 16 insertions(+), 4 deletions(-)

diff --git a/MdeModulePkg/Core/Dxe/Mem/HeapGuard.c b/MdeModulePkg/Core/Dxe/Mem/HeapGuard.c
index 3a85c38d121..42ee9300c71 100644
--- a/MdeModulePkg/Core/Dxe/Mem/HeapGuard.c
+++ b/MdeModulePkg/Core/Dxe/Mem/HeapGuard.c
@@ -252,8 +252,14 @@ FindGuardedMemoryMap (
                  &MapMemory,
                  FALSE
                  );
-      ASSERT_EFI_ERROR (Status);
-      ASSERT (MapMemory != 0);
+      // MU_CHANGE START: Check if memory was successfully allocated
+      if (EFI_ERROR (Status) || (MapMemory == 0)) {
+        ASSERT_EFI_ERROR (Status);
+        ASSERT (MapMemory != 0);
+        return 0;
+      }
+
+      // MU_CHANGE END
       // MU_CHANGE START: Apply Protection policy to the allocated memory
       ApplyMemoryProtectionPolicy (
         EfiConventionalMemory,
@@ -291,8 +297,14 @@ FindGuardedMemoryMap (
                  &MapMemory,
                  FALSE
                  );
-      ASSERT_EFI_ERROR (Status);
-      ASSERT (MapMemory != 0);
+      // MU_CHANGE START: Check if memory was successfully allocated
+      if (EFI_ERROR (Status) || (MapMemory == 0)) {
+        ASSERT_EFI_ERROR (Status);
+        ASSERT (MapMemory != 0);
+        return 0;
+      }
+
+      // MU_CHANGE END
       // MU_CHANGE START: Apply Protection policy to the allocated memory
       ApplyMemoryProtectionPolicy (
         EfiConventionalMemory,