pvmfw: Receive payload location from DT

To unify the cases where pvmfw is loaded into pVM address space by pKVM
(the VMM is then allowed to set registers x1 and x2) or by the VMM
itself (all registers except x0 are silently zeroed by pKVM), locate the
pre-loaded kernel through the DT, in a similar way to the ramdisk.

Keep supporting the x1/x2 ABI when the legacy feature is enabled.

Bug: 253616035
Test: Boot pvmfw & check logs with a crosvm with crrev.com/c/4064759
Change-Id: I74d3e52025fa1fbc5e74246e9fb09a7fb46222a0
diff --git a/pvmfw/src/entry.rs b/pvmfw/src/entry.rs
index b840488..342c8ed 100644
--- a/pvmfw/src/entry.rs
+++ b/pvmfw/src/entry.rs
@@ -72,8 +72,8 @@
 impl<'a> MemorySlices<'a> {
     fn new(
         fdt: usize,
-        payload: usize,
-        payload_size: usize,
+        kernel: usize,
+        kernel_size: usize,
         memory: &mut MemoryTracker,
     ) -> Result<Self, RebootReason> {
         // SAFETY - SIZE_2MB is non-zero.
@@ -118,18 +118,36 @@
             RebootReason::InvalidFdt
         })?;
 
-        let payload_size = NonZeroUsize::new(payload_size).ok_or_else(|| {
-            error!("Invalid payload size: {payload_size:#x}");
-            RebootReason::InvalidPayload
+        let kernel_range = fdt::kernel_range(fdt).map_err(|e| {
+            error!("Error while attempting to read the kernel range from the DT: {e}");
+            RebootReason::InvalidFdt
         })?;
 
-        let payload_range = memory.alloc(payload, payload_size).map_err(|e| {
-            error!("Failed to obtain the payload range: {e}");
-            RebootReason::InternalError
-        })?;
+        let kernel_range = if let Some(r) = kernel_range {
+            memory.alloc_range(&r).map_err(|e| {
+                error!("Failed to obtain the kernel range with DT range: {e}");
+                RebootReason::InternalError
+            })?
+        } else if cfg!(feature = "legacy") {
+            warn!("Failed to find the kernel range in the DT; falling back to legacy ABI");
+
+            let kernel_size = NonZeroUsize::new(kernel_size).ok_or_else(|| {
+                error!("Invalid kernel size: {kernel_size:#x}");
+                RebootReason::InvalidPayload
+            })?;
+
+            memory.alloc(kernel, kernel_size).map_err(|e| {
+                error!("Failed to obtain the kernel range with legacy range: {e}");
+                RebootReason::InternalError
+            })?
+        } else {
+            error!("Failed to locate the kernel from the DT");
+            return Err(RebootReason::InvalidPayload);
+        };
+
         // SAFETY - The tracker validated the range to be in main memory, mapped, and not overlap.
         let kernel =
-            unsafe { slice::from_raw_parts(payload_range.start as *const u8, payload_range.len()) };
+            unsafe { slice::from_raw_parts(kernel_range.start as *const u8, kernel_range.len()) };
 
         let ramdisk_range = fdt::initrd_range(fdt).map_err(|e| {
             error!("An error occurred while locating the ramdisk in the device tree: {e}");
diff --git a/pvmfw/src/fdt.rs b/pvmfw/src/fdt.rs
index 5b9efd2..dcd17b7 100644
--- a/pvmfw/src/fdt.rs
+++ b/pvmfw/src/fdt.rs
@@ -17,6 +17,24 @@
 use core::ffi::CStr;
 use core::ops::Range;
 
+/// Extract from /config the address range containing the pre-loaded kernel.
+pub fn kernel_range(fdt: &libfdt::Fdt) -> libfdt::Result<Option<Range<usize>>> {
+    let config = CStr::from_bytes_with_nul(b"/config\0").unwrap();
+    let addr = CStr::from_bytes_with_nul(b"kernel-address\0").unwrap();
+    let size = CStr::from_bytes_with_nul(b"kernel-size\0").unwrap();
+
+    if let Some(config) = fdt.node(config)? {
+        if let (Some(addr), Some(size)) = (config.getprop_u32(addr)?, config.getprop_u32(size)?) {
+            let addr = addr as usize;
+            let size = size as usize;
+
+            return Ok(Some(addr..(addr + size)));
+        }
+    }
+
+    Ok(None)
+}
+
 /// Extract from /chosen the address range containing the pre-loaded ramdisk.
 pub fn initrd_range(fdt: &libfdt::Fdt) -> libfdt::Result<Option<Range<usize>>> {
     let start = CStr::from_bytes_with_nul(b"linux,initrd-start\0").unwrap();