[PATCH 3/4] ArmPkg: Fix Unsafe ASSERTs in MMU Logic

["Taylor Beebe" <t@taylorbeebe.com>]

From: Taylor Beebe <tabeebe@microsoft.com>

There are ASSERTs present in the MMU logic to ensure various
functions return successfully, but these ASSERTs may be ignored
on release builds causing unsafe behavior. This patch updates
the logic to handle unexpected return values and branch safely.

Cc: Leif Lindholm <quic_llindhol@quicinc.com>
Cc: Ard Biesheuvel <ardb+tianocore@kernel.org>

Signed-off-by: Taylor Beebe <t@taylorbeebe.com>
---
 ArmPkg/Drivers/CpuDxe/AArch64/Mmu.c | 21 ++++++++++++-----
 ArmPkg/Drivers/CpuDxe/Arm/Mmu.c     | 36 ++++++++++++++++++++++++-----
 2 files changed, 45 insertions(+), 12 deletions(-)

diff --git a/ArmPkg/Drivers/CpuDxe/AArch64/Mmu.c b/ArmPkg/Drivers/CpuDxe/AArch64/Mmu.c
index 0d3bc2809682..d9d386dbed6b 100644
--- a/ArmPkg/Drivers/CpuDxe/AArch64/Mmu.c
+++ b/ArmPkg/Drivers/CpuDxe/AArch64/Mmu.c
@@ -148,10 +148,12 @@ GetNextEntryAttribute (
   // Get the memory space map from GCD
   MemorySpaceMap = NULL;
   Status         = gDS->GetMemorySpaceMap (&NumberOfDescriptors, &MemorySpaceMap);
-  ASSERT_EFI_ERROR (Status);
 
-  // We cannot get more than 3-level page table
-  ASSERT (TableLevel <= 3);
+  if (EFI_ERROR (Status) || (TableLevel > 3)) {
+    ASSERT_EFI_ERROR (Status);
+    ASSERT (TableLevel <= 3);
+    return 0;
+  }
 
   // While the top level table might not contain TT_ENTRY_COUNT entries;
   // the subsequent ones should be filled up
@@ -243,7 +245,11 @@ SyncCacheConfig (
   //
   MemorySpaceMap = NULL;
   Status         = gDS->GetMemorySpaceMap (&NumberOfDescriptors, &MemorySpaceMap);
-  ASSERT_EFI_ERROR (Status);
+
+  if (EFI_ERROR (Status)) {
+    ASSERT_EFI_ERROR (Status);
+    return Status;
+  }
 
   // The GCD implementation maintains its own copy of the state of memory space attributes.  GCD needs
   // to know what the initial memory space attributes are.  The CPU Arch. Protocol does not provide a
@@ -277,7 +283,7 @@ SyncCacheConfig (
                            );
 
   // Update GCD with the last region if valid
-  if (PageAttribute != INVALID_ENTRY) {
+  if ((PageAttribute != INVALID_ENTRY) && (EndAddressGcdRegion > BaseAddressGcdRegion)) {
     SetGcdMemorySpaceAttributes (
       MemorySpaceMap,
       NumberOfDescriptors,
@@ -430,7 +436,10 @@ GetMemoryRegion (
   UINTN       EntryCount;
   UINTN       T0SZ;
 
-  ASSERT ((BaseAddress != NULL) && (RegionLength != NULL) && (RegionAttributes != NULL));
+  if ((BaseAddress == NULL) || (RegionLength == NULL) || (RegionAttributes == NULL)) {
+    ASSERT ((BaseAddress != NULL) && (RegionLength != NULL) && (RegionAttributes != NULL));
+    return EFI_INVALID_PARAMETER;
+  }
 
   TranslationTable = ArmGetTTBR0BaseAddress ();
 
diff --git a/ArmPkg/Drivers/CpuDxe/Arm/Mmu.c b/ArmPkg/Drivers/CpuDxe/Arm/Mmu.c
index 268c0bf3f9ae..5a2f36d06086 100644
--- a/ArmPkg/Drivers/CpuDxe/Arm/Mmu.c
+++ b/ArmPkg/Drivers/CpuDxe/Arm/Mmu.c
@@ -217,7 +217,10 @@ SyncCacheConfigPage (
       } else if (PageAttributes != NextPageAttributes) {
         // Convert Section Attributes into GCD Attributes
         Status = PageToGcdAttributes (NextPageAttributes, &GcdAttributes);
-        ASSERT_EFI_ERROR (Status);
+        if (EFI_ERROR (Status)) {
+          ASSERT_EFI_ERROR (Status);
+          GcdAttributes = 0;
+        }
 
         // update GCD with these changes (this will recurse into our own CpuSetMemoryAttributes below which is OK)
         SetGcdMemorySpaceAttributes (MemorySpaceMap, NumberOfDescriptors, *NextRegionBase, *NextRegionLength, GcdAttributes);
@@ -230,7 +233,10 @@ SyncCacheConfigPage (
     } else if (NextPageAttributes != 0) {
       // Convert Page Attributes into GCD Attributes
       Status = PageToGcdAttributes (NextPageAttributes, &GcdAttributes);
-      ASSERT_EFI_ERROR (Status);
+      if (EFI_ERROR (Status)) {
+        ASSERT_EFI_ERROR (Status);
+        GcdAttributes = 0;
+      }
 
       // update GCD with these changes (this will recurse into our own CpuSetMemoryAttributes below which is OK)
       SetGcdMemorySpaceAttributes (MemorySpaceMap, NumberOfDescriptors, *NextRegionBase, *NextRegionLength, GcdAttributes);
@@ -278,7 +284,12 @@ SyncCacheConfig (
   //
   MemorySpaceMap = NULL;
   Status         = gDS->GetMemorySpaceMap (&NumberOfDescriptors, &MemorySpaceMap);
-  ASSERT_EFI_ERROR (Status);
+
+  if (EFI_ERROR (Status)) {
+    DEBUG ((DEBUG_ERROR, "SyncCacheConfig - GetMemorySpaceMap() failed! Status: %r\n", Status));
+    ASSERT_EFI_ERROR (Status);
+    return Status;
+  }
 
   // The GCD implementation maintains its own copy of the state of memory space attributes.  GCD needs
   // to know what the initial memory space attributes are.  The CPU Arch. Protocol does not provide a
@@ -307,7 +318,12 @@ SyncCacheConfig (
       } else if (SectionAttributes != NextSectionAttributes) {
         // Convert Section Attributes into GCD Attributes
         Status = SectionToGcdAttributes (NextSectionAttributes, &GcdAttributes);
-        ASSERT_EFI_ERROR (Status);
+
+        if (EFI_ERROR (Status)) {
+          DEBUG ((DEBUG_ERROR, "SyncCacheConfig - SectionToGcdAttributes() failed! Status: %r\n", Status));
+          ASSERT_EFI_ERROR (Status);
+          GcdAttributes = 0;
+        }
 
         // update GCD with these changes (this will recurse into our own CpuSetMemoryAttributes below which is OK)
         SetGcdMemorySpaceAttributes (MemorySpaceMap, NumberOfDescriptors, NextRegionBase, NextRegionLength, GcdAttributes);
@@ -343,7 +359,11 @@ SyncCacheConfig (
       if (NextSectionAttributes != 0) {
         // Convert Section Attributes into GCD Attributes
         Status = SectionToGcdAttributes (NextSectionAttributes, &GcdAttributes);
-        ASSERT_EFI_ERROR (Status);
+        if (EFI_ERROR (Status)) {
+          DEBUG ((DEBUG_ERROR, "SyncCacheConfig - SectionToGcdAttributes() failed! Status: %r\n", Status));
+          ASSERT_EFI_ERROR (Status);
+          GcdAttributes = 0;
+        }
 
         // update GCD with these changes (this will recurse into our own CpuSetMemoryAttributes below which is OK)
         SetGcdMemorySpaceAttributes (MemorySpaceMap, NumberOfDescriptors, NextRegionBase, NextRegionLength, GcdAttributes);
@@ -360,7 +380,11 @@ SyncCacheConfig (
   if (NextSectionAttributes != 0) {
     // Convert Section Attributes into GCD Attributes
     Status = SectionToGcdAttributes (NextSectionAttributes, &GcdAttributes);
-    ASSERT_EFI_ERROR (Status);
+    if (EFI_ERROR (Status)) {
+      DEBUG ((DEBUG_ERROR, "SyncCacheConfig - SectionToGcdAttributes() failed! Status: %r\n", Status));
+      ASSERT_EFI_ERROR (Status);
+      GcdAttributes = 0;
+    }
 
     // update GCD with these changes (this will recurse into our own CpuSetMemoryAttributes below which is OK)
     SetGcdMemorySpaceAttributes (MemorySpaceMap, NumberOfDescriptors, NextRegionBase, NextRegionLength, GcdAttributes);
-- 
2.41.0.windows.1