From f8bea4233199f3776452aa1a4b9217cf26483d37 Mon Sep 17 00:00:00 2001
From: Michael Kubacki <michael.kubacki@microsoft.com>
Date: Wed, 15 Nov 2023 19:00:28 -0500
Subject: [PATCH] ArmPkg/Drivers/CpuDxe: Check integer before conversion

GetNextEntryAttribute() assigns a 64-bit integer to 32-bit integers.
This change checks that the value fits in a 32-bit integer and
fixes the following Visual Studio compiler warning:

'=': conversion from 'UINT64' to 'UINT32', possible loss of data

Signed-off-by: Michael Kubacki <michael.kubacki@microsoft.com>
---
 ArmPkg/Drivers/CpuDxe/AArch64/Mmu.c | 27 +++++++++++++++++++++++----
 ArmPkg/Drivers/CpuDxe/CpuDxe.inf    |  1 +
 2 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/ArmPkg/Drivers/CpuDxe/AArch64/Mmu.c b/ArmPkg/Drivers/CpuDxe/AArch64/Mmu.c
index 7325c30164..4ac784eb2b 100644
--- a/ArmPkg/Drivers/CpuDxe/AArch64/Mmu.c
+++ b/ArmPkg/Drivers/CpuDxe/AArch64/Mmu.c
@@ -11,6 +11,8 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 --*/
 
 #include <Library/MemoryAllocationLib.h>
+#include <Library/SafeIntLib.h>         // MU_CHANGE: Convert integers safely
+#include <Chipset/AArch64Mmu.h>         // MU_CHANGE: Include header used in file
 #include "CpuDxe.h"
 
 #define INVALID_ENTRY  ((UINT32)~0)
@@ -148,7 +150,15 @@ GetNextEntryAttribute (
   // Get the memory space map from GCD
   MemorySpaceMap = NULL;
   Status         = gDS->GetMemorySpaceMap (&NumberOfDescriptors, &MemorySpaceMap);
-  ASSERT_EFI_ERROR (Status);
+  // MU_CHANGE [BEGIN]: Check if the memory space map is valid
+  if (EFI_ERROR (Status)) {
+    // This function needs to define what is returned when an error occurs.
+    // Callers need to actually check the return value and add error handling.
+    ASSERT_EFI_ERROR (Status);
+    return 0;
+  }
+
+  // MU_CHANGE [END]: Check if the memory space map is valid
 
   // We cannot get more than 3-level page table
   ASSERT (TableLevel <= 3);
@@ -156,9 +166,18 @@ GetNextEntryAttribute (
   // While the top level table might not contain TT_ENTRY_COUNT entries;
   // the subsequent ones should be filled up
   for (Index = 0; Index < EntryCount; Index++) {
-    Entry          = TableAddress[Index];
-    EntryType      = Entry & TT_TYPE_MASK;
-    EntryAttribute = Entry & TT_ATTRIBUTES_MASK;  // MU_CHANGE: Return all attributes from page table
+    Entry = TableAddress[Index];
+
+    // MU_CHANGE [BEGIN]: Convert integers safely
+    Status = SafeUint64ToUint32 (Entry, &EntryType);
+    if (EFI_ERROR (Status)) {
+      DEBUG ((DEBUG_ERROR, "[%a] - Table address entry exceeds 32-bit.\n", __func__));
+      return 0;
+    }
+
+    EntryAttribute = EntryType & TT_ATTRIBUTES_MASK;  // MU_CHANGE: Return all attributes from page table
+    EntryType     &= TT_TYPE_MASK;
+    // MU_CHANGE [END]: Convert integers safely
 
     // If Entry is a Table Descriptor type entry then go through the sub-level table
     if ((EntryType == TT_TYPE_BLOCK_ENTRY) ||
diff --git a/ArmPkg/Drivers/CpuDxe/CpuDxe.inf b/ArmPkg/Drivers/CpuDxe/CpuDxe.inf
index a91d7097e5..897b20fe6d 100644
--- a/ArmPkg/Drivers/CpuDxe/CpuDxe.inf
+++ b/ArmPkg/Drivers/CpuDxe/CpuDxe.inf
@@ -50,6 +50,7 @@
   HobLib
   MemoryAllocationLib
   PeCoffGetEntryPointLib
+  SafeIntLib                # MU_CHANGE: Convert integers safely
   UefiDriverEntryPoint
   UefiLib
   DxeMemoryProtectionHobLib # MU_CHANGE