[pvmfw] Introduce FdtValidationError during the FDT validation am: abc7d63b9e am: 053c53a250

Original change: https://android-review.googlesource.com/c/platform/packages/modules/Virtualization/+/2624551

Change-Id: Ic1f48d6aa9be93bb881fef6990880e431e67fe65
Signed-off-by: Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com>
diff --git a/pvmfw/src/fdt.rs b/pvmfw/src/fdt.rs
index ab851a1..311f41f 100644
--- a/pvmfw/src/fdt.rs
+++ b/pvmfw/src/fdt.rs
@@ -24,6 +24,7 @@
 use core::cmp::max;
 use core::cmp::min;
 use core::ffi::CStr;
+use core::fmt;
 use core::mem::size_of;
 use core::ops::Range;
 use fdtpci::PciMemoryFlags;
@@ -43,6 +44,21 @@
 use vmbase::util::flatten;
 use vmbase::util::RangeExt as _;
 
+/// An enumeration of errors that can occur during the FDT validation.
+#[derive(Clone, Debug)]
+pub enum FdtValidationError {
+    /// Invalid CPU count.
+    InvalidCpuCount(usize),
+}
+
+impl fmt::Display for FdtValidationError {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            Self::InvalidCpuCount(num_cpus) => write!(f, "Invalid CPU count: {num_cpus}"),
+        }
+    }
+}
+
 /// Extract from /config the address range containing the pre-loaded kernel. Absence of /config is
 /// not an error.
 fn read_kernel_range_from(fdt: &Fdt) -> libfdt::Result<Option<Range<usize>>> {
@@ -140,16 +156,12 @@
 }
 
 /// Validate number of CPUs
-fn validate_num_cpus(num_cpus: usize) -> Result<(), RebootReason> {
-    if num_cpus == 0 {
-        error!("Number of CPU can't be 0");
-        return Err(RebootReason::InvalidFdt);
+fn validate_num_cpus(num_cpus: usize) -> Result<(), FdtValidationError> {
+    if num_cpus == 0 || DeviceTreeInfo::gic_patched_size(num_cpus).is_none() {
+        Err(FdtValidationError::InvalidCpuCount(num_cpus))
+    } else {
+        Ok(())
     }
-    if DeviceTreeInfo::GIC_REDIST_SIZE_PER_CPU.checked_mul(num_cpus.try_into().unwrap()).is_none() {
-        error!("Too many CPUs for gic: {}", num_cpus);
-        return Err(RebootReason::InvalidFdt);
-    }
-    Ok(())
 }
 
 /// Patch DT by keeping `num_cpus` number of arm,arm-v8 compatible nodes, and pruning the rest.
@@ -516,9 +528,8 @@
     let mut range1 = ranges.next().ok_or(FdtError::NotFound)?;
 
     let addr = range0.addr;
-    // SAFETY - doesn't overflow. checked in validate_num_cpus
-    let size: u64 =
-        DeviceTreeInfo::GIC_REDIST_SIZE_PER_CPU.checked_mul(num_cpus.try_into().unwrap()).unwrap();
+    // `validate_num_cpus()` checked that this wouldn't panic
+    let size = u64::try_from(DeviceTreeInfo::gic_patched_size(num_cpus).unwrap()).unwrap();
 
     // range1 is just below range0
     range1.addr = addr - size;
@@ -581,7 +592,11 @@
 }
 
 impl DeviceTreeInfo {
-    const GIC_REDIST_SIZE_PER_CPU: u64 = (32 * SIZE_4KB) as u64;
+    fn gic_patched_size(num_cpus: usize) -> Option<usize> {
+        const GIC_REDIST_SIZE_PER_CPU: usize = 32 * SIZE_4KB;
+
+        GIC_REDIST_SIZE_PER_CPU.checked_mul(num_cpus)
+    }
 }
 
 pub fn sanitize_device_tree(fdt: &mut Fdt) -> Result<DeviceTreeInfo, RebootReason> {
@@ -623,7 +638,10 @@
         error!("Failed to read num cpus from DT: {e}");
         RebootReason::InvalidFdt
     })?;
-    validate_num_cpus(num_cpus)?;
+    validate_num_cpus(num_cpus).map_err(|e| {
+        error!("Failed to validate num cpus from DT: {e}");
+        RebootReason::InvalidFdt
+    })?;
 
     let pci_info = read_pci_info_from(fdt).map_err(|e| {
         error!("Failed to read pci info from DT: {e}");