From fc6cfb3472030a16ddfcb485eb505e2cd3b7fa16 Mon Sep 17 00:00:00 2001
From: jackabald <jwarchibald@wisc.edu>
Date: Wed, 31 Jul 2024 00:23:47 -0500
Subject: [PATCH] write fdt to memory in configure_system

From: jackabald <jwarchibald@wisc.edu>

- removes writing to memory in create_fdt
- add functionality of writing to memory in configure_system
- uses previously unused return value of create_fdt to write fdt to
  memory

Signed-off-by: jackabald <jwarchibald@wisc.edu>
Signed-off-by: Patrick Roy <roypat@amazon.co.uk>
---
 src/vmm/src/arch/aarch64/fdt.rs | 12 ++++--------
 src/vmm/src/arch/aarch64/mod.rs | 16 ++++++++++++----
 2 files changed, 16 insertions(+), 12 deletions(-)

diff --git a/src/vmm/src/arch/aarch64/fdt.rs b/src/vmm/src/arch/aarch64/fdt.rs
index 71dc401fb32b..82963d1b0862 100644
--- a/src/vmm/src/arch/aarch64/fdt.rs
+++ b/src/vmm/src/arch/aarch64/fdt.rs
@@ -14,10 +14,9 @@ use vm_memory::GuestMemoryError;
 
 use super::super::{DeviceType, InitrdConfig};
 use super::cache_info::{read_cache_config, CacheEntry};
-use super::get_fdt_addr;
 use super::gic::GICDevice;
 use crate::devices::acpi::vmgenid::{VmGenId, VMGENID_MEM_SIZE};
-use crate::vstate::memory::{Address, Bytes, GuestAddress, GuestMemory, GuestMemoryMmap};
+use crate::vstate::memory::{Address, GuestMemory, GuestMemoryMmap};
 
 // This is a value for uniquely identifying the FDT node declaring the interrupt controller.
 const GIC_PHANDLE: u32 = 1;
@@ -65,11 +64,11 @@ pub enum FdtError {
 }
 
 /// Creates the flattened device tree for this aarch64 microVM.
-pub fn create_fdt<T: DeviceInfoForFDT + Clone + Debug, S: std::hash::BuildHasher>(
+pub fn create_fdt<T: DeviceInfoForFDT + Clone + Debug>(
     guest_mem: &GuestMemoryMmap,
     vcpu_mpidr: Vec<u64>,
     cmdline: CString,
-    device_info: &HashMap<(DeviceType, String), T, S>,
+    device_info: &HashMap<(DeviceType, String), T>,
     gic_device: &GICDevice,
     vmgenid: &Option<VmGenId>,
     initrd: &Option<InitrdConfig>,
@@ -106,10 +105,6 @@ pub fn create_fdt<T: DeviceInfoForFDT + Clone + Debug, S: std::hash::BuildHasher
 
     // Allocate another buffer so we can format and then write fdt to guest.
     let fdt_final = fdt_writer.finish()?;
-
-    // Write FDT to memory.
-    let fdt_address = GuestAddress(get_fdt_addr(guest_mem));
-    guest_mem.write_slice(fdt_final.as_slice(), fdt_address)?;
     Ok(fdt_final)
 }
 
@@ -461,6 +456,7 @@ mod tests {
     use crate::arch::aarch64::layout;
     use crate::device_manager::resources::ResourceAllocator;
     use crate::test_utils::arch_mem;
+    use crate::vmstate::memory::GuestAddress;
 
     const LEN: u64 = 4096;
 
diff --git a/src/vmm/src/arch/aarch64/mod.rs b/src/vmm/src/arch/aarch64/mod.rs
index 0cb6f1ed9ba8..17ff63f5b879 100644
--- a/src/vmm/src/arch/aarch64/mod.rs
+++ b/src/vmm/src/arch/aarch64/mod.rs
@@ -17,11 +17,13 @@ use std::collections::HashMap;
 use std::ffi::CString;
 use std::fmt::Debug;
 
+use vm_memory::GuestMemoryError;
+
 pub use self::fdt::DeviceInfoForFDT;
 use self::gic::GICDevice;
 use crate::arch::DeviceType;
 use crate::devices::acpi::vmgenid::VmGenId;
-use crate::vstate::memory::{Address, GuestAddress, GuestMemory, GuestMemoryMmap};
+use crate::vstate::memory::{Address, Bytes, GuestAddress, GuestMemory, GuestMemoryMmap};
 
 /// Errors thrown while configuring aarch64 system.
 #[derive(Debug, thiserror::Error, displaydoc::Display)]
@@ -30,6 +32,8 @@ pub enum ConfigurationError {
     SetupFDT(#[from] fdt::FdtError),
     /// Failed to compute the initrd address.
     InitrdAddress,
+    /// Failed to write to guest memory.
+    MemoryError(GuestMemoryError),
 }
 
 /// The start of the memory area reserved for MMIO devices.
@@ -55,16 +59,16 @@ pub fn arch_memory_regions(size: usize) -> Vec<(GuestAddress, usize)> {
 /// * `device_info` - A hashmap containing the attached devices for building FDT device nodes.
 /// * `gic_device` - The GIC device.
 /// * `initrd` - Information about an optional initrd.
-pub fn configure_system<T: DeviceInfoForFDT + Clone + Debug, S: std::hash::BuildHasher>(
+pub fn configure_system<T: DeviceInfoForFDT + Clone + Debug>(
     guest_mem: &GuestMemoryMmap,
     cmdline_cstring: CString,
     vcpu_mpidr: Vec<u64>,
-    device_info: &HashMap<(DeviceType, String), T, S>,
+    device_info: &HashMap<(DeviceType, String), T>,
     gic_device: &GICDevice,
     vmgenid: &Option<VmGenId>,
     initrd: &Option<super::InitrdConfig>,
 ) -> Result<(), ConfigurationError> {
-    fdt::create_fdt(
+    let fdt = fdt::create_fdt(
         guest_mem,
         vcpu_mpidr,
         cmdline_cstring,
@@ -73,6 +77,10 @@ pub fn configure_system<T: DeviceInfoForFDT + Clone + Debug, S: std::hash::Build
         vmgenid,
         initrd,
     )?;
+    let fdt_address = GuestAddress(get_fdt_addr(guest_mem));
+    guest_mem
+        .write_slice(fdt.as_slice(), fdt_address)
+        .map_err(ConfigurationError::MemoryError)?;
     Ok(())
 }