[RFC] Use guest memory collection and other refactors

src/vmm/src/arch/aarch64/fdt.rs

@@ -9,7 +9,7 @@ use std::ffi::CString;
 use std::fmt::Debug;
 
 use vm_fdt::{Error as VmFdtError, FdtWriter, FdtWriterNode};
-use vm_memory::GuestMemoryError;
+use vm_memory::{GuestMemoryError, GuestMemoryRegion};
 
 use super::cache_info::{CacheEntry, read_cache_config};
 use super::gic::GICDevice;
@@ -22,7 +22,7 @@ use crate::device_manager::mmio::MMIODeviceInfo;
 use crate::device_manager::pci_mngr::PciDevices;
 use crate::devices::acpi::vmgenid::{VMGENID_MEM_SIZE, VmGenId};
 use crate::initrd::InitrdConfig;
-use crate::vstate::memory::{Address, GuestMemory, GuestMemoryMmap};
+use crate::vstate::memory::{Address, GuestMemory, GuestMemoryMmap, GuestRegionType};
 
 // This is a value for uniquely identifying the FDT node declaring the interrupt controller.
 const GIC_PHANDLE: u32 = 1;
@@ -227,14 +227,16 @@ fn create_memory_node(fdt: &mut FdtWriter, guest_mem: &GuestMemoryMmap) -> Resul
     // The reason we do this is that Linux does not allow remapping system memory. However, without
     // remap, kernel drivers cannot get virtual addresses to read data from device memory. Leaving
     // this memory region out allows Linux kernel modules to remap and thus read this region.
-    let mem_size = guest_mem.last_addr().raw_value()
-        - super::layout::DRAM_MEM_START
-        - super::layout::SYSTEM_MEM_SIZE
-        + 1;
-    let mem_reg_prop = &[
-        super::layout::DRAM_MEM_START + super::layout::SYSTEM_MEM_SIZE,
-        mem_size,
-    ];
+    // Pick the first (and only) memory region
+    let dram_region = guest_mem
+        .iter()
+        .find(|region| region.region_type == GuestRegionType::Dram)
+        .unwrap();
+    let start_addr = dram_region
+        .start_addr()
+        .unchecked_add(super::layout::SYSTEM_MEM_SIZE);
+    let mem_size = dram_region.last_addr().unchecked_offset_from(start_addr) + 1;
+    let mem_reg_prop = &[start_addr.raw_value(), mem_size];
     let mem = fdt.begin_node("memory@ram")?;
     fdt.property_string("device_type", "memory")?;
     fdt.property_array_u64("reg", mem_reg_prop)?;

src/vmm/src/arch/aarch64/mod.rs

@@ -22,15 +22,17 @@ use std::fs::File;
 
 use linux_loader::loader::pe::PE as Loader;
 use linux_loader::loader::{Cmdline, KernelLoader};
-use vm_memory::GuestMemoryError;
+use vm_memory::{GuestMemoryError, GuestMemoryRegion};
 
 use crate::arch::{BootProtocol, EntryPoint, arch_memory_regions_with_gap};
 use crate::cpu_config::aarch64::{CpuConfiguration, CpuConfigurationError};
 use crate::cpu_config::templates::CustomCpuTemplate;
 use crate::initrd::InitrdConfig;
 use crate::utils::{align_up, u64_to_usize, usize_to_u64};
 use crate::vmm_config::machine_config::MachineConfig;
-use crate::vstate::memory::{Address, Bytes, GuestAddress, GuestMemory, GuestMemoryMmap};
+use crate::vstate::memory::{
+    Address, Bytes, GuestAddress, GuestMemory, GuestMemoryMmap, GuestRegionType,
+};
 use crate::vstate::vcpu::KvmVcpuError;
 use crate::{DeviceManager, Kvm, Vcpu, VcpuConfig, Vm, logger};
 
@@ -151,31 +153,29 @@ pub fn initrd_load_addr(guest_mem: &GuestMemoryMmap, initrd_size: usize) -> Opti
         usize_to_u64(initrd_size),
         usize_to_u64(super::GUEST_PAGE_SIZE),
     );
-    match GuestAddress(get_fdt_addr(guest_mem)).checked_sub(rounded_size) {
-        Some(offset) => {
-            if guest_mem.address_in_range(offset) {
-                Some(offset.raw_value())
-            } else {
-                None
-            }
-        }
-        None => None,
-    }
+    GuestAddress(get_fdt_addr(guest_mem))
+        .checked_sub(rounded_size)
+        .filter(|&addr| guest_mem.address_in_range(addr))
+        .map(|addr| addr.raw_value())
 }
 
 // Auxiliary function to get the address where the device tree blob is loaded.
 fn get_fdt_addr(mem: &GuestMemoryMmap) -> u64 {
+    // Find the first (and only) DRAM region.
+    let dram_region = mem
+        .iter()
+        .find(|region| region.region_type == GuestRegionType::Dram)
+        .unwrap();
+
     // If the memory allocated is smaller than the size allocated for the FDT,
     // we return the start of the DRAM so that
     // we allow the code to try and load the FDT.
-
-    if let Some(addr) = mem.last_addr().checked_sub(layout::FDT_MAX_SIZE as u64 - 1)
-        && mem.address_in_range(addr)
-    {
-        return addr.raw_value();
-    }
-
-    layout::DRAM_MEM_START
+    dram_region
+        .last_addr()
+        .checked_sub(layout::FDT_MAX_SIZE as u64 - 1)
+        .filter(|&addr| mem.address_in_range(addr))
+        .map(|addr| addr.raw_value())
+        .unwrap_or(layout::DRAM_MEM_START)
 }
 
 /// Load linux kernel into guest memory.