Merge changes I90fa0e62,I8df67629,Iefc54070 into main

* changes:
  vmbase: Enter clients with dynamic PTs live
  vmbase: Default to largest stack size possible
  vmbase_example: Upgrade to new vmbase::memory API
diff --git a/guest/pvmfw/src/entry.rs b/guest/pvmfw/src/entry.rs
index 0607a5a..2f0b391 100644
--- a/guest/pvmfw/src/entry.rs
+++ b/guest/pvmfw/src/entry.rs
@@ -26,11 +26,10 @@
 use vmbase::util::RangeExt as _;
 use vmbase::{
     arch::aarch64::min_dcache_line_size,
-    configure_heap, console_writeln, layout, main,
+    configure_heap, console_writeln, layout, limit_stack_size, main,
     memory::{
-        deactivate_dynamic_page_tables, map_image_footer, switch_to_dynamic_page_tables,
-        unshare_all_memory, unshare_all_mmio_except_uart, unshare_uart, MemoryTrackerError,
-        SIZE_128KB, SIZE_4KB,
+        deactivate_dynamic_page_tables, map_image_footer, unshare_all_memory,
+        unshare_all_mmio_except_uart, unshare_uart, MemoryTrackerError, SIZE_128KB, SIZE_4KB,
     },
     power::reboot,
 };
@@ -73,6 +72,7 @@
 
 main!(start);
 configure_heap!(SIZE_128KB);
+limit_stack_size!(SIZE_4KB * 12);
 
 /// Entry point for pVM firmware.
 pub fn start(fdt_address: u64, payload_start: u64, payload_size: u64, _arg3: u64) {
@@ -108,13 +108,6 @@
 
     log::set_max_level(LevelFilter::Info);
 
-    let page_table = memory::init_page_table().map_err(|e| {
-        error!("Failed to set up the dynamic page tables: {e}");
-        RebootReason::InternalError
-    })?;
-    // Up to this point, we were using the built-in static (from .rodata) page tables.
-    switch_to_dynamic_page_tables(page_table);
-
     let appended_data = get_appended_data_slice().map_err(|e| {
         error!("Failed to map the appended data: {e}");
         RebootReason::InternalError
@@ -190,7 +183,7 @@
     assert_eq!(bcc.start % ASM_STP_ALIGN, 0, "Misaligned guest BCC.");
     assert_eq!(bcc.end % ASM_STP_ALIGN, 0, "Misaligned guest BCC.");
 
-    let stack = memory::stack_range();
+    let stack = layout::stack_range();
 
     assert_ne!(stack.end - stack.start, 0, "stack region is empty.");
     assert_eq!(stack.start.0 % ASM_STP_ALIGN, 0, "Misaligned stack region.");
diff --git a/guest/pvmfw/src/helpers.rs b/guest/pvmfw/src/helpers.rs
index 8981408..0552640 100644
--- a/guest/pvmfw/src/helpers.rs
+++ b/guest/pvmfw/src/helpers.rs
@@ -14,7 +14,6 @@
 
 //! Miscellaneous helper functions.
 
-use vmbase::memory::{PAGE_SIZE, SIZE_4KB};
+use vmbase::memory::SIZE_4KB;
 
 pub const GUEST_PAGE_SIZE: usize = SIZE_4KB;
-pub const PVMFW_PAGE_SIZE: usize = PAGE_SIZE;
diff --git a/guest/pvmfw/src/memory.rs b/guest/pvmfw/src/memory.rs
index b54f014..35bfd3a 100644
--- a/guest/pvmfw/src/memory.rs
+++ b/guest/pvmfw/src/memory.rs
@@ -16,46 +16,17 @@
 
 use crate::entry::RebootReason;
 use crate::fdt;
-use crate::helpers::PVMFW_PAGE_SIZE;
-use aarch64_paging::paging::VirtualAddress;
-use aarch64_paging::MapError;
 use core::num::NonZeroUsize;
-use core::ops::Range;
-use core::result;
 use core::slice;
 use log::debug;
 use log::error;
 use log::info;
 use log::warn;
 use vmbase::{
-    layout::{self, crosvm},
-    memory::{init_shared_pool, map_data, map_rodata, resize_available_memory, PageTable},
+    layout::crosvm,
+    memory::{init_shared_pool, map_data, map_rodata, resize_available_memory},
 };
 
-/// Region allocated for the stack.
-pub fn stack_range() -> Range<VirtualAddress> {
-    const STACK_PAGES: usize = 12;
-
-    layout::stack_range(STACK_PAGES * PVMFW_PAGE_SIZE)
-}
-
-pub fn init_page_table() -> result::Result<PageTable, MapError> {
-    let mut page_table = PageTable::default();
-
-    // Stack and scratch ranges are explicitly zeroed and flushed before jumping to payload,
-    // so dirty state management can be omitted.
-    page_table.map_data(&layout::data_bss_range().into())?;
-    page_table.map_data(&layout::eh_stack_range().into())?;
-    page_table.map_data(&stack_range().into())?;
-    page_table.map_code(&layout::text_range().into())?;
-    page_table.map_rodata(&layout::rodata_range().into())?;
-    if let Err(e) = page_table.map_device(&layout::console_uart_page().into()) {
-        error!("Failed to remap the UART as a dynamic page table entry: {e}");
-        return Err(e);
-    }
-    Ok(page_table)
-}
-
 pub(crate) struct MemorySlices<'a> {
     pub fdt: &'a mut libfdt::Fdt,
     pub kernel: &'a [u8],
diff --git a/guest/rialto/src/main.rs b/guest/rialto/src/main.rs
index 8095a1f..04d18be 100644
--- a/guest/rialto/src/main.rs
+++ b/guest/rialto/src/main.rs
@@ -46,11 +46,11 @@
     fdt::pci::PciInfo,
     fdt::SwiotlbInfo,
     generate_image_header,
-    layout::{self, crosvm},
+    layout::crosvm,
     main,
     memory::{
         init_shared_pool, map_rodata, map_rodata_outside_main_memory, resize_available_memory,
-        switch_to_dynamic_page_tables, PageTable, PAGE_SIZE, SIZE_128KB,
+        SIZE_128KB,
     },
     power::reboot,
     virtio::{
@@ -71,28 +71,12 @@
     }
 }
 
-fn new_page_table() -> Result<PageTable> {
-    let mut page_table = PageTable::default();
-
-    page_table.map_data(&layout::data_bss_range().into())?;
-    page_table.map_data(&layout::eh_stack_range().into())?;
-    page_table.map_data(&layout::stack_range(40 * PAGE_SIZE).into())?;
-    page_table.map_code(&layout::text_range().into())?;
-    page_table.map_rodata(&layout::rodata_range().into())?;
-    page_table.map_device(&layout::console_uart_page().into())?;
-
-    Ok(page_table)
-}
-
 /// # Safety
 ///
 /// Behavior is undefined if any of the following conditions are violated:
 /// * The `fdt_addr` must be a valid pointer and points to a valid `Fdt`.
 unsafe fn try_main(fdt_addr: usize) -> Result<()> {
     info!("Welcome to Rialto!");
-    let page_table = new_page_table()?;
-
-    switch_to_dynamic_page_tables(page_table);
 
     let fdt_size = NonZeroUsize::new(crosvm::FDT_MAX_SIZE).unwrap();
     map_rodata(fdt_addr, fdt_size)?;
diff --git a/guest/vmbase_example/src/layout.rs b/guest/vmbase_example/src/layout.rs
index 4e87e4e..bafce10 100644
--- a/guest/vmbase_example/src/layout.rs
+++ b/guest/vmbase_example/src/layout.rs
@@ -14,15 +14,8 @@
 
 //! Memory layout.
 
-use aarch64_paging::paging::VirtualAddress;
-use core::ops::Range;
 use log::info;
-use vmbase::{layout, memory::PAGE_SIZE};
-
-/// Writable data region for the stack.
-pub fn boot_stack_range() -> Range<VirtualAddress> {
-    layout::stack_range(40 * PAGE_SIZE)
-}
+use vmbase::layout;
 
 pub fn print_addresses() {
     let text = layout::text_range();
@@ -40,7 +33,7 @@
     );
     let bss = layout::bss_range();
     info!("bss:        {}..{} ({} bytes)", bss.start, bss.end, bss.end - bss.start);
-    let boot_stack = boot_stack_range();
+    let boot_stack = layout::stack_range();
     info!(
         "boot_stack: {}..{} ({} bytes)",
         boot_stack.start,
diff --git a/guest/vmbase_example/src/main.rs b/guest/vmbase_example/src/main.rs
index f00effa..4c5e880 100644
--- a/guest/vmbase_example/src/main.rs
+++ b/guest/vmbase_example/src/main.rs
@@ -23,12 +23,9 @@
 
 extern crate alloc;
 
-use crate::layout::{boot_stack_range, print_addresses};
-use crate::pci::{check_pci, get_bar_region, get_cam_region};
-use aarch64_paging::paging::VirtualAddress;
-use aarch64_paging::MapError;
+use crate::layout::print_addresses;
+use crate::pci::check_pci;
 use alloc::{vec, vec::Vec};
-use core::mem;
 use core::ptr::addr_of_mut;
 use cstr::cstr;
 use libfdt::Fdt;
@@ -37,12 +34,9 @@
     bionic, configure_heap,
     fdt::pci::PciInfo,
     generate_image_header,
-    layout::{
-        console_uart_page, crosvm::FDT_MAX_SIZE, data_bss_range, eh_stack_range, rodata_range,
-        text_range,
-    },
+    layout::crosvm::FDT_MAX_SIZE,
     linker, logger, main,
-    memory::{PageTable, SIZE_64KB},
+    memory::{deactivate_dynamic_page_tables, map_data, SIZE_64KB},
 };
 
 static INITIALISED_DATA: [u32; 4] = [1, 2, 3, 4];
@@ -53,25 +47,6 @@
 main!(main);
 configure_heap!(SIZE_64KB);
 
-fn init_page_table(page_table: &mut PageTable) -> Result<(), MapError> {
-    page_table.map_device(&console_uart_page().into())?;
-    page_table.map_code(&text_range().into())?;
-    page_table.map_rodata(&rodata_range().into())?;
-    page_table.map_data(&data_bss_range().into())?;
-    page_table.map_data(&eh_stack_range().into())?;
-    page_table.map_data(&boot_stack_range().into())?;
-
-    info!("Activating IdMap...");
-    // SAFETY: page_table duplicates the static mappings for everything that the Rust code is
-    // aware of so activating it shouldn't have any visible effect.
-    unsafe {
-        page_table.activate();
-    }
-    info!("Activated.");
-
-    Ok(())
-}
-
 /// Entry point for VM bootloader.
 pub fn main(arg0: u64, arg1: u64, arg2: u64, arg3: u64) {
     log::set_max_level(LevelFilter::Debug);
@@ -82,15 +57,11 @@
     check_data();
     check_stack_guard();
 
-    let mut page_table = PageTable::default();
-    init_page_table(&mut page_table).unwrap();
-
     info!("Checking FDT...");
     let fdt_addr = usize::try_from(arg0).unwrap();
     // SAFETY: The DTB range is valid, writable memory, and we don't construct any aliases to it.
     let fdt = unsafe { core::slice::from_raw_parts_mut(fdt_addr as *mut u8, FDT_MAX_SIZE) };
-    let fdt_region = (VirtualAddress(fdt_addr)..VirtualAddress(fdt_addr + fdt.len())).into();
-    page_table.map_data(&fdt_region).unwrap();
+    map_data(fdt_addr, FDT_MAX_SIZE.try_into().unwrap()).unwrap();
     let fdt = Fdt::from_mut_slice(fdt).unwrap();
     info!("FDT passed verification.");
     check_fdt(fdt);
@@ -101,23 +72,16 @@
     modify_fdt(fdt);
 
     check_alloc();
-
-    let cam_region = get_cam_region(&pci_info);
-    page_table.map_device(&cam_region).unwrap();
-    let bar_region = get_bar_region(&pci_info);
-    page_table.map_device(&bar_region).unwrap();
-
     check_data();
     check_dice();
 
-    // SAFETY: This is the only place where `make_pci_root` is called.
-    let mut pci_root = unsafe { pci_info.make_pci_root() };
+    let mut pci_root = vmbase::virtio::pci::initialize(pci_info).unwrap();
     check_pci(&mut pci_root);
 
     emit_suppressed_log();
 
     info!("De-activating IdMap...");
-    mem::drop(page_table); // Release PageTable and switch back to idmap.S
+    deactivate_dynamic_page_tables();
     info!("De-activated.");
 }
 
diff --git a/guest/vmbase_example/src/pci.rs b/guest/vmbase_example/src/pci.rs
index 379425d..32ab9f6 100644
--- a/guest/vmbase_example/src/pci.rs
+++ b/guest/vmbase_example/src/pci.rs
@@ -14,7 +14,6 @@
 
 //! Functions to scan the PCI bus for VirtIO device.
 
-use aarch64_paging::paging::MemoryRegion;
 use alloc::alloc::{alloc_zeroed, dealloc, handle_alloc_error, Layout};
 use core::{mem::size_of, ptr::NonNull};
 use log::{debug, info};
@@ -26,10 +25,7 @@
     },
     BufferDirection, Error, Hal, PhysAddr, PAGE_SIZE,
 };
-use vmbase::{
-    fdt::pci::PciInfo,
-    virtio::pci::{self, PciTransportIterator},
-};
+use vmbase::virtio::pci::{self, PciTransportIterator};
 
 /// The standard sector size of a VirtIO block device, in bytes.
 const SECTOR_SIZE_BYTES: usize = 512;
@@ -115,16 +111,6 @@
     info!("Wrote to VirtIO console.");
 }
 
-/// Gets the memory region in which BARs are allocated.
-pub fn get_bar_region(pci_info: &PciInfo) -> MemoryRegion {
-    MemoryRegion::new(pci_info.bar_range.start as usize, pci_info.bar_range.end as usize)
-}
-
-/// Gets the PCI CAM memory region.
-pub fn get_cam_region(pci_info: &PciInfo) -> MemoryRegion {
-    MemoryRegion::new(pci_info.cam_range.start, pci_info.cam_range.end)
-}
-
 struct HalImpl;
 
 /// SAFETY: See the 'Implementation Safety' comments on methods below for how they fulfill the
diff --git a/libs/libvmbase/sections.ld b/libs/libvmbase/sections.ld
index 222edae..9d69935 100644
--- a/libs/libvmbase/sections.ld
+++ b/libs/libvmbase/sections.ld
@@ -132,3 +132,10 @@
 		*(.note.gnu.build-id)
 	}
 }
+
+/*
+ * Make calling the limit_stack_size!() macro optional by providing a default.
+ */
+PROVIDE(vmbase_stack_limit = DEFINED(vmbase_stack_limit_client) ?
+                                     vmbase_stack_limit_client :
+                                     vmbase_stack_limit_default);
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/layout.rs b/libs/libvmbase/src/layout.rs
index a8f7827..9a702b0 100644
--- a/libs/libvmbase/src/layout.rs
+++ b/libs/libvmbase/src/layout.rs
@@ -17,7 +17,7 @@
 pub mod crosvm;
 
 use crate::linker::__stack_chk_guard;
-use crate::memory::{page_4kb_of, PAGE_SIZE};
+use crate::memory::{max_stack_size, page_4kb_of, PAGE_SIZE};
 use aarch64_paging::paging::VirtualAddress;
 use core::ops::Range;
 use core::ptr::addr_of;
@@ -91,10 +91,16 @@
 }
 
 /// Writable data region for the stack.
-pub fn stack_range(stack_size: usize) -> Range<VirtualAddress> {
+pub fn stack_range() -> Range<VirtualAddress> {
     let end = linker_addr!(init_stack_pointer);
-    let start = VirtualAddress(end.0.checked_sub(stack_size).unwrap());
-    assert!(start >= linker_addr!(stack_limit));
+    let start = if let Some(stack_size) = max_stack_size() {
+        assert_eq!(stack_size % PAGE_SIZE, 0);
+        let start = VirtualAddress(end.0.checked_sub(stack_size).unwrap());
+        assert!(start >= linker_addr!(stack_limit));
+        start
+    } else {
+        linker_addr!(stack_limit)
+    };
 
     start..end
 }
diff --git a/libs/libvmbase/src/memory.rs b/libs/libvmbase/src/memory.rs
index 145f766..fd4706f 100644
--- a/libs/libvmbase/src/memory.rs
+++ b/libs/libvmbase/src/memory.rs
@@ -18,6 +18,7 @@
 mod error;
 mod page_table;
 mod shared;
+mod stack;
 mod tracker;
 mod util;
 
@@ -26,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,
@@ -35,5 +36,6 @@
 };
 
 pub(crate) use shared::{alloc_shared, dealloc_shared};
-pub(crate) use tracker::MEMORY;
+pub(crate) use stack::max_stack_size;
+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/stack.rs b/libs/libvmbase/src/memory/stack.rs
new file mode 100644
index 0000000..639029e
--- /dev/null
+++ b/libs/libvmbase/src/memory/stack.rs
@@ -0,0 +1,41 @@
+// Copyright 2024, 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.
+
+//! Low-level stack support.
+
+/// Configures the maximum size of the stack.
+#[macro_export]
+macro_rules! limit_stack_size {
+    ($len:expr) => {
+        #[export_name = "vmbase_stack_limit_client"]
+        fn __vmbase_stack_limit_client() -> Option<usize> {
+            Some($len)
+        }
+    };
+}
+
+pub(crate) fn max_stack_size() -> Option<usize> {
+    extern "Rust" {
+        fn vmbase_stack_limit() -> Option<usize>;
+    }
+    // SAFETY: This function is safe to call as the linker script aliases it to either:
+    // - the safe vmbase_stack_limit_default();
+    // - the safe vmbase_stack_limit_client() potentially defined using limit_stack_size!()
+    unsafe { vmbase_stack_limit() }
+}
+
+#[no_mangle]
+fn vmbase_stack_limit_default() -> Option<usize> {
+    None
+}
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 {