MdeModulePkg: Compatibility Mode: Only Remap System Memory Regions

When we enter memory protections compatibility mode, we attempt to
disable null protection and remap 0 - 0xA0000 as RWX. This was done
for x86 systems with broken shim/grubs on Linux that would attempt
to use those regions. This resolved that issue and we could boot
non-memory protection safe Linux images on x86 HW. However, this
approach did not take into account systems that do not have that
range marked as system memory, for example ARM64 systems do not
have this requirement. As such, this would inappropriately map
these regions as RWX when they were not system memory.

This patch updates the remapping to only remap and disable null
protection if these ranges are marked as system memory, otherwise
it will leave them alone.

MdeModulePkg/Core/Dxe/DxeMain.inf

@@ -99,6 +99,7 @@
   ImagePropertiesRecordLib
   DxeMemoryProtectionHobLib        ## MU_CHANGE
   MemoryBinOverrideLib      ## MU_CHANGE
+  SafeIntLib                ## MU_CHANGE
 
 [Guids]
   gEfiEventMemoryMapChangeGuid                  ## PRODUCES             ## Event

MdeModulePkg/Core/Dxe/Misc/MemoryProtectionSupport.c

@@ -5,6 +5,7 @@
   SPDX-License-Identifier: BSD-2-Clause-Patent
 **/
 
+#include <Library/SafeIntLib.h>
 #include "MemoryProtectionSupport.h"
 
 IMAGE_PROPERTIES_PRIVATE_DATA  mImagePropertiesPrivate = {
@@ -66,6 +67,8 @@ BOOLEAN                        mGcdSyncComplete                = FALSE;
 // TRUE if A is a bitwise subset of B
 #define CHECK_SUBSET(A, B)  ((A | B) == B)
 
+#define LEGACY_BIOS_WB_LENGTH  0xA0000
+
 /**
   Return the section alignment requirement for the PE image section type.
 
@@ -2855,6 +2858,17 @@ DisableNullDetection (
     return;
   }
 
+  // Only re-enable the null page if it is system memory. If this page belongs to
+  // another memory type or is unmapped in general, leave it RP
+  if (Desc.GcdMemoryType != EfiGcdMemoryTypeSystemMemory) {
+    DEBUG ((
+      DEBUG_WARN,
+      "%a - Not disabling null detection as page 0 is not marked as system memory\n",
+      __func__
+      ));
+    return;
+  }
+
   Status = CoreSetMemorySpaceAttributes (
              0,
              EFI_PAGE_SIZE,
@@ -2941,21 +2955,122 @@ MapLegacyBiosMemoryRWX (
   VOID
   )
 {
-  EFI_STATUS  Status = EFI_SUCCESS;
+  EFI_STATUS                       Status = EFI_SUCCESS;
+  EFI_GCD_MEMORY_SPACE_DESCRIPTOR  Desc;
+  EFI_PHYSICAL_ADDRESS             Start = 0x0;
+  UINT64                           Length;
+  UINT64                           DescLengthFromStart;
 
-  // https://wiki.osdev.org/Memory_Map_(x86)
-  //
-  // Map the legacy BIOS write-back memory as RWX.
-  if (gCpu != NULL) {
+  if (gCpu == NULL) {
+    DEBUG ((
+      DEBUG_ERROR,
+      "%a cannot remap Legacy BIOS memory as RWX because gCpu is NULL\n",
+      __func__
+      ));
+    return;
+  }
+
+  // Ensure that this memory is marked as system memory. If it is not system memory, do not change
+  // the memory attributes as we do not want to map something that shouldn't be mapped or map something
+  // incorrectly
+  while (Start < LEGACY_BIOS_WB_LENGTH) {
+    Status = CoreGetMemorySpaceDescriptor (Start, &Desc);
+    if (EFI_ERROR (Status)) {
+      DEBUG ((
+        DEBUG_ERROR,
+        "%a - Failed to get memory space descriptor for address 0x%llx! Status = %r\n",
+        __func__,
+        Start,
+        Status
+        ));
+      return;
+    }
+
+    // find the length from Start to the end of this descriptor, in the case Start != Desc.BaseAddress
+    // these should all be well formed descriptors, but use SafeIntLib functions just to be sure we don't
+    // over/underflow
+    Status = SafeUint64Add (Desc.BaseAddress, Desc.Length, &DescLengthFromStart);
+    if (EFI_ERROR (Status)) {
+      DEBUG ((
+        DEBUG_ERROR,
+        "%a - Memory space descriptor has UINT64 overflowing address 0x%llx + length 0x%llx! Status = %r\n",
+        __func__,
+        Desc.BaseAddress,
+        Desc.Length,
+        Status
+        ));
+
+      // if we fail here this is very bad and means the GCD is malformed
+      ASSERT (FALSE);
+      return;
+    }
+
+    Status = SafeUint64Sub (DescLengthFromStart, Start, &DescLengthFromStart);
+    if (EFI_ERROR (Status)) {
+      DEBUG ((
+        DEBUG_ERROR,
+        "%a - Memory space descriptor has UINT64 underflowing address 0x%llx + length 0x%llx Start: 0x%llx! Status = %r\n",
+        __func__,
+        Desc.BaseAddress,
+        Desc.Length,
+        Start,
+        Status
+        ));
+
+      // if we fail here this is very bad and means the GCD is malformed
+      ASSERT (FALSE);
+      return;
+    }
+
+    // Ensure we only go up to LEGACY_BIOS_WB_LENGTH here, we know Start is less than it due to while condition
+    // We also know Start + DescLengthFromStart won't overflow because above we did a safe subtraction of Start from
+    // DescLengthFromStart
+    if (Start + DescLengthFromStart > LEGACY_BIOS_WB_LENGTH) {
+      Length = LEGACY_BIOS_WB_LENGTH - Start;
+    } else {
+      Length = DescLengthFromStart;
+    }
+
+    // remove this chunk from being remapped if it is not system memory
+    if (Desc.GcdMemoryType != EfiGcdMemoryTypeSystemMemory) {
+      DEBUG ((
+        DEBUG_WARN,
+        "%a Not mapping 0x%llx for 0x%llx as RWX because it is not system memory\n",
+        __func__,
+        Desc.BaseAddress,
+        Length
+        ));
+
+      // we know this doesn't overflow because we did a safe subtraction of Start from DescLengthFromStart above
+      Start += DescLengthFromStart;
+      continue;
+    }
+
+    // https://wiki.osdev.org/Memory_Map_(x86)
+    //
+    // Map the legacy BIOS write-back memory as RWX.
     Status = gCpu->SetMemoryAttributes (
                      gCpu,
-                     0x0,
-                     0xa0000,
+                     Start,
+                     Length,
                      0
                      );
-  }
 
-  ASSERT_EFI_ERROR (Status);
+    if (EFI_ERROR (Status)) {
+      DEBUG ((
+        DEBUG_ERROR,
+        "%a failed to map 0x%llx for length 0x%llx as RWX\n",
+        __func__,
+        Start,
+        Length
+        ));
+
+      ASSERT_EFI_ERROR (Status);
+    }
+
+    // we know this doesn't overflow because we did a safe subtraction of Start from DescLengthFromStart above
+    Start += DescLengthFromStart;
+  }
 }
 
 /**