Merge "Add VirtualizationModuleFrameworkInitializer"
diff --git a/compos/src/compsvc_main.rs b/compos/src/compsvc_main.rs
index 206dd4b..77e2daa 100644
--- a/compos/src/compsvc_main.rs
+++ b/compos/src/compsvc_main.rs
@@ -23,11 +23,13 @@
 mod fsverity;
 
 use anyhow::Result;
+use binder::unstable_api::AsNative;
 use compos_common::COMPOS_VSOCK_PORT;
 use log::{debug, error};
-use rpcbinder::RpcServer;
+use std::os::raw::c_void;
 use std::panic;
-use vm_payload_bindgen::AVmPayload_notifyPayloadReady;
+use std::ptr;
+use vm_payload_bindgen::{AIBinder, AVmPayload_notifyPayloadReady, AVmPayload_runVsockRpcServer};
 
 fn main() {
     if let Err(e) = try_main() {
@@ -46,10 +48,20 @@
     }));
 
     debug!("compsvc is starting as a rpc service.");
-    let service = compsvc::new_binder()?.as_binder();
-    let server = RpcServer::new_vsock(service, COMPOS_VSOCK_PORT)?;
-    // SAFETY: Invokes a method from the bindgen library `vm_payload_bindgen`.
-    unsafe { AVmPayload_notifyPayloadReady() };
-    server.join();
+    let param = ptr::null_mut();
+    let mut service = compsvc::new_binder()?.as_binder();
+    unsafe {
+        // SAFETY: We hold a strong pointer, so the raw pointer remains valid. The bindgen AIBinder
+        // is the same type as sys::AIBinder.
+        let service = service.as_native_mut() as *mut AIBinder;
+        // SAFETY: It is safe for on_ready to be invoked at any time, with any parameter.
+        AVmPayload_runVsockRpcServer(service, COMPOS_VSOCK_PORT, Some(on_ready), param);
+    }
     Ok(())
 }
+
+extern "C" fn on_ready(_param: *mut c_void) {
+    // SAFETY: Invokes a method from the bindgen library `vm_payload_bindgen` which is safe to
+    // call at any time.
+    unsafe { AVmPayload_notifyPayloadReady() };
+}
diff --git a/libs/libfdt/Android.bp b/libs/libfdt/Android.bp
new file mode 100644
index 0000000..72399b0
--- /dev/null
+++ b/libs/libfdt/Android.bp
@@ -0,0 +1,45 @@
+package {
+    default_applicable_licenses: ["Android-Apache-2.0"],
+}
+
+rust_bindgen {
+    name: "liblibfdt_bindgen",
+    crate_name: "libfdt_bindgen",
+    wrapper_src: "bindgen/fdt.h",
+    source_stem: "bindings",
+    bindgen_flags: [
+        "--size_t-is-usize",
+        "--allowlist-type=fdt_.*",
+        "--allowlist-function=fdt_.*",
+        "--allowlist-var=FDT_.*",
+        "--use-core",
+        "--raw-line=#![no_std]",
+        "--ctypes-prefix=core::ffi",
+    ],
+    static_libs: [
+        "libfdt",
+    ],
+    apex_available: ["com.android.virt"],
+}
+
+rust_library_rlib {
+    name: "liblibfdt",
+    crate_name: "libfdt",
+    srcs: [
+        "src/lib.rs",
+        ":liblibfdt_bindgen",
+    ],
+    edition: "2021",
+    no_stdlibs: true,
+    prefer_rlib: true,
+    stdlibs: [
+        "libcore.rust_sysroot",
+    ],
+    rustlibs: [
+        "liblibfdt_bindgen",
+    ],
+    whole_static_libs: [
+        "libfdt",
+    ],
+    apex_available: ["com.android.virt"],
+}
diff --git a/libs/libfdt/bindgen/fdt.h b/libs/libfdt/bindgen/fdt.h
new file mode 100644
index 0000000..16f2784
--- /dev/null
+++ b/libs/libfdt/bindgen/fdt.h
@@ -0,0 +1,17 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <libfdt.h>
diff --git a/libs/libfdt/src/lib.rs b/libs/libfdt/src/lib.rs
new file mode 100644
index 0000000..01f7b36
--- /dev/null
+++ b/libs/libfdt/src/lib.rs
@@ -0,0 +1,502 @@
+// Copyright 2022, The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Wrapper around libfdt library. Provides parsing/generating functionality
+//! to a bare-metal environment.
+
+#![no_std]
+
+use core::ffi::{c_int, c_void, CStr};
+use core::fmt;
+use core::mem;
+use core::ops::Range;
+use core::result;
+use core::slice;
+
+/// Error type corresponding to libfdt error codes.
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub enum FdtError {
+    /// FDT_ERR_NOTFOUND
+    NotFound,
+    /// FDT_ERR_EXISTS
+    Exists,
+    /// FDT_ERR_NOSPACE
+    NoSpace,
+    /// FDT_ERR_BADOFFSET
+    BadOffset,
+    /// FDT_ERR_BADPATH
+    BadPath,
+    /// FDT_ERR_BADPHANDLE
+    BadPhandle,
+    /// FDT_ERR_BADSTATE
+    BadState,
+    /// FDT_ERR_TRUNCATED
+    Truncated,
+    /// FDT_ERR_BADMAGIC
+    BadMagic,
+    /// FDT_ERR_BADVERSION
+    BadVersion,
+    /// FDT_ERR_BADSTRUCTURE
+    BadStructure,
+    /// FDT_ERR_BADLAYOUT
+    BadLayout,
+    /// FDT_ERR_INTERNAL
+    Internal,
+    /// FDT_ERR_BADNCELLS
+    BadNCells,
+    /// FDT_ERR_BADVALUE
+    BadValue,
+    /// FDT_ERR_BADOVERLAY
+    BadOverlay,
+    /// FDT_ERR_NOPHANDLES
+    NoPhandles,
+    /// FDT_ERR_BADFLAGS
+    BadFlags,
+    /// FDT_ERR_ALIGNMENT
+    Alignment,
+    /// Unexpected error code
+    Unknown(i32),
+}
+
+impl fmt::Display for FdtError {
+    /// Prints error messages from libfdt.h documentation.
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            Self::NotFound => write!(f, "The requested node or property does not exist"),
+            Self::Exists => write!(f, "Attempted to create an existing node or property"),
+            Self::NoSpace => write!(f, "Insufficient buffer space to contain the expanded tree"),
+            Self::BadOffset => write!(f, "Structure block offset is out-of-bounds or invalid"),
+            Self::BadPath => write!(f, "Badly formatted path"),
+            Self::BadPhandle => write!(f, "Invalid phandle length or value"),
+            Self::BadState => write!(f, "Received incomplete device tree"),
+            Self::Truncated => write!(f, "Device tree or sub-block is improperly terminated"),
+            Self::BadMagic => write!(f, "Device tree header missing its magic number"),
+            Self::BadVersion => write!(f, "Device tree has a version which can't be handled"),
+            Self::BadStructure => write!(f, "Device tree has a corrupt structure block"),
+            Self::BadLayout => write!(f, "Device tree sub-blocks in unsupported order"),
+            Self::Internal => write!(f, "libfdt has failed an internal assertion"),
+            Self::BadNCells => write!(f, "Bad format or value of #address-cells or #size-cells"),
+            Self::BadValue => write!(f, "Unexpected property value"),
+            Self::BadOverlay => write!(f, "Overlay cannot be applied"),
+            Self::NoPhandles => write!(f, "Device tree doesn't have any phandle available anymore"),
+            Self::BadFlags => write!(f, "Invalid flag or invalid combination of flags"),
+            Self::Alignment => write!(f, "Device tree base address is not 8-byte aligned"),
+            Self::Unknown(e) => write!(f, "Unknown libfdt error '{e}'"),
+        }
+    }
+}
+
+/// Result type with FdtError enum.
+pub type Result<T> = result::Result<T, FdtError>;
+
+fn fdt_err(val: c_int) -> Result<c_int> {
+    if val >= 0 {
+        Ok(val)
+    } else {
+        Err(match -val as _ {
+            libfdt_bindgen::FDT_ERR_NOTFOUND => FdtError::NotFound,
+            libfdt_bindgen::FDT_ERR_EXISTS => FdtError::Exists,
+            libfdt_bindgen::FDT_ERR_NOSPACE => FdtError::NoSpace,
+            libfdt_bindgen::FDT_ERR_BADOFFSET => FdtError::BadOffset,
+            libfdt_bindgen::FDT_ERR_BADPATH => FdtError::BadPath,
+            libfdt_bindgen::FDT_ERR_BADPHANDLE => FdtError::BadPhandle,
+            libfdt_bindgen::FDT_ERR_BADSTATE => FdtError::BadState,
+            libfdt_bindgen::FDT_ERR_TRUNCATED => FdtError::Truncated,
+            libfdt_bindgen::FDT_ERR_BADMAGIC => FdtError::BadMagic,
+            libfdt_bindgen::FDT_ERR_BADVERSION => FdtError::BadVersion,
+            libfdt_bindgen::FDT_ERR_BADSTRUCTURE => FdtError::BadStructure,
+            libfdt_bindgen::FDT_ERR_BADLAYOUT => FdtError::BadLayout,
+            libfdt_bindgen::FDT_ERR_INTERNAL => FdtError::Internal,
+            libfdt_bindgen::FDT_ERR_BADNCELLS => FdtError::BadNCells,
+            libfdt_bindgen::FDT_ERR_BADVALUE => FdtError::BadValue,
+            libfdt_bindgen::FDT_ERR_BADOVERLAY => FdtError::BadOverlay,
+            libfdt_bindgen::FDT_ERR_NOPHANDLES => FdtError::NoPhandles,
+            libfdt_bindgen::FDT_ERR_BADFLAGS => FdtError::BadFlags,
+            libfdt_bindgen::FDT_ERR_ALIGNMENT => FdtError::Alignment,
+            _ => FdtError::Unknown(val),
+        })
+    }
+}
+
+fn fdt_err_expect_zero(val: c_int) -> Result<()> {
+    match fdt_err(val)? {
+        0 => Ok(()),
+        _ => Err(FdtError::Unknown(val)),
+    }
+}
+
+/// Value of a #address-cells property.
+#[derive(Copy, Clone, Debug)]
+enum AddrCells {
+    Single = 1,
+    Double = 2,
+}
+
+impl TryFrom<c_int> for AddrCells {
+    type Error = FdtError;
+
+    fn try_from(res: c_int) -> Result<Self> {
+        match fdt_err(res)? {
+            x if x == Self::Single as c_int => Ok(Self::Single),
+            x if x == Self::Double as c_int => Ok(Self::Double),
+            _ => Err(FdtError::BadNCells),
+        }
+    }
+}
+
+/// Value of a #size-cells property.
+#[derive(Copy, Clone, Debug)]
+enum SizeCells {
+    None = 0,
+    Single = 1,
+    Double = 2,
+}
+
+impl TryFrom<c_int> for SizeCells {
+    type Error = FdtError;
+
+    fn try_from(res: c_int) -> Result<Self> {
+        match fdt_err(res)? {
+            x if x == Self::None as c_int => Ok(Self::None),
+            x if x == Self::Single as c_int => Ok(Self::Single),
+            x if x == Self::Double as c_int => Ok(Self::Double),
+            _ => Err(FdtError::BadNCells),
+        }
+    }
+}
+
+/// Iterator over cells of a DT property.
+#[derive(Debug)]
+pub struct CellIterator<'a> {
+    chunks: slice::ChunksExact<'a, u8>,
+}
+
+impl<'a> CellIterator<'a> {
+    fn new(bytes: &'a [u8]) -> Self {
+        const CHUNK_SIZE: usize = mem::size_of::<<CellIterator as Iterator>::Item>();
+
+        Self { chunks: bytes.chunks_exact(CHUNK_SIZE) }
+    }
+}
+
+impl<'a> Iterator for CellIterator<'a> {
+    type Item = u32;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        Some(Self::Item::from_be_bytes(self.chunks.next()?.try_into().ok()?))
+    }
+}
+
+/// Iterator over a 'reg' property of a DT node.
+#[derive(Debug)]
+pub struct RegIterator<'a> {
+    cells: CellIterator<'a>,
+    addr_cells: AddrCells,
+    size_cells: SizeCells,
+}
+
+/// Represents a contiguous region within the address space defined by the parent bus.
+/// Commonly means the offsets and lengths of MMIO blocks, but may have a different meaning on some
+/// bus types. Addresses in the address space defined by the root node are CPU real addresses.
+#[derive(Copy, Clone, Debug)]
+pub struct Reg<T> {
+    /// Base address of the region.
+    pub addr: T,
+    /// Size of the region (optional).
+    pub size: Option<T>,
+}
+
+impl<'a> RegIterator<'a> {
+    fn new(cells: CellIterator<'a>, addr_cells: AddrCells, size_cells: SizeCells) -> Self {
+        Self { cells, addr_cells, size_cells }
+    }
+}
+
+impl<'a> Iterator for RegIterator<'a> {
+    type Item = Reg<u64>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let make_double = |a, b| (u64::from(a) << 32) | u64::from(b);
+
+        let addr = match self.addr_cells {
+            AddrCells::Single => self.cells.next()?.into(),
+            AddrCells::Double => make_double(self.cells.next()?, self.cells.next()?),
+        };
+        // If the parent node specifies a value of 0 for #size-cells, 'size' shall be omitted.
+        let size = match self.size_cells {
+            SizeCells::None => None,
+            SizeCells::Single => Some(self.cells.next()?.into()),
+            SizeCells::Double => Some(make_double(self.cells.next()?, self.cells.next()?)),
+        };
+
+        Some(Self::Item { addr, size })
+    }
+}
+
+/// Iterator over the address ranges defined by the /memory/ node.
+#[derive(Debug)]
+pub struct MemRegIterator<'a> {
+    reg: RegIterator<'a>,
+}
+
+impl<'a> MemRegIterator<'a> {
+    fn new(reg: RegIterator<'a>) -> Result<Self> {
+        Ok(Self { reg })
+    }
+}
+
+impl<'a> Iterator for MemRegIterator<'a> {
+    type Item = Range<usize>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let next = self.reg.next()?;
+        let addr = usize::try_from(next.addr).ok()?;
+        let size = usize::try_from(next.size?).ok()?;
+
+        Some(addr..addr.checked_add(size)?)
+    }
+}
+
+/// DT node.
+#[derive(Clone, Copy)]
+pub struct FdtNode<'a> {
+    fdt: &'a Fdt,
+    offset: c_int,
+}
+
+impl<'a> FdtNode<'a> {
+    /// Find parent node.
+    pub fn parent(&self) -> Result<Self> {
+        // SAFETY - Accesses (read-only) are constrained to the DT totalsize.
+        let ret = unsafe { libfdt_bindgen::fdt_parent_offset(self.fdt.as_ptr(), self.offset) };
+
+        Ok(Self { fdt: self.fdt, offset: fdt_err(ret)? })
+    }
+
+    /// Retrieve the standard (deprecated) device_type <string> property.
+    pub fn device_type(&self) -> Result<&CStr> {
+        self.getprop_str(CStr::from_bytes_with_nul(b"device_type\0").unwrap())
+    }
+
+    /// Retrieve the standard reg <prop-encoded-array> property.
+    pub fn reg(&self) -> Result<RegIterator<'a>> {
+        let parent = self.parent()?;
+
+        let addr_cells = parent.address_cells()?;
+        let size_cells = parent.size_cells()?;
+        let cells = self.getprop_cells(CStr::from_bytes_with_nul(b"reg\0").unwrap())?;
+
+        Ok(RegIterator::new(cells, addr_cells, size_cells))
+    }
+
+    /// Retrieve the value of a given <string> property.
+    pub fn getprop_str(&self, name: &CStr) -> Result<&CStr> {
+        CStr::from_bytes_with_nul(self.getprop(name)?).map_err(|_| FdtError::BadValue)
+    }
+
+    /// Retrieve the value of a given property as an array of cells.
+    pub fn getprop_cells(&self, name: &CStr) -> Result<CellIterator<'a>> {
+        Ok(CellIterator::new(self.getprop(name)?))
+    }
+
+    /// Retrieve the value of a given <u32> property.
+    pub fn getprop_u32(&self, name: &CStr) -> Result<u32> {
+        let prop = self.getprop(name)?.try_into().map_err(|_| FdtError::BadValue)?;
+        Ok(u32::from_be_bytes(prop))
+    }
+
+    /// Retrieve the value of a given <u64> property.
+    pub fn getprop_u64(&self, name: &CStr) -> Result<u64> {
+        let prop = self.getprop(name)?.try_into().map_err(|_| FdtError::BadValue)?;
+        Ok(u64::from_be_bytes(prop))
+    }
+
+    /// Retrieve the value of a given property.
+    pub fn getprop(&self, name: &CStr) -> Result<&'a [u8]> {
+        let mut len: i32 = 0;
+        // SAFETY - Accesses are constrained to the DT totalsize (validated by ctor) and the
+        // function respects the passed number of characters.
+        let prop = unsafe {
+            libfdt_bindgen::fdt_getprop_namelen(
+                self.fdt.as_ptr(),
+                self.offset,
+                name.as_ptr(),
+                // *_namelen functions don't include the trailing nul terminator in 'len'.
+                name.to_bytes().len().try_into().map_err(|_| FdtError::BadPath)?,
+                &mut len as *mut i32,
+            )
+        } as *const u8;
+        if prop.is_null() {
+            return fdt_err(len).and(Err(FdtError::Internal));
+        }
+        let len = usize::try_from(fdt_err(len)?).map_err(|_| FdtError::Internal)?;
+        let base =
+            (prop as usize).checked_sub(self.fdt.as_ptr() as usize).ok_or(FdtError::Internal)?;
+
+        self.fdt.bytes.get(base..(base + len)).ok_or(FdtError::Internal)
+    }
+
+    /// Get reference to the containing device tree.
+    pub fn fdt(&self) -> &Fdt {
+        self.fdt
+    }
+
+    fn next_compatible(self, compatible: &CStr) -> Result<Option<Self>> {
+        // SAFETY - Accesses (read-only) are constrained to the DT totalsize.
+        let ret = unsafe {
+            libfdt_bindgen::fdt_node_offset_by_compatible(
+                self.fdt.as_ptr(),
+                self.offset,
+                compatible.as_ptr(),
+            )
+        };
+
+        match fdt_err(ret) {
+            Ok(offset) => Ok(Some(Self { fdt: self.fdt, offset })),
+            Err(FdtError::NotFound) => Ok(None),
+            Err(e) => Err(e),
+        }
+    }
+
+    fn address_cells(&self) -> Result<AddrCells> {
+        // SAFETY - Accesses are constrained to the DT totalsize (validated by ctor).
+        unsafe { libfdt_bindgen::fdt_address_cells(self.fdt.as_ptr(), self.offset) }
+            .try_into()
+            .map_err(|_| FdtError::Internal)
+    }
+
+    fn size_cells(&self) -> Result<SizeCells> {
+        // SAFETY - Accesses are constrained to the DT totalsize (validated by ctor).
+        unsafe { libfdt_bindgen::fdt_size_cells(self.fdt.as_ptr(), self.offset) }
+            .try_into()
+            .map_err(|_| FdtError::Internal)
+    }
+}
+
+/// Iterator over nodes sharing a same compatible string.
+pub struct CompatibleIterator<'a> {
+    node: FdtNode<'a>,
+    compatible: &'a CStr,
+}
+
+impl<'a> CompatibleIterator<'a> {
+    fn new(fdt: &'a Fdt, compatible: &'a CStr) -> Result<Self> {
+        let node = fdt.root()?;
+        Ok(Self { node, compatible })
+    }
+}
+
+impl<'a> Iterator for CompatibleIterator<'a> {
+    type Item = FdtNode<'a>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let next = self.node.next_compatible(self.compatible).ok()?;
+
+        if let Some(node) = next {
+            self.node = node;
+        }
+
+        next
+    }
+}
+
+/// Wrapper around low-level read-only libfdt functions.
+#[repr(transparent)]
+pub struct Fdt {
+    bytes: [u8],
+}
+
+impl Fdt {
+    /// Wraps a slice containing a Flattened Device Tree.
+    ///
+    /// Fails if the FDT does not pass validation.
+    pub fn from_slice(fdt: &[u8]) -> Result<&Self> {
+        // SAFETY - The FDT will be validated before it is returned.
+        let fdt = unsafe { Self::unchecked_from_slice(fdt) };
+        fdt.check_full()?;
+        Ok(fdt)
+    }
+
+    /// Wraps a slice containing a Flattened Device Tree.
+    ///
+    /// # Safety
+    ///
+    /// The returned FDT might be invalid, only use on slices containing a valid DT.
+    pub unsafe fn unchecked_from_slice(fdt: &[u8]) -> &Self {
+        mem::transmute::<&[u8], &Self>(fdt)
+    }
+
+    /// Return an iterator of memory banks specified the "/memory" node.
+    ///
+    /// NOTE: This does not support individual "/memory@XXXX" banks.
+    pub fn memory(&self) -> Result<MemRegIterator> {
+        let memory = CStr::from_bytes_with_nul(b"/memory\0").unwrap();
+        let device_type = CStr::from_bytes_with_nul(b"memory\0").unwrap();
+
+        let node = self.node(memory)?;
+        if node.device_type()? != device_type {
+            return Err(FdtError::BadValue);
+        }
+
+        MemRegIterator::new(node.reg()?)
+    }
+
+    /// Retrieve the standard /chosen node.
+    pub fn chosen(&self) -> Result<FdtNode> {
+        self.node(CStr::from_bytes_with_nul(b"/chosen\0").unwrap())
+    }
+
+    /// Get the root node of the tree.
+    pub fn root(&self) -> Result<FdtNode> {
+        self.node(CStr::from_bytes_with_nul(b"/\0").unwrap())
+    }
+
+    /// Find a tree node by its full path.
+    pub fn node(&self, path: &CStr) -> Result<FdtNode> {
+        let offset = self.path_offset(path)?;
+        Ok(FdtNode { fdt: self, offset })
+    }
+
+    /// Iterate over nodes with a given compatible string.
+    pub fn compatible_nodes<'a>(&'a self, compatible: &'a CStr) -> Result<CompatibleIterator<'a>> {
+        CompatibleIterator::new(self, compatible)
+    }
+
+    fn path_offset(&self, path: &CStr) -> Result<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
+        // function respects the passed number of characters.
+        let ret = unsafe {
+            // *_namelen functions don't include the trailing nul terminator in 'len'.
+            libfdt_bindgen::fdt_path_offset_namelen(self.as_ptr(), path.as_ptr(), len)
+        };
+
+        fdt_err(ret)
+    }
+
+    fn check_full(&self) -> Result<()> {
+        let len = self.bytes.len();
+        // SAFETY - Only performs read accesses within the limits of the slice. If successful, this
+        // call guarantees to other unsafe calls that the header contains a valid totalsize (w.r.t.
+        // 'len' i.e. the self.fdt slice) that those C functions can use to perform bounds
+        // checking. The library doesn't maintain an internal state (such as pointers) between
+        // calls as it expects the client code to keep track of the objects (DT, nodes, ...).
+        let ret = unsafe { libfdt_bindgen::fdt_check_full(self.as_ptr(), len) };
+        fdt_err_expect_zero(ret)
+    }
+
+    fn as_ptr(&self) -> *const c_void {
+        self as *const _ as *const c_void
+    }
+}
diff --git a/tests/benchmark/src/native/benchmarkbinary.cpp b/tests/benchmark/src/native/benchmarkbinary.cpp
index 66b41a1..70ec7db 100644
--- a/tests/benchmark/src/native/benchmarkbinary.cpp
+++ b/tests/benchmark/src/native/benchmarkbinary.cpp
@@ -161,10 +161,8 @@
 Result<void> run_io_benchmark_tests() {
     auto test_service = ndk::SharedRefBase::make<IOBenchmarkService>();
     auto callback = []([[maybe_unused]] void* param) { AVmPayload_notifyPayloadReady(); };
-    if (!AVmPayload_runVsockRpcServer(test_service->asBinder().get(), test_service->SERVICE_PORT,
-                                      callback, nullptr)) {
-        return Error() << "RPC Server failed to run";
-    }
+    AVmPayload_runVsockRpcServer(test_service->asBinder().get(), test_service->SERVICE_PORT,
+                                 callback, nullptr);
     return {};
 }
 } // Anonymous namespace
diff --git a/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java b/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
index f0c89c9..11b3e84 100644
--- a/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
+++ b/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
@@ -66,6 +66,7 @@
 import java.io.ByteArrayInputStream;
 import java.io.File;
 import java.io.IOException;
+import java.io.InputStream;
 import java.io.InputStreamReader;
 import java.io.PipedInputStream;
 import java.io.PipedOutputStream;
@@ -94,6 +95,16 @@
 
     private static final int BOOT_COMPLETE_TIMEOUT = 30000; // 30 seconds
 
+    private static class VmInfo {
+        final Process mProcess;
+        final String mCid;
+
+        VmInfo(Process process, String cid) {
+            mProcess = process;
+            mCid = cid;
+        }
+    }
+
     @Rule public TestLogData mTestLogs = new TestLogData();
     @Rule public TestName mTestName = new TestName();
     @Rule public TestMetrics mMetrics = new TestMetrics();
@@ -135,23 +146,16 @@
                                 .collect(toList())));
         FileUtil.writeToFile(config.toString(), configFile);
 
-        File mkPayload = findTestFile("mk_payload");
         RunUtil runUtil = new RunUtil();
-        // Set the parent dir on the PATH (e.g. <workdir>/bin)
-        String separator = System.getProperty("path.separator");
-        String path = mkPayload.getParentFile().getPath() + separator + System.getenv("PATH");
-        runUtil.setEnvVariable("PATH", path);
-
-        List<String> command = new ArrayList<>();
-        command.add("mk_payload");
-        command.add("--metadata-only");
-        command.add(configFile.toString());
-        command.add(payloadMetadata.toString());
-
-        CommandResult result =
-                runUtil.runTimedCmd(
-                        // mk_payload should run fast enough
-                        5 * 1000, "/bin/bash", "-c", String.join(" ", command));
+        String command =
+                String.join(
+                        " ",
+                        findTestFile("mk_payload").getAbsolutePath(),
+                        "--metadata-only",
+                        configFile.getAbsolutePath(),
+                        payloadMetadata.getAbsolutePath());
+        // mk_payload should run fast enough
+        CommandResult result = runUtil.runTimedCmd(5000, "/bin/bash", "-c", command);
         String out = result.getStdout();
         String err = result.getStderr();
         assertWithMessage(
@@ -175,12 +179,10 @@
         runUtil.setEnvVariable("PATH", path);
 
         List<String> command = new ArrayList<>();
-        command.add("sign_virt_apex");
-        for (Map.Entry<String, File> entry : keyOverrides.entrySet()) {
-            String filename = entry.getKey();
-            File overridingKey = entry.getValue();
-            command.add("--key_override " + filename + "=" + overridingKey.getPath());
-        }
+        command.add(signVirtApex.getAbsolutePath());
+        keyOverrides.forEach(
+                (filename, keyFile) ->
+                        command.add("--key_override " + filename + "=" + keyFile.getPath()));
         command.add(signingKey.getPath());
         command.add(virtApexDir.getPath());
 
@@ -202,18 +204,11 @@
             long timeoutMillis, Callable<T> callable, org.hamcrest.Matcher<T> matcher)
             throws Exception {
         long start = System.currentTimeMillis();
-        while (true) {
-            try {
-                assertThat(callable.call(), matcher);
-                return;
-            } catch (Throwable e) {
-                if (System.currentTimeMillis() - start < timeoutMillis) {
-                    Thread.sleep(500);
-                } else {
-                    throw e;
-                }
-            }
+        while ((System.currentTimeMillis() - start < timeoutMillis)
+                && !matcher.matches(callable.call())) {
+            Thread.sleep(500);
         }
+        assertThat(callable.call(), matcher);
     }
 
     static class ActiveApexInfo {
@@ -236,12 +231,10 @@
         }
 
         ActiveApexInfo get(String apexName) {
-            for (ActiveApexInfo info : mList) {
-                if (info.name.equals(apexName)) {
-                    return info;
-                }
-            }
-            return null;
+            return mList.stream()
+                    .filter(info -> apexName.equals(info.name))
+                    .findFirst()
+                    .orElse(null);
         }
 
         List<ActiveApexInfo> getSharedLibApexes() {
@@ -273,13 +266,8 @@
         return new ActiveApexInfoList(list);
     }
 
-    private Process runMicrodroidWithResignedImages(
-            File key,
-            Map<String, File> keyOverrides,
-            boolean isProtected,
-            boolean waitForOnline,
-            String consolePath)
-            throws Exception {
+    private VmInfo runMicrodroidWithResignedImages(
+            File key, Map<String, File> keyOverrides, boolean isProtected) throws Exception {
         CommandRunner android = new CommandRunner(getDevice());
 
         File virtApexDir = FileUtil.createTempDir("virt_apex");
@@ -389,20 +377,21 @@
                         getDevice().getSerialNumber(),
                         "shell",
                         VIRT_APEX + "bin/vm run",
-                        (consolePath != null) ? "--console " + consolePath : "",
+                        "--console " + CONSOLE_PATH,
                         "--log " + LOG_PATH,
                         configPath);
 
         PipedInputStream pis = new PipedInputStream();
         Process process = RunUtil.getDefault().runCmdInBackground(args, new PipedOutputStream(pis));
-        BufferedReader stdout = new BufferedReader(new InputStreamReader(pis));
+        return new VmInfo(process, extractCidFrom(pis));
+    }
 
-        // Retrieve the CID from the vm tool output
+    private static String extractCidFrom(InputStream input) throws IOException {
         String cid = null;
         Pattern pattern = Pattern.compile("with CID (\\d+)");
-        try {
-            String line;
-            while ((line = stdout.readLine()) != null) {
+        String line;
+        try (BufferedReader out = new BufferedReader(new InputStreamReader(input))) {
+            while ((line = out.readLine()) != null) {
                 CLog.i("VM output: " + line);
                 Matcher matcher = pattern.matcher(line);
                 if (matcher.find()) {
@@ -410,21 +399,11 @@
                     break;
                 }
             }
-        } catch (IOException ex) {
-            throw new IllegalStateException(
-                    "Could not find the CID of the VM. The process probably died.", ex);
         }
-        if (cid == null) {
-            throw new IllegalStateException(
-                    "Could not find the CID of the VM. Output does not contain the expected"
-                            + " pattern.");
-        }
-
-        if (waitForOnline) {
-            adbConnectToMicrodroid(getDevice(), cid);
-        }
-
-        return process;
+        assertWithMessage("The output does not contain the expected pattern for CID.")
+                .that(cid)
+                .isNotNull();
+        return cid;
     }
 
     @Test
@@ -432,20 +411,16 @@
     @CddTest(requirements = {"9.17/C-2-1", "9.17/C-2-2", "9.17/C-2-6"})
     public void testBootFailsWhenProtectedVmStartsWithImagesSignedWithDifferentKey()
             throws Exception {
+        boolean protectedVm = true;
         assumeTrue(
-                "Protected VMs are not supported",
-                getAndroidDevice().supportsMicrodroid(/*protectedVm=*/ true));
+                "Skip if protected VMs are not supported",
+                getAndroidDevice().supportsMicrodroid(protectedVm));
 
         File key = findTestFile("test.com.android.virt.pem");
         Map<String, File> keyOverrides = Map.of();
-        boolean isProtected = true;
-        boolean waitForOnline = false; // VM should shut down due to boot failure.
-        String consolePath = TEST_ROOT + "console";
-        Process process =
-                runMicrodroidWithResignedImages(
-                        key, keyOverrides, isProtected, waitForOnline, consolePath);
-        process.waitFor(5L, TimeUnit.SECONDS);
-        assertThat(getDevice().pullFileContents(consolePath), containsString("pvmfw boot failed"));
+        VmInfo vmInfo = runMicrodroidWithResignedImages(key, keyOverrides, protectedVm);
+        vmInfo.mProcess.waitFor(5L, TimeUnit.SECONDS);
+        assertThat(getDevice().pullFileContents(CONSOLE_PATH), containsString("pvmfw boot failed"));
     }
 
     // TODO(b/245277660): Resigning the system/vendor image changes the vbmeta hash.
@@ -457,13 +432,10 @@
             throws Exception {
         File key = findTestFile("test.com.android.virt.pem");
         Map<String, File> keyOverrides = Map.of();
-        boolean isProtected = false;
-        boolean waitForOnline = true; // Device online means that boot must have succeeded.
-        String consolePath = TEST_ROOT + "console";
-        Process process =
-                runMicrodroidWithResignedImages(
-                        key, keyOverrides, isProtected, waitForOnline, consolePath);
-        process.destroy();
+        VmInfo vmInfo = runMicrodroidWithResignedImages(key, keyOverrides, /*isProtected=*/ false);
+        // Device online means that boot must have succeeded.
+        adbConnectToMicrodroid(getDevice(), vmInfo.mCid);
+        vmInfo.mProcess.destroy();
     }
 
     @Test
@@ -473,19 +445,14 @@
         File key = findTestFile("test.com.android.virt.pem");
         File key2 = findTestFile("test2.com.android.virt.pem");
         Map<String, File> keyOverrides = Map.of("microdroid_vbmeta.img", key2);
-        boolean isProtected = false; // Not interested in pvwfw
-        boolean waitForOnline = false; // Bootloader fails and enters prompts.
         // To be able to stop it, it should be a daemon.
-        String consolePath = TEST_ROOT + "console";
-        Process process =
-                runMicrodroidWithResignedImages(
-                        key, keyOverrides, isProtected, waitForOnline, consolePath);
+        VmInfo vmInfo = runMicrodroidWithResignedImages(key, keyOverrides, /*isProtected=*/ false);
         // Wait so that init can print errors to console (time in cuttlefish >> in real device)
         assertThatEventually(
                 100000,
-                () -> getDevice().pullFileContents(consolePath),
+                () -> getDevice().pullFileContents(CONSOLE_PATH),
                 containsString("init: [libfs_avb]Failed to verify vbmeta digest"));
-        process.destroy();
+        vmInfo.mProcess.destroy();
     }
 
     private boolean isTombstoneGeneratedWithConfig(String configPath) throws Exception {
diff --git a/tests/testapk/src/native/testbinary.cpp b/tests/testapk/src/native/testbinary.cpp
index 5c217ff..c0a8c0e 100644
--- a/tests/testapk/src/native/testbinary.cpp
+++ b/tests/testapk/src/native/testbinary.cpp
@@ -119,10 +119,8 @@
     auto testService = ndk::SharedRefBase::make<TestService>();
 
     auto callback = []([[maybe_unused]] void* param) { AVmPayload_notifyPayloadReady(); };
-    if (!AVmPayload_runVsockRpcServer(testService->asBinder().get(), testService->SERVICE_PORT,
-                                      callback, nullptr)) {
-        return Error() << "RPC Server failed to run";
-    }
+    AVmPayload_runVsockRpcServer(testService->asBinder().get(), testService->SERVICE_PORT, callback,
+                                 nullptr);
 
     return {};
 }
diff --git a/virtualizationservice/src/crosvm.rs b/virtualizationservice/src/crosvm.rs
index 76e18db..85a57c9 100644
--- a/virtualizationservice/src/crosvm.rs
+++ b/virtualizationservice/src/crosvm.rs
@@ -24,6 +24,7 @@
 use semver::{Version, VersionReq};
 use nix::{fcntl::OFlag, unistd::pipe2};
 use regex::{Captures, Regex};
+use rustutils::system_properties;
 use shared_child::SharedChild;
 use std::borrow::Cow;
 use std::cmp::max;
@@ -590,10 +591,15 @@
         .arg("info,disk=off")
         .arg("run")
         .arg("--disable-sandbox")
-        .arg("--no-balloon")
         .arg("--cid")
         .arg(config.cid.to_string());
 
+    if system_properties::read_bool("hypervisor.memory_reclaim.supported", false)? {
+        command.arg("--balloon-page-reporting");
+    } else {
+        command.arg("--no-balloon");
+    }
+
     if config.protected {
         command.arg("--protected-vm");
 
diff --git a/vm_payload/include/vm_payload.h b/vm_payload/include/vm_payload.h
index 0ad4c64..7c224f6 100644
--- a/vm_payload/include/vm_payload.h
+++ b/vm_payload/include/vm_payload.h
@@ -18,6 +18,7 @@
 
 #include <stdbool.h>
 #include <stddef.h>
+#include <stdnoreturn.h>
 #include <sys/cdefs.h>
 
 #include "vm_main.h"
@@ -46,18 +47,17 @@
  * called to allow appropriate action to be taken - e.g. to notify clients that they may now
  * attempt to connect with `AVmPayload_notifyPayloadReady`.
  *
- * The current thread is joined to the binder thread pool to handle incoming messages.
+ * Note that this function does not return. The calling thread joins the binder
+ * thread pool to handle incoming messages.
  *
  * \param service the service to bind to the given port.
  * \param port vsock port.
  * \param on_ready the callback to execute once the server is ready for connections. The callback
  *                 will be called at most once.
  * \param param param for the `on_ready` callback.
- *
- * \return true if the server has shutdown normally, false if it failed in some way.
  */
-bool AVmPayload_runVsockRpcServer(AIBinder *service, unsigned int port,
-                                  void (*on_ready)(void *param), void *param);
+noreturn void AVmPayload_runVsockRpcServer(AIBinder *service, unsigned int port,
+                                           void (*on_ready)(void *param), void *param);
 
 /**
  * Get a secret that is uniquely bound to this VM instance. The secrets are
diff --git a/vm_payload/src/api.rs b/vm_payload/src/api.rs
index febc2be..a79c0bb 100644
--- a/vm_payload/src/api.rs
+++ b/vm_payload/src/api.rs
@@ -14,13 +14,17 @@
 
 //! This module handles the interaction with virtual machine payload service.
 
+// We're implementing unsafe functions, but we still want warnings on unsafe usage within them.
+#![warn(unsafe_op_in_unsafe_fn)]
+
 use android_system_virtualization_payload::aidl::android::system::virtualization::payload::IVmPayloadService::{
     IVmPayloadService, VM_PAYLOAD_SERVICE_SOCKET_NAME, VM_APK_CONTENTS_PATH};
-use anyhow::{ensure, Context, Result};
+use anyhow::{ensure, bail, Context, Result};
 use binder::{Strong, unstable_api::{AIBinder, new_spibinder}};
 use lazy_static::lazy_static;
 use log::{error, info, Level};
 use rpcbinder::{get_unix_domain_rpc_interface, RpcServer};
+use std::convert::Infallible;
 use std::ffi::CString;
 use std::fmt::Debug;
 use std::os::raw::{c_char, c_void};
@@ -96,44 +100,55 @@
 /// called to allow appropriate action to be taken - e.g. to notify clients that they may now
 /// attempt to connect.
 ///
-/// The current thread is joined to the binder thread pool to handle incoming messages.
+/// The current thread joins the binder thread pool to handle incoming messages.
+/// This function never returns.
 ///
-/// Returns true if the server has shutdown normally, false if it failed in some way.
+/// Panics on error (including unexpected server exit).
 ///
 /// # Safety
 ///
-/// The `on_ready` callback is only called inside `run_vsock_rpc_server`, within the lifetime of
-/// `ReadyNotifier` (the last parameter of `run_vsock_rpc_server`). If `on_ready` is called with
-/// wrong param, the callback execution could go wrong.
+/// If present, the `on_ready` callback must be a valid function pointer, which will be called at
+/// most once, while this function is executing, with the `param` parameter.
 #[no_mangle]
 pub unsafe extern "C" fn AVmPayload_runVsockRpcServer(
     service: *mut AIBinder,
     port: u32,
     on_ready: Option<unsafe extern "C" fn(param: *mut c_void)>,
     param: *mut c_void,
-) -> bool {
+) -> Infallible {
     initialize_logging();
 
+    // SAFETY: try_run_vsock_server has the same requirements as this function
+    unwrap_or_abort(unsafe { try_run_vsock_server(service, port, on_ready, param) })
+}
+
+/// # Safety: Same as `AVmPayload_runVsockRpcServer`.
+unsafe fn try_run_vsock_server(
+    service: *mut AIBinder,
+    port: u32,
+    on_ready: Option<unsafe extern "C" fn(param: *mut c_void)>,
+    param: *mut c_void,
+) -> Result<Infallible> {
     // SAFETY: AIBinder returned has correct reference count, and the ownership can
     // safely be taken by new_spibinder.
-    let service = new_spibinder(service);
+    let service = unsafe { new_spibinder(service) };
     if let Some(service) = service {
         match RpcServer::new_vsock(service, port) {
             Ok(server) => {
                 if let Some(on_ready) = on_ready {
-                    on_ready(param);
+                    // SAFETY: We're calling the callback with the parameter specified within the
+                    // allowed lifetime.
+                    unsafe { on_ready(param) };
                 }
                 server.join();
-                true
+                bail!("RpcServer unexpectedly terminated");
             }
             Err(err) => {
-                error!("Failed to start RpcServer: {:?}", err);
-                false
+                bail!("Failed to start RpcServer: {:?}", err);
             }
         }
     } else {
-        error!("Failed to convert the given service from AIBinder to SpIBinder.");
-        false
+        bail!("Failed to convert the given service from AIBinder to SpIBinder.");
     }
 }
 
@@ -157,9 +172,15 @@
 ) {
     initialize_logging();
 
-    let identifier = std::slice::from_raw_parts(identifier, identifier_size);
+    // SAFETY: See the requirements on `identifier` above.
+    let identifier = unsafe { std::slice::from_raw_parts(identifier, identifier_size) };
     let vm_secret = unwrap_or_abort(try_get_vm_instance_secret(identifier, size));
-    ptr::copy_nonoverlapping(vm_secret.as_ptr(), secret, size);
+
+    // SAFETY: See the requirements on `secret` above; `vm_secret` is known to have length `size`,
+    // and cannot overlap `secret` because we just allocated it.
+    unsafe {
+        ptr::copy_nonoverlapping(vm_secret.as_ptr(), secret, size);
+    }
 }
 
 fn try_get_vm_instance_secret(identifier: &[u8], size: usize) -> Result<Vec<u8>> {
@@ -190,7 +211,9 @@
     initialize_logging();
 
     let chain = unwrap_or_abort(try_get_dice_attestation_chain());
-    ptr::copy_nonoverlapping(chain.as_ptr(), data, std::cmp::min(chain.len(), size));
+    // SAFETY: See the requirements on `data` above. The number of bytes copied doesn't exceed
+    // the length of either buffer, and `chain` cannot overlap `data` because we just allocated it.
+    unsafe { ptr::copy_nonoverlapping(chain.as_ptr(), data, std::cmp::min(chain.len(), size)) };
     chain.len()
 }
 
@@ -213,7 +236,9 @@
     initialize_logging();
 
     let cdi = unwrap_or_abort(try_get_dice_attestation_cdi());
-    ptr::copy_nonoverlapping(cdi.as_ptr(), data, std::cmp::min(cdi.len(), size));
+    // SAFETY: See the requirements on `data` above. The number of bytes copied doesn't exceed
+    // the length of either buffer, and `cdi` cannot overlap `data` because we just allocated it.
+    unsafe { ptr::copy_nonoverlapping(cdi.as_ptr(), data, std::cmp::min(cdi.len(), size)) };
     cdi.len()
 }
 
diff --git a/vmbase/example/Android.bp b/vmbase/example/Android.bp
index 0f1e66a..505de6b 100644
--- a/vmbase/example/Android.bp
+++ b/vmbase/example/Android.bp
@@ -12,9 +12,13 @@
         "libaarch64_paging",
         "libbuddy_system_allocator",
         "libdice_nostd",
+        "liblibfdt",
         "liblog_rust_nostd",
         "libvmbase",
     ],
+    static_libs: [
+        "libarm-optimized-routines-mem",
+    ],
     apex_available: ["com.android.virt"],
 }
 
diff --git a/vmbase/example/src/main.rs b/vmbase/example/src/main.rs
index 03f0603..dcff6e1 100644
--- a/vmbase/example/src/main.rs
+++ b/vmbase/example/src/main.rs
@@ -27,9 +27,14 @@
     bionic_tls, dtb_range, print_addresses, rodata_range, stack_chk_guard, text_range,
     writable_region, DEVICE_REGION,
 };
-use aarch64_paging::{idmap::IdMap, paging::Attributes};
+use aarch64_paging::{
+    idmap::IdMap,
+    paging::{Attributes, MemoryRegion},
+};
 use alloc::{vec, vec::Vec};
 use buddy_system_allocator::LockedHeap;
+use core::ffi::CStr;
+use libfdt::Fdt;
 use log::{info, LevelFilter};
 use vmbase::{logger, main, println};
 
@@ -57,6 +62,7 @@
     assert_eq!(arg0, dtb_range().start.0 as u64);
     check_data();
     check_stack_guard();
+    check_fdt();
 
     unsafe {
         HEAP_ALLOCATOR.lock().init(HEAP.as_mut_ptr() as usize, HEAP.len());
@@ -138,6 +144,27 @@
     info!("Data looks good");
 }
 
+fn check_fdt() {
+    info!("Checking FDT...");
+    let fdt = MemoryRegion::from(layout::dtb_range());
+    let fdt = unsafe { core::slice::from_raw_parts_mut(fdt.start().0 as *mut u8, fdt.len()) };
+
+    let reader = Fdt::from_slice(fdt).unwrap();
+    info!("FDT passed verification.");
+    for reg in reader.memory().unwrap() {
+        info!("memory @ {reg:#x?}");
+    }
+
+    let compatible = CStr::from_bytes_with_nul(b"ns16550a\0").unwrap();
+
+    for c in reader.compatible_nodes(compatible).unwrap() {
+        let reg = c.reg().unwrap().next().unwrap();
+        info!("node compatible with '{}' at {reg:?}", compatible.to_str().unwrap());
+    }
+
+    info!("FDT checks done.");
+}
+
 fn check_alloc() {
     info!("Allocating a Vec...");
     let mut vector: Vec<u32> = vec![1, 2, 3, 4];