vmbase: Enter clients with dynamic PTs live

The layout and linker sections of clients are controlled by a linker
script provided by vmbase so the library should know how to replicate
the client-specific static PTs to generate the dynamic PTs. Therefore,
teach vmbase to use the information it gets from the linker to replicate
the client static PTs, with client-agnostic code, and switch to the
dynamic PTs before entering its clients.

Bug: 377276983
Test: m {pvmfw,rialto,vmbase_example_{bios,kernel}}_bin
Test: atest rialto_test vmbase_example.integration_test
Change-Id: I90fa0e62e6e4cb8474f3294633cee775b1191c71
diff --git a/libs/libvmbase/src/entry.rs b/libs/libvmbase/src/entry.rs
index 2433722..b681aea 100644
--- a/libs/libvmbase/src/entry.rs
+++ b/libs/libvmbase/src/entry.rs
@@ -18,7 +18,7 @@
     bionic, console, heap,
     layout::{UART_ADDRESSES, UART_PAGE_ADDR},
     logger,
-    memory::{PAGE_SIZE, SIZE_16KB, SIZE_4KB},
+    memory::{switch_to_dynamic_page_tables, PAGE_SIZE, SIZE_16KB, SIZE_4KB},
     power::{reboot, shutdown},
     rand,
 };
@@ -82,6 +82,8 @@
 
     bionic::__get_tls().stack_guard = u64::from_ne_bytes(stack_guard);
 
+    switch_to_dynamic_page_tables();
+
     // Note: If rust_entry ever returned (which it shouldn't by being -> !), the compiler-injected
     // stack guard comparison would detect a mismatch and call __stack_chk_fail.
 
diff --git a/libs/libvmbase/src/memory.rs b/libs/libvmbase/src/memory.rs
index e903858..fd4706f 100644
--- a/libs/libvmbase/src/memory.rs
+++ b/libs/libvmbase/src/memory.rs
@@ -27,8 +27,8 @@
 pub use shared::MemoryRange;
 pub use tracker::{
     deactivate_dynamic_page_tables, init_shared_pool, map_data, map_device, map_image_footer,
-    map_rodata, map_rodata_outside_main_memory, resize_available_memory,
-    switch_to_dynamic_page_tables, unshare_all_memory, unshare_all_mmio_except_uart, unshare_uart,
+    map_rodata, map_rodata_outside_main_memory, resize_available_memory, unshare_all_memory,
+    unshare_all_mmio_except_uart, unshare_uart,
 };
 pub use util::{
     flush, flushed_zeroize, page_4kb_of, PAGE_SIZE, SIZE_128KB, SIZE_16KB, SIZE_2MB, SIZE_4KB,
@@ -37,5 +37,5 @@
 
 pub(crate) use shared::{alloc_shared, dealloc_shared};
 pub(crate) use stack::max_stack_size;
-pub(crate) use tracker::MEMORY;
+pub(crate) use tracker::{switch_to_dynamic_page_tables, MEMORY};
 pub(crate) use util::{phys_to_virt, virt_to_phys};
diff --git a/libs/libvmbase/src/memory/tracker.rs b/libs/libvmbase/src/memory/tracker.rs
index e75dc11..3416dc6 100644
--- a/libs/libvmbase/src/memory/tracker.rs
+++ b/libs/libvmbase/src/memory/tracker.rs
@@ -52,14 +52,13 @@
 /// Switch the MMU to the provided PageTable.
 ///
 /// Panics if called more than once.
-pub fn switch_to_dynamic_page_tables(pt: PageTable) {
+pub(crate) fn switch_to_dynamic_page_tables() {
     let mut locked_tracker = try_lock_memory_tracker().unwrap();
     if locked_tracker.is_some() {
         panic!("switch_to_dynamic_page_tables() called more than once.");
     }
 
     locked_tracker.replace(MemoryTracker::new(
-        pt,
         layout::crosvm::MEM_START..layout::MAX_VIRT_ADDR,
         layout::crosvm::MMIO_RANGE,
     ));
@@ -207,12 +206,13 @@
     const MMIO_CAPACITY: usize = 5;
 
     /// Creates a new instance from an active page table, covering the maximum RAM size.
-    fn new(mut page_table: PageTable, total: MemoryRange, mmio_range: MemoryRange) -> Self {
+    fn new(total: MemoryRange, mmio_range: MemoryRange) -> Self {
         assert!(
             !total.overlaps(&mmio_range),
             "MMIO space should not overlap with the main memory region."
         );
 
+        let mut page_table = Self::initialize_dynamic_page_tables();
         // Activate dirty state management first, otherwise we may get permission faults immediately
         // after activating the new page table. This has no effect before the new page table is
         // activated because none of the entries in the initial idmap have the DBM flag.
@@ -503,6 +503,28 @@
             .modify_range(&(addr..addr + 1).into(), &mark_dirty_block)
             .map_err(|_| MemoryTrackerError::SetPteDirtyFailed)
     }
+
+    // TODO(ptosi): Move this and `PageTable` references to crate::arch::aarch64
+    /// Produces a `PageTable` that can safely replace the static PTs.
+    fn initialize_dynamic_page_tables() -> PageTable {
+        let text = layout::text_range();
+        let rodata = layout::rodata_range();
+        let data_bss = layout::data_bss_range();
+        let eh_stack = layout::eh_stack_range();
+        let stack = layout::stack_range();
+        let console_uart_page = layout::console_uart_page();
+
+        let mut page_table = PageTable::default();
+
+        page_table.map_device(&console_uart_page.into()).unwrap();
+        page_table.map_code(&text.into()).unwrap();
+        page_table.map_rodata(&rodata.into()).unwrap();
+        page_table.map_data(&data_bss.into()).unwrap();
+        page_table.map_data(&eh_stack.into()).unwrap();
+        page_table.map_data(&stack.into()).unwrap();
+
+        page_table
+    }
 }
 
 impl Drop for MemoryTracker {