vmbase: Initialize heap in rust_entry

Initialize the heap as we enter Rust so that any code that follows can
make use of alloc, preventing bugs such as aosp/2567870 where client
code used the heap before initializing it.

Test: TH
Change-Id: Ibe6629a22cd828d8b9c6418b91b5971dbbeb4c3d
diff --git a/pvmfw/src/entry.rs b/pvmfw/src/entry.rs
index 56f9462..10a3e5a 100644
--- a/pvmfw/src/entry.rs
+++ b/pvmfw/src/entry.rs
@@ -31,7 +31,7 @@
 use log::LevelFilter;
 use vmbase::util::RangeExt as _;
 use vmbase::{
-    configure_global_allocator_size, console, heap,
+    configure_global_allocator_size, console,
     layout::{self, crosvm},
     logger, main,
     memory::{min_dcache_line_size, MemoryTracker, MEMORY, SIZE_128KB, SIZE_4KB},
@@ -69,9 +69,6 @@
     // - can't access non-pvmfw memory (only statically-mapped memory)
     // - can't access MMIO (therefore, no logging)
 
-    // SAFETY - This function should and will only be called once, here.
-    unsafe { heap::init() };
-
     match main_wrapper(fdt_address as usize, payload_start as usize, payload_size as usize) {
         Ok((entry, bcc)) => jump_to_payload(fdt_address, entry.try_into().unwrap(), bcc),
         Err(_) => reboot(), // TODO(b/220071963) propagate the reason back to the host.
diff --git a/rialto/src/main.rs b/rialto/src/main.rs
index f3f29fe..7b3783a 100644
--- a/rialto/src/main.rs
+++ b/rialto/src/main.rs
@@ -33,7 +33,6 @@
 use vmbase::{
     configure_global_allocator_size,
     fdt::SwiotlbInfo,
-    heap,
     layout::{self, crosvm},
     main,
     memory::{MemoryTracker, PageTable, MEMORY, PAGE_SIZE, SIZE_64KB},
@@ -149,9 +148,6 @@
 
 /// Entry point for Rialto.
 pub fn main(fdt_addr: u64, _a1: u64, _a2: u64, _a3: u64) {
-    // SAFETY - Only called once, from here.
-    unsafe { heap::init() };
-
     let Ok(mmio_guard_supported) = try_init_logger() else {
         // Don't log anything if the logger initialization fails.
         reboot();
diff --git a/vmbase/example/src/main.rs b/vmbase/example/src/main.rs
index d604509..ccae6ca 100644
--- a/vmbase/example/src/main.rs
+++ b/vmbase/example/src/main.rs
@@ -33,7 +33,7 @@
 use fdtpci::PciInfo;
 use libfdt::Fdt;
 use log::{debug, error, info, trace, warn, LevelFilter};
-use vmbase::{configure_global_allocator_size, cstr, heap, logger, main, memory::SIZE_64KB};
+use vmbase::{configure_global_allocator_size, cstr, logger, main, memory::SIZE_64KB};
 
 static INITIALISED_DATA: [u32; 4] = [1, 2, 3, 4];
 static mut ZEROED_DATA: [u32; 10] = [0; 10];
@@ -69,9 +69,6 @@
 
     modify_fdt(fdt);
 
-    // SAFETY - Only called once, from here.
-    unsafe { heap::init() };
-
     check_alloc();
 
     let mut idmap = IdMap::new(ASID, ROOT_LEVEL);
diff --git a/vmbase/src/entry.rs b/vmbase/src/entry.rs
index 8cdfe77..df0bb7c 100644
--- a/vmbase/src/entry.rs
+++ b/vmbase/src/entry.rs
@@ -14,11 +14,13 @@
 
 //! Rust entry point.
 
-use crate::{console, power::shutdown};
+use crate::{console, heap, power::shutdown};
 
 /// This is the entry point to the Rust code, called from the binary entry point in `entry.S`.
 #[no_mangle]
 extern "C" fn rust_entry(x0: u64, x1: u64, x2: u64, x3: u64) -> ! {
+    // SAFETY - Only called once, from here, and inaccessible to client code.
+    unsafe { heap::init() };
     console::init();
     unsafe {
         main(x0, x1, x2, x3);
diff --git a/vmbase/src/heap.rs b/vmbase/src/heap.rs
index a28a02c..08240b9 100644
--- a/vmbase/src/heap.rs
+++ b/vmbase/src/heap.rs
@@ -51,7 +51,7 @@
 /// # Safety
 ///
 /// Must be called no more than once.
-pub unsafe fn init() {
+pub(crate) unsafe fn init() {
     // SAFETY: Nothing else accesses this memory, and we hand it over to the heap to manage and
     // never touch it again. The heap is locked, so there cannot be any races.
     let (start, size) = unsafe { (HEAP.as_mut_ptr() as usize, HEAP.len()) };