Merge "vmbase: Improve safety of console::init()" into main
diff --git a/pvmfw/src/entry.rs b/pvmfw/src/entry.rs
index cc6e302..36ead46 100644
--- a/pvmfw/src/entry.rs
+++ b/pvmfw/src/entry.rs
@@ -66,7 +66,7 @@
 pub fn start(fdt_address: u64, payload_start: u64, payload_size: u64, _arg3: u64) {
     // Limitations in this function:
     // - can't access non-pvmfw memory (only statically-mapped memory)
-    // - can't access MMIO (therefore, no logging)
+    // - can't access MMIO (except the console, already configured by vmbase)
 
     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),
diff --git a/vmbase/src/console.rs b/vmbase/src/console.rs
index 657a62b..7356d7f 100644
--- a/vmbase/src/console.rs
+++ b/vmbase/src/console.rs
@@ -15,26 +15,33 @@
 //! Console driver for 8250 UART.
 
 use crate::uart::Uart;
-use core::fmt::{write, Arguments, Write};
+use core::{
+    cell::OnceCell,
+    fmt::{write, Arguments, Write},
+};
 use spin::mutex::SpinMutex;
 
-/// Base memory-mapped address of the primary UART device.
-pub const BASE_ADDRESS: usize = 0x3f8;
-
+// ADDRESS is the base of the MMIO region for a UART and must be mapped as device memory.
+static ADDRESS: SpinMutex<OnceCell<usize>> = SpinMutex::new(OnceCell::new());
 static CONSOLE: SpinMutex<Option<Uart>> = SpinMutex::new(None);
 
-/// Initialises a new instance of the UART driver and returns it.
-fn create() -> Uart {
-    // SAFETY: BASE_ADDRESS is the base of the MMIO region for a UART and is mapped as device
-    // memory.
-    unsafe { Uart::new(BASE_ADDRESS) }
-}
+/// Initialises the global instance of the UART driver.
+///
+/// This must be called before using the `print!` and `println!` macros.
+///
+/// # Safety
+///
+/// This must be called with the base of a UART, mapped as device memory and (if necessary) shared
+/// with the host as MMIO.
+pub unsafe fn init(base_address: usize) {
+    // Remember the valid address, for emergency console accesses.
+    ADDRESS.lock().set(base_address).expect("console::init() called more than once");
 
-/// Initialises the global instance of the UART driver. This must be called before using
-/// the `print!` and `println!` macros.
-pub fn init() {
-    let uart = create();
-    CONSOLE.lock().replace(uart);
+    // Initialize the console driver, for normal console accesses.
+    let mut console = CONSOLE.lock();
+    assert!(console.is_none(), "console::init() called more than once");
+    // SAFETY: base_address must be the base of a mapped UART.
+    console.replace(unsafe { Uart::new(base_address) });
 }
 
 /// Writes a formatted string followed by a newline to the console.
@@ -53,7 +60,12 @@
 /// This is intended for use in situations where the UART may be in an unknown state or the global
 /// instance may be locked, such as in an exception handler or panic handler.
 pub fn ewriteln(format_args: Arguments) {
-    let mut uart = create();
+    let Some(cell) = ADDRESS.try_lock() else { return };
+    let Some(addr) = cell.get() else { return };
+
+    // SAFETY: addr contains the base of a mapped UART, passed in init().
+    let mut uart = unsafe { Uart::new(*addr) };
+
     let _ = write(&mut uart, format_args);
     let _ = uart.write_str("\n");
 }
diff --git a/vmbase/src/entry.rs b/vmbase/src/entry.rs
index dedc6ae..4a436d2 100644
--- a/vmbase/src/entry.rs
+++ b/vmbase/src/entry.rs
@@ -16,7 +16,7 @@
 
 use crate::{
     bionic, console, heap, hyp,
-    layout::UART_PAGE_ADDR,
+    layout::{UART_ADDR, UART_PAGE_ADDR},
     logger,
     memory::{SIZE_16KB, SIZE_4KB},
     power::{reboot, shutdown},
@@ -26,8 +26,6 @@
 use static_assertions::const_assert_eq;
 
 fn try_console_init() -> Result<(), hyp::Error> {
-    console::init();
-
     if let Some(mmio_guard) = hyp::get_mmio_guard() {
         mmio_guard.enroll()?;
 
@@ -49,6 +47,9 @@
         mmio_guard.map(UART_PAGE_ADDR)?;
     }
 
+    // SAFETY: UART_PAGE is mapped at stage-1 (see entry.S) and was just MMIO-guarded.
+    unsafe { console::init(UART_ADDR) };
+
     Ok(())
 }
 
diff --git a/vmbase/src/layout.rs b/vmbase/src/layout.rs
index 993141d..e65abda 100644
--- a/vmbase/src/layout.rs
+++ b/vmbase/src/layout.rs
@@ -16,7 +16,6 @@
 
 pub mod crosvm;
 
-use crate::console;
 use crate::linker::__stack_chk_guard;
 use crate::memory::{page_4kb_of, PAGE_SIZE};
 use aarch64_paging::paging::VirtualAddress;
@@ -27,9 +26,14 @@
 /// First address that can't be translated by a level 1 TTBR0_EL1.
 pub const MAX_VIRT_ADDR: usize = 1 << 40;
 
+/// Base memory-mapped address of the primary UART device.
+///
+/// See SERIAL_ADDR in https://crosvm.dev/book/appendix/memory_layout.html#common-layout.
+pub const UART_ADDR: usize = 0x3f8;
+
 /// Address of the single page containing all the UART devices.
 pub const UART_PAGE_ADDR: usize = 0;
-const_assert_eq!(UART_PAGE_ADDR, page_4kb_of(console::BASE_ADDRESS));
+const_assert_eq!(UART_PAGE_ADDR, page_4kb_of(UART_ADDR));
 
 /// Get an address from a linker-defined symbol.
 #[macro_export]