Merge "[pvmfw] Separate SWIOTLB info parsing and validation properly"
diff --git a/pvmfw/src/fdt.rs b/pvmfw/src/fdt.rs
index d4c8385..738675c 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.
@@ -521,9 +533,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;
@@ -586,7 +597,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> {
@@ -628,7 +643,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}");
diff --git a/vmbase/example/src/pci.rs b/vmbase/example/src/pci.rs
index 41a3ff4..384a9c1 100644
--- a/vmbase/example/src/pci.rs
+++ b/vmbase/example/src/pci.rs
@@ -15,7 +15,7 @@
 //! Functions to scan the PCI bus for VirtIO device.
 
 use aarch64_paging::paging::MemoryRegion;
-use alloc::alloc::{alloc, dealloc, handle_alloc_error, Layout};
+use alloc::alloc::{alloc_zeroed, dealloc, handle_alloc_error, Layout};
 use core::{mem::size_of, ptr::NonNull};
 use fdtpci::PciInfo;
 use log::{debug, info};
@@ -103,7 +103,7 @@
         debug!("dma_alloc: pages={}", pages);
         let layout = Layout::from_size_align(pages * PAGE_SIZE, PAGE_SIZE).unwrap();
         // Safe because the layout has a non-zero size.
-        let vaddr = unsafe { alloc(layout) };
+        let vaddr = unsafe { alloc_zeroed(layout) };
         let vaddr =
             if let Some(vaddr) = NonNull::new(vaddr) { vaddr } else { handle_alloc_error(layout) };
         let paddr = virt_to_phys(vaddr);
diff --git a/vmbase/src/virtio/hal.rs b/vmbase/src/virtio/hal.rs
index ac5b967..36f9e56 100644
--- a/vmbase/src/virtio/hal.rs
+++ b/vmbase/src/virtio/hal.rs
@@ -42,13 +42,14 @@
     /// `dma_alloc` ensures the returned DMA buffer is not aliased with any other allocation or
     /// reference in the program until it is deallocated by `dma_dealloc` by allocating a unique
     /// block of memory using `alloc_shared`, which is guaranteed to allocate valid and unique
-    /// memory. We request an alignment of at least `PAGE_SIZE` from `alloc_shared`.
+    /// memory. We request an alignment of at least `PAGE_SIZE` from `alloc_shared`. We zero the
+    /// buffer before returning it.
     fn dma_alloc(pages: usize, _direction: BufferDirection) -> (PhysAddr, NonNull<u8>) {
-        let vaddr = alloc_shared(dma_layout(pages))
-            .expect("Failed to allocate and share VirtIO DMA range with host");
-        // TODO(ptosi): Move this zeroing to virtio_drivers, if it silently wants a zeroed region.
+        let layout = dma_layout(pages);
+        let vaddr =
+            alloc_shared(layout).expect("Failed to allocate and share VirtIO DMA range with host");
         // SAFETY - vaddr points to a region allocated for the caller so is safe to access.
-        unsafe { core::ptr::write_bytes(vaddr.as_ptr(), 0, dma_layout(pages).size()) };
+        unsafe { core::ptr::write_bytes(vaddr.as_ptr(), 0, layout.size()) };
         let paddr = virt_to_phys(vaddr);
         (paddr, vaddr)
     }