Patch the template DT

The original DT from crosvm is now overwritten by the template DT, which
is then patched.

add_dice_node is modified to patch_dice_node as google,open-dice node is
already in the template DT and therefore we don't need to (actually
can't) add new one.

Bug: 249054080
Test: TH
Change-Id: Ie1c93b56af6afd1a7844478f514bc42db3b2cff0

libs/libfdt/src/lib.rs

@@ -21,6 +21,7 @@ mod iterators;
 
 pub use iterators::{AddressRange, CellIterator, MemRegIterator, RangesIterator, Reg, RegIterator};
 
+use core::cmp::max;
 use core::ffi::{c_int, c_void, CStr};
 use core::fmt;
 use core::mem;
@@ -623,6 +624,21 @@ impl Fdt {
         mem::transmute::<&mut [u8], &mut Self>(fdt)
     }
 
+    /// Update this FDT from a slice containing another FDT
+    pub fn copy_from_slice(&mut self, new_fdt: &[u8]) -> Result<()> {
+        if self.buffer.len() < new_fdt.len() {
+            Err(FdtError::NoSpace)
+        } else {
+            let totalsize = self.totalsize();
+            self.buffer[..new_fdt.len()].clone_from_slice(new_fdt);
+            // Zeroize the remaining part. We zeroize up to the size of the original DT because
+            // zeroizing the entire buffer (max 2MB) is not necessary and may increase the VM boot
+            // time.
+            self.buffer[new_fdt.len()..max(new_fdt.len(), totalsize)].fill(0_u8);
+            Ok(())
+        }
+    }
+
     /// Make the whole slice containing the DT available to libfdt.
     pub fn unpack(&mut self) -> Result<()> {
         // SAFETY - "Opens" the DT in-place (supported use-case) by updating its header and

pvmfw/src/fdt.rs

@@ -19,6 +19,7 @@ use crate::helpers::flatten;
 use crate::helpers::GUEST_PAGE_SIZE;
 use crate::helpers::SIZE_4KB;
 use crate::RebootReason;
+use alloc::ffi::CString;
 use core::ffi::CStr;
 use core::mem::size_of;
 use core::ops::Range;
@@ -76,6 +77,24 @@ fn patch_initrd_range(fdt: &mut Fdt, initrd_range: &Range<usize>) -> libfdt::Res
     Ok(())
 }
 
+fn read_bootargs_from(fdt: &Fdt) -> libfdt::Result<Option<CString>> {
+    if let Some(chosen) = fdt.chosen()? {
+        if let Some(bootargs) = chosen.getprop_str(cstr!("bootargs"))? {
+            // We need to copy the string to heap because the original fdt will be invalidated
+            // by the templated DT
+            let copy = CString::new(bootargs.to_bytes()).map_err(|_| FdtError::BadValue)?;
+            return Ok(Some(copy));
+        }
+    }
+    Ok(None)
+}
+
+fn patch_bootargs(fdt: &mut Fdt, bootargs: &CStr) -> libfdt::Result<()> {
+    let mut node = fdt.chosen_mut()?.ok_or(FdtError::NotFound)?;
+    // TODO(b/275306568) filter out dangerous options
+    node.setprop(cstr!("bootargs"), bootargs.to_bytes_with_nul())
+}
+
 /// Read the first range in /memory node in DT
 fn read_memory_range_from(fdt: &Fdt) -> libfdt::Result<Range<usize>> {
     fdt.memory()?.ok_or(FdtError::NotFound)?.next().ok_or(FdtError::NotFound)
@@ -490,6 +509,7 @@ pub struct DeviceTreeInfo {
     pub kernel_range: Option<Range<usize>>,
     pub initrd_range: Option<Range<usize>>,
     pub memory_range: Range<usize>,
+    bootargs: Option<CString>,
     num_cpus: usize,
     pci_info: PciInfo,
     serial_info: SerialInfo,
@@ -505,7 +525,11 @@ pub fn sanitize_device_tree(fdt: &mut Fdt) -> Result<DeviceTreeInfo, RebootReaso
     let info = parse_device_tree(fdt)?;
     debug!("Device tree info: {:?}", info);
 
-    // TODO: replace fdt with the template DT
+    fdt.copy_from_slice(pvmfw_fdt_template::RAW).map_err(|e| {
+        error!("Failed to instantiate FDT from the template DT: {e}");
+        RebootReason::InvalidFdt
+    })?;
+
     patch_device_tree(fdt, &info)?;
     Ok(info)
 }
@@ -527,6 +551,11 @@ fn parse_device_tree(fdt: &libfdt::Fdt) -> Result<DeviceTreeInfo, RebootReason>
     })?;
     validate_memory_range(&memory_range)?;
 
+    let bootargs = read_bootargs_from(fdt).map_err(|e| {
+        error!("Failed to read bootargs from DT: {e}");
+        RebootReason::InvalidFdt
+    })?;
+
     let num_cpus = read_num_cpus_from(fdt).map_err(|e| {
         error!("Failed to read num cpus from DT: {e}");
         RebootReason::InvalidFdt
@@ -554,6 +583,7 @@ fn parse_device_tree(fdt: &libfdt::Fdt) -> Result<DeviceTreeInfo, RebootReason>
         kernel_range,
         initrd_range,
         memory_range,
+        bootargs,
         num_cpus,
         pci_info,
         serial_info,
@@ -562,6 +592,11 @@ fn parse_device_tree(fdt: &libfdt::Fdt) -> Result<DeviceTreeInfo, RebootReason>
 }
 
 fn patch_device_tree(fdt: &mut Fdt, info: &DeviceTreeInfo) -> Result<(), RebootReason> {
+    fdt.unpack().map_err(|e| {
+        error!("Failed to unpack DT for patching: {e}");
+        RebootReason::InvalidFdt
+    })?;
+
     if let Some(initrd_range) = &info.initrd_range {
         patch_initrd_range(fdt, initrd_range).map_err(|e| {
             error!("Failed to patch initrd range to DT: {e}");
@@ -572,6 +607,12 @@ fn patch_device_tree(fdt: &mut Fdt, info: &DeviceTreeInfo) -> Result<(), RebootR
         error!("Failed to patch memory range to DT: {e}");
         RebootReason::InvalidFdt
     })?;
+    if let Some(bootargs) = &info.bootargs {
+        patch_bootargs(fdt, bootargs.as_c_str()).map_err(|e| {
+            error!("Failed to patch bootargs to DT: {e}");
+            RebootReason::InvalidFdt
+        })?;
+    }
     patch_num_cpus(fdt, info.num_cpus).map_err(|e| {
         error!("Failed to patch cpus to DT: {e}");
         RebootReason::InvalidFdt
@@ -596,6 +637,12 @@ fn patch_device_tree(fdt: &mut Fdt, info: &DeviceTreeInfo) -> Result<(), RebootR
         error!("Failed to patch timer info to DT: {e}");
         RebootReason::InvalidFdt
     })?;
+
+    fdt.pack().map_err(|e| {
+        error!("Failed to pack DT after patching: {e}");
+        RebootReason::InvalidFdt
+    })?;
+
     Ok(())
 }
 
@@ -608,7 +655,7 @@ pub fn modify_for_next_stage(
 ) -> libfdt::Result<()> {
     fdt.unpack()?;
 
-    add_dice_node(fdt, bcc.as_ptr() as usize, bcc.len())?;
+    patch_dice_node(fdt, bcc.as_ptr() as usize, bcc.len())?;
 
     set_or_clear_chosen_flag(fdt, cstr!("avf,strict-boot"), strict_boot)?;
     set_or_clear_chosen_flag(fdt, cstr!("avf,new-instance"), new_instance)?;
@@ -618,24 +665,17 @@ pub fn modify_for_next_stage(
     Ok(())
 }
 
-/// Add a "google,open-dice"-compatible reserved-memory node to the tree.
-fn add_dice_node(fdt: &mut Fdt, addr: usize, size: usize) -> libfdt::Result<()> {
+/// Patch the "google,open-dice"-compatible reserved-memory node to point to the bcc range
+fn patch_dice_node(fdt: &mut Fdt, addr: usize, size: usize) -> libfdt::Result<()> {
     // We reject DTs with missing reserved-memory node as validation should have checked that the
     // "swiotlb" subnode (compatible = "restricted-dma-pool") was present.
-    let mut reserved_memory =
-        fdt.node_mut(cstr!("/reserved-memory"))?.ok_or(libfdt::FdtError::NotFound)?;
+    let node = fdt.node_mut(cstr!("/reserved-memory"))?.ok_or(libfdt::FdtError::NotFound)?;
 
-    let mut dice = reserved_memory.add_subnode(cstr!("dice"))?;
+    let mut node = node.next_compatible(cstr!("google,open-dice"))?.ok_or(FdtError::NotFound)?;
 
-    dice.appendprop(cstr!("compatible"), b"google,open-dice\0")?;
-
-    dice.appendprop(cstr!("no-map"), &[])?;
-
-    let addr = addr.try_into().unwrap();
-    let size = size.try_into().unwrap();
-    dice.appendprop_addrrange(cstr!("reg"), addr, size)?;
-
-    Ok(())
+    let addr: u64 = addr.try_into().unwrap();
+    let size: u64 = size.try_into().unwrap();
+    node.setprop_inplace(cstr!("reg"), flatten(&[addr.to_be_bytes(), size.to_be_bytes()]))
 }
 
 fn set_or_clear_chosen_flag(fdt: &mut Fdt, flag: &CStr, value: bool) -> libfdt::Result<()> {