pvmfw: Pass BCC to next stage through DT

Generate the next stage BCC in a heap-allocated page-aligned buffer that
our allocator leaks so that it outlives the execution of pvmfw and can
be accessed by the next stage. Flush the cache to ensure that it isn't
destroyed during invalidation (by the next stage) or missed if accessed
with the caches disabled.

Pass the size and location of the region through a pKVM-standard
device tree node.

Bug: 256827715
Test: atest MicrodroidHostTests
Change-Id: I5931054f74063eac3b3b21a6bcbe4881af2e1e8e
diff --git a/libs/libfdt/src/lib.rs b/libs/libfdt/src/lib.rs
index 64e6746..7c72fab 100644
--- a/libs/libfdt/src/lib.rs
+++ b/libs/libfdt/src/lib.rs
@@ -557,6 +557,11 @@
         Ok(self.path_offset(path)?.map(|offset| FdtNodeMut { fdt: self, offset }))
     }
 
+    /// Return the device tree as a slice (may be smaller than the containing buffer).
+    pub fn as_slice(&self) -> &[u8] {
+        &self.buffer[..self.totalsize()]
+    }
+
     fn path_offset(&self, path: &CStr) -> Result<Option<c_int>> {
         let len = path.to_bytes().len().try_into().map_err(|_| FdtError::BadPath)?;
         // SAFETY - Accesses are constrained to the DT totalsize (validated by ctor) and the
@@ -591,4 +596,13 @@
     fn capacity(&self) -> usize {
         self.buffer.len()
     }
+
+    fn header(&self) -> &libfdt_bindgen::fdt_header {
+        // SAFETY - A valid FDT (verified by constructor) must contain a valid fdt_header.
+        unsafe { &*(&self as *const _ as *const libfdt_bindgen::fdt_header) }
+    }
+
+    fn totalsize(&self) -> usize {
+        u32::from_be(self.header().totalsize) as usize
+    }
 }
diff --git a/pvmfw/src/entry.rs b/pvmfw/src/entry.rs
index 52e7a22..bfcb423 100644
--- a/pvmfw/src/entry.rs
+++ b/pvmfw/src/entry.rs
@@ -250,6 +250,7 @@
     crate::main(slices.fdt, slices.kernel, slices.ramdisk, &bcc, &mut memory)?;
 
     helpers::flushed_zeroize(bcc_slice);
+    helpers::flush(slices.fdt.as_slice());
 
     info!("Expecting a bug making MMIO_GUARD_UNMAP return NOT_SUPPORTED on success");
     memory.mmio_unmap_all().map_err(|e| {
diff --git a/pvmfw/src/fdt.rs b/pvmfw/src/fdt.rs
index dcd17b7..b735b9c 100644
--- a/pvmfw/src/fdt.rs
+++ b/pvmfw/src/fdt.rs
@@ -48,3 +48,29 @@
 
     Ok(None)
 }
+
+/// Add a "google,open-dice"-compatible reserved-memory node to the tree.
+pub fn add_dice_node(fdt: &mut libfdt::Fdt, addr: usize, size: usize) -> libfdt::Result<()> {
+    fdt.unpack()?;
+
+    let reserved_memory = CStr::from_bytes_with_nul(b"/reserved-memory\0").unwrap();
+    // 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(reserved_memory)?.ok_or(libfdt::FdtError::NotFound)?;
+
+    let dice = CStr::from_bytes_with_nul(b"dice\0").unwrap();
+    let mut dice = reserved_memory.add_subnode(dice)?;
+
+    let compatible = CStr::from_bytes_with_nul(b"compatible\0").unwrap();
+    dice.appendprop(compatible, b"google,open-dice\0")?;
+
+    let no_map = CStr::from_bytes_with_nul(b"no-map\0").unwrap();
+    dice.appendprop(no_map, &[])?;
+
+    let reg = CStr::from_bytes_with_nul(b"reg\0").unwrap();
+    dice.appendprop_addrrange(reg, addr as u64, size as u64)?;
+
+    fdt.pack()?;
+
+    Ok(())
+}
diff --git a/pvmfw/src/heap.rs b/pvmfw/src/heap.rs
index eab3bc4..e412f69 100644
--- a/pvmfw/src/heap.rs
+++ b/pvmfw/src/heap.rs
@@ -14,8 +14,11 @@
 
 //! Heap implementation.
 
+use alloc::alloc::alloc;
+use alloc::alloc::Layout;
+use alloc::boxed::Box;
+
 use core::alloc::GlobalAlloc as _;
-use core::alloc::Layout;
 use core::ffi::c_void;
 use core::mem;
 use core::num::NonZeroUsize;
@@ -33,6 +36,19 @@
     HEAP_ALLOCATOR.lock().init(HEAP.as_mut_ptr() as usize, HEAP.len());
 }
 
+/// Allocate an aligned but uninitialized slice of heap.
+pub fn aligned_boxed_slice(size: usize, align: usize) -> Option<Box<[u8]>> {
+    let size = NonZeroUsize::new(size)?.get();
+    let layout = Layout::from_size_align(size, align).ok()?;
+    // SAFETY - We verify that `size` and the returned `ptr` are non-null.
+    let ptr = unsafe { alloc(layout) };
+    let ptr = NonNull::new(ptr)?.as_ptr();
+    let slice_ptr = ptr::slice_from_raw_parts_mut(ptr, size);
+
+    // SAFETY - The memory was allocated using the proper layout by our global_allocator.
+    Some(unsafe { Box::from_raw(slice_ptr) })
+}
+
 #[no_mangle]
 unsafe extern "C" fn malloc(size: usize) -> *mut c_void {
     malloc_(size).map_or(ptr::null_mut(), |p| p.cast::<c_void>().as_ptr())
diff --git a/pvmfw/src/helpers.rs b/pvmfw/src/helpers.rs
index ad980f5..40266f7 100644
--- a/pvmfw/src/helpers.rs
+++ b/pvmfw/src/helpers.rs
@@ -91,8 +91,14 @@
 }
 
 #[inline]
+/// Flushes the slice to the point of unification.
+pub fn flush(reg: &[u8]) {
+    flush_region(reg.as_ptr() as usize, reg.len())
+}
+
+#[inline]
 /// Overwrites the slice with zeroes, to the point of unification.
 pub fn flushed_zeroize(reg: &mut [u8]) {
     reg.zeroize();
-    flush_region(reg.as_ptr() as usize, reg.len())
+    flush(reg)
 }
diff --git a/pvmfw/src/main.rs b/pvmfw/src/main.rs
index efd3e7c..9c62e03 100644
--- a/pvmfw/src/main.rs
+++ b/pvmfw/src/main.rs
@@ -19,6 +19,8 @@
 #![feature(default_alloc_error_handler)]
 #![feature(ptr_const_cast)] // Stabilized in 1.65.0
 
+extern crate alloc;
+
 mod avb;
 mod config;
 mod dice;
@@ -34,10 +36,14 @@
 mod pci;
 mod smccc;
 
+use alloc::boxed::Box;
+
 use crate::{
     avb::PUBLIC_KEY,
     dice::derive_next_bcc,
     entry::RebootReason,
+    fdt::add_dice_node,
+    helpers::flush,
     helpers::GUEST_PAGE_SIZE,
     memory::MemoryTracker,
     pci::{find_virtio_devices, map_mmio},
@@ -51,7 +57,7 @@
 const NEXT_BCC_SIZE: usize = GUEST_PAGE_SIZE;
 
 fn main(
-    fdt: &Fdt,
+    fdt: &mut Fdt,
     signed_kernel: &[u8],
     ramdisk: Option<&[u8]>,
     bcc: &bcc::Handover,
@@ -82,8 +88,6 @@
         RebootReason::PayloadVerificationError
     })?;
 
-    let mut scratch_bcc = [0; NEXT_BCC_SIZE];
-    let next_bcc = &mut scratch_bcc; // TODO(b/256827715): Pass result BCC to next stage.
     let debug_mode = false; // TODO(b/256148034): Derive the DICE mode from the received initrd.
     const HASH_SIZE: usize = 64;
     let mut hashes = [0; HASH_SIZE * 2]; // TODO(b/256148034): Extract AvbHashDescriptor digests.
@@ -101,6 +105,12 @@
     } else {
         &hashes[..HASH_SIZE]
     };
+    let next_bcc = heap::aligned_boxed_slice(NEXT_BCC_SIZE, GUEST_PAGE_SIZE).ok_or_else(|| {
+        error!("Failed to allocate the next-stage BCC");
+        RebootReason::InternalError
+    })?;
+    // By leaking the slice, its content will be left behind for the next stage.
+    let next_bcc = Box::leak(next_bcc);
     let next_bcc_size =
         derive_next_bcc(bcc, next_bcc, code_hash, debug_mode, PUBLIC_KEY).map_err(|e| {
             error!("Failed to derive next-stage DICE secrets: {e:?}");
@@ -108,6 +118,13 @@
         })?;
     trace!("Next BCC: {:x?}", bcc::Handover::new(&next_bcc[..next_bcc_size]));
 
+    flush(next_bcc);
+
+    add_dice_node(fdt, next_bcc.as_ptr() as usize, NEXT_BCC_SIZE).map_err(|e| {
+        error!("Failed to add DICE node to device tree: {e}");
+        RebootReason::InternalError
+    })?;
+
     info!("Starting payload...");
     Ok(())
 }