pvmfw: Add MemoryTracker & MemorySlices

Add support for dynamically validating, mapping, and allocating
non-overlapping slices of main RAM in a safe manner in MemoryTracker and
use it in MemorySlices for the regions of interest, located through the
inputs.

Bug: 256148034
Bug: 249054080
Bug: 256827715
Test: atest MicrodroidTestApp
Change-Id: I25f8db10df763c0473c281bf96f30e6b5fbe1619
diff --git a/pvmfw/src/entry.rs b/pvmfw/src/entry.rs
index 07dd0c5..18ff1c1 100644
--- a/pvmfw/src/entry.rs
+++ b/pvmfw/src/entry.rs
@@ -14,14 +14,18 @@
 
 //! Low-level entry and exit points of pvmfw.
 
+use crate::fdt;
 use crate::heap;
 use crate::helpers;
+use crate::memory::MemoryTracker;
 use crate::mmio_guard;
 use crate::mmu;
 use core::arch::asm;
+use core::num::NonZeroUsize;
 use core::slice;
 use log::debug;
 use log::error;
+use log::info;
 use log::LevelFilter;
 use vmbase::{console, layout, logger, main, power::reboot};
 
@@ -31,6 +35,12 @@
     InvalidBcc,
     /// An unexpected internal error happened.
     InternalError,
+    /// The provided FDT was invalid.
+    InvalidFdt,
+    /// The provided payload was invalid.
+    InvalidPayload,
+    /// The provided ramdisk was invalid.
+    InvalidRamdisk,
 }
 
 main!(start);
@@ -49,6 +59,98 @@
     // if we reach this point and return, vmbase::entry::rust_entry() will call power::shutdown().
 }
 
+struct MemorySlices<'a> {
+    fdt: &'a mut libfdt::Fdt,
+    kernel: &'a [u8],
+    ramdisk: Option<&'a [u8]>,
+}
+
+impl<'a> MemorySlices<'a> {
+    fn new(
+        fdt: usize,
+        payload: usize,
+        payload_size: usize,
+        memory: &mut MemoryTracker,
+    ) -> Result<Self, RebootReason> {
+        // SAFETY - SIZE_2MB is non-zero.
+        const FDT_SIZE: NonZeroUsize = unsafe { NonZeroUsize::new_unchecked(helpers::SIZE_2MB) };
+        // TODO - Only map the FDT as read-only, until we modify it right before jump_to_payload()
+        // e.g. by generating a DTBO for a template DT in main() and, on return, re-map DT as RW,
+        // overwrite with the template DT and apply the DTBO.
+        let range = memory.alloc_mut(fdt, FDT_SIZE).map_err(|e| {
+            error!("Failed to allocate the FDT range: {e}");
+            RebootReason::InternalError
+        })?;
+
+        // SAFETY - The tracker validated the range to be in main memory, mapped, and not overlap.
+        let fdt = unsafe { slice::from_raw_parts_mut(range.start as *mut u8, range.len()) };
+        let fdt = libfdt::Fdt::from_mut_slice(fdt).map_err(|e| {
+            error!("Failed to spawn the FDT wrapper: {e}");
+            RebootReason::InvalidFdt
+        })?;
+
+        debug!("Fdt passed validation!");
+
+        let memory_range = fdt
+            .memory()
+            .map_err(|e| {
+                error!("Failed to get /memory from the DT: {e}");
+                RebootReason::InvalidFdt
+            })?
+            .ok_or_else(|| {
+                error!("Node /memory was found empty");
+                RebootReason::InvalidFdt
+            })?
+            .next()
+            .ok_or_else(|| {
+                error!("Failed to read the memory size from the FDT");
+                RebootReason::InternalError
+            })?;
+
+        debug!("Resizing MemoryTracker to range {memory_range:#x?}");
+
+        memory.shrink(&memory_range).map_err(|_| {
+            error!("Failed to use memory range value from DT: {memory_range:#x?}");
+            RebootReason::InvalidFdt
+        })?;
+
+        let payload_size = NonZeroUsize::new(payload_size).ok_or_else(|| {
+            error!("Invalid payload size: {payload_size:#x}");
+            RebootReason::InvalidPayload
+        })?;
+
+        let payload_range = memory.alloc(payload, payload_size).map_err(|e| {
+            error!("Failed to obtain the payload range: {e}");
+            RebootReason::InternalError
+        })?;
+        // 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()) };
+
+        let ramdisk_range = fdt::initrd_range(fdt).map_err(|e| {
+            error!("An error occurred while locating the ramdisk in the device tree: {e}");
+            RebootReason::InternalError
+        })?;
+
+        let ramdisk = if let Some(r) = ramdisk_range {
+            debug!("Located ramdisk at {r:?}");
+            let r = memory.alloc_range(&r).map_err(|e| {
+                error!("Failed to obtain the initrd range: {e}");
+                RebootReason::InvalidRamdisk
+            })?;
+
+            // SAFETY - The region was validated by memory to be in main memory, mapped, and
+            // not overlap.
+            Some(unsafe { slice::from_raw_parts(r.start as *const u8, r.len()) })
+        } else {
+            info!("Couldn't locate the ramdisk from the device tree");
+            None
+        };
+
+        Ok(Self { fdt, kernel, ramdisk })
+    }
+}
+
 /// Sets up the environment for main() and wraps its result for start().
 ///
 /// Provide the abstractions necessary for start() to abort the pVM boot and for main() to run with
@@ -64,14 +166,6 @@
 
     logger::init(LevelFilter::Info).map_err(|_| RebootReason::InternalError)?;
 
-    const FDT_MAX_SIZE: usize = helpers::SIZE_2MB;
-    // TODO: Check that the FDT is fully contained in RAM.
-    // SAFETY - We trust the VMM, for now.
-    let fdt = unsafe { slice::from_raw_parts_mut(fdt as *mut u8, FDT_MAX_SIZE) };
-    // TODO: Check that the payload is fully contained in RAM and doesn't overlap with the FDT.
-    // SAFETY - We trust the VMM, for now.
-    let payload = unsafe { slice::from_raw_parts(payload as *const u8, payload_size) };
-
     // Use debug!() to avoid printing to the UART if we failed to configure it as only local
     // builds that have tweaked the logger::init() call will actually attempt to log the message.
 
@@ -121,8 +215,11 @@
     unsafe { page_table.activate() };
     debug!("... Success!");
 
+    let mut memory = MemoryTracker::new(page_table);
+    let slices = MemorySlices::new(fdt, payload, payload_size, &mut memory)?;
+
     // This wrapper allows main() to be blissfully ignorant of platform details.
-    crate::main(fdt, payload, bcc);
+    crate::main(slices.fdt, slices.kernel, slices.ramdisk, bcc);
 
     // TODO: Overwrite BCC before jumping to payload to avoid leaking our sealing key.