vmbase: Separate console form aarch64 UART implementation

Number and and initialization code of UART may be different between platforms, for example x86_64 may use I/O ports other platform may use MMIO. This commit separate console logic from platform dependent UART implementations.

Bug: 362733888
Test: m {libvmbase, pvmfw, vmbase_example_{bios, kernel}}_bin
test: atest rialto_test vmbase_example.integration_test

Change-Id: I82360003e4511aa760831135f826e9a4a28cc902
diff --git a/libs/libvmbase/src/arch/aarch64/platform.rs b/libs/libvmbase/src/arch/aarch64/platform.rs
index 6985445..6a62f8c 100644
--- a/libs/libvmbase/src/arch/aarch64/platform.rs
+++ b/libs/libvmbase/src/arch/aarch64/platform.rs
@@ -14,10 +14,109 @@
 
 //! Definition of platform
 
+use crate::{
+    arch::aarch64::{
+        layout::{UART_ADDRESSES, UART_PAGE_ADDR},
+        uart::Uart,
+    },
+    memory::{SIZE_16KB, SIZE_4KB},
+};
 use smccc::{
     psci::{system_off, system_reset},
     Hvc,
 };
+use spin::{mutex::SpinMutex, once::Once};
+use static_assertions::const_assert_eq;
+
+// Arbitrary limit on the number of consoles that can be registered.
+//
+// Matches the UART count in crosvm.
+const MAX_CONSOLES: usize = 4;
+
+static CONSOLES: [Once<SpinMutex<Uart>>; MAX_CONSOLES] =
+    [Once::new(), Once::new(), Once::new(), Once::new()];
+static ADDRESSES: [Once<usize>; MAX_CONSOLES] =
+    [Once::new(), Once::new(), Once::new(), Once::new()];
+
+/// Index of the console used by default for logging.
+pub const DEFAULT_CONSOLE_INDEX: usize = 0;
+
+/// Index of the console used by default for emergency logging.
+pub const DEFAULT_EMERGENCY_CONSOLE_INDEX: usize = DEFAULT_CONSOLE_INDEX;
+
+/// Initialises the global instance(s) of the UART driver.
+///
+/// # Safety
+///
+/// This must be called before using the `print!` and `println!` macros.
+/// The only safe place to execute this function is in rust initialization code.
+///
+/// This must be called once with the bases of UARTs, mapped as device memory and (if necessary)
+/// shared with the host as MMIO, to which no other references must be held.
+pub unsafe fn init_all_uart(base_addresses: &[usize]) {
+    for (i, &base_address) in base_addresses.iter().enumerate() {
+        // Remember the valid address, for emergency console accesses.
+        ADDRESSES[i].call_once(|| base_address);
+
+        // Initialize the console driver, for normal console accesses.
+        assert!(!CONSOLES[i].is_completed(), "console::init() called more than once");
+        // SAFETY: The caller promised that base_address is the base of a mapped UART with no
+        // aliases.
+        CONSOLES[i].call_once(|| SpinMutex::new(unsafe { Uart::new(base_address) }));
+    }
+}
+
+/// Initialize console by mapping MMIO memory
+pub fn map_uarts_mmio() -> Result<(), hypervisor_backends::Error> {
+    if let Some(mmio_guard) = hypervisor_backends::get_mmio_guard() {
+        mmio_guard.enroll()?;
+
+        // TODO(ptosi): Use MmioSharer::share() to properly track this MMIO_GUARD_MAP.
+        //
+        // The following call shares the UART but also anything else present in 0..granule.
+        //
+        // For 4KiB, that's only the UARTs. For 16KiB, it also covers the RTC and watchdog but, as
+        // neither is used by vmbase clients (and as both are outside of the UART page), they
+        // will never have valid stage-1 mappings to those devices. As a result, this
+        // MMIO_GUARD_MAP isn't affected by the granule size in any visible way. Larger granule
+        // sizes will need to be checked separately, if needed.
+        assert!({
+            let granule = mmio_guard.granule()?;
+            granule == SIZE_4KB || granule == SIZE_16KB
+        });
+        // Validate the assumption above by ensuring that the UART is not moved to another page:
+        const_assert_eq!(UART_PAGE_ADDR, 0);
+        mmio_guard.map(UART_PAGE_ADDR)?;
+    }
+    Ok(())
+}
+
+/// Initialize platform specific device drivers. If this function fails the reboot is issued.
+pub fn init_console() {
+    if map_uarts_mmio().is_err() {
+        // UART mapping failed platform can't provide any output.
+        // Reboot to prevent printing any message.
+        reboot()
+    }
+    // SAFETY: UART_PAGE is mapped at stage-1 (see entry.S) and was just MMIO-guarded.
+    unsafe { init_all_uart(&UART_ADDRESSES) };
+}
+
+/// Return platform uart with specific index
+///
+/// Panics if console was not initialized by calling [`init`] first.
+pub fn uart(id: usize) -> &'static spin::mutex::SpinMutex<Uart> {
+    return CONSOLES[id].get().unwrap();
+}
+
+/// Reinitializes the emergency UART driver and returns it.
+///
+/// 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 emergency_uart() -> Uart {
+    // SAFETY: Initialization of UART using dedicated const address.
+    unsafe { Uart::new(UART_ADDRESSES[DEFAULT_EMERGENCY_CONSOLE_INDEX]) }
+}
 
 /// Makes a `PSCI_SYSTEM_OFF` call to shutdown the VM.
 ///
diff --git a/libs/libvmbase/src/console.rs b/libs/libvmbase/src/console.rs
index 533f2b6..6d9a4fe 100644
--- a/libs/libvmbase/src/console.rs
+++ b/libs/libvmbase/src/console.rs
@@ -14,67 +14,25 @@
 
 //! Console driver for 8250 UART.
 
-use crate::arch::uart::Uart;
+use crate::arch::platform;
 use core::fmt::{write, Arguments, Write};
-use spin::{mutex::SpinMutex, Once};
-
-// Arbitrary limit on the number of consoles that can be registered.
-//
-// Matches the UART count in crosvm.
-const MAX_CONSOLES: usize = 4;
-
-static CONSOLES: [Once<SpinMutex<Uart>>; MAX_CONSOLES] =
-    [Once::new(), Once::new(), Once::new(), Once::new()];
-static ADDRESSES: [Once<usize>; MAX_CONSOLES] =
-    [Once::new(), Once::new(), Once::new(), Once::new()];
-
-/// Index of the console used by default for logging.
-pub const DEFAULT_CONSOLE_INDEX: usize = 0;
-
-/// Index of the console used by default for emergency logging.
-pub const DEFAULT_EMERGENCY_CONSOLE_INDEX: usize = DEFAULT_CONSOLE_INDEX;
-
-/// Initialises the global instance(s) of the UART driver.
-///
-/// This must be called before using the `print!` and `println!` macros.
-///
-/// # Safety
-///
-/// This must be called once with the bases of UARTs, mapped as device memory and (if necessary)
-/// shared with the host as MMIO, to which no other references must be held.
-pub unsafe fn init(base_addresses: &[usize]) {
-    for (i, &base_address) in base_addresses.iter().enumerate() {
-        // Remember the valid address, for emergency console accesses.
-        ADDRESSES[i].call_once(|| base_address);
-
-        // Initialize the console driver, for normal console accesses.
-        assert!(!CONSOLES[i].is_completed(), "console::init() called more than once");
-        // SAFETY: The caller promised that base_address is the base of a mapped UART with no
-        // aliases.
-        CONSOLES[i].call_once(|| SpinMutex::new(unsafe { Uart::new(base_address) }));
-    }
-}
 
 /// Writes a formatted string followed by a newline to the n-th console.
 ///
 /// Panics if the n-th console was not initialized by calling [`init`] first.
 pub fn writeln(n: usize, format_args: Arguments) {
-    let uart = &mut *CONSOLES[n].get().unwrap().lock();
-
+    let uart = &mut *platform::uart(n).lock();
     write(uart, format_args).unwrap();
     let _ = uart.write_str("\n");
 }
 
-/// Reinitializes the n-th UART driver and writes a formatted string followed by a newline to it.
+/// Reinitializes the emergency UART driver and writes a formatted string followed by a newline to
+/// it.
 ///
 /// 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(n: usize, format_args: Arguments) {
-    let Some(addr) = ADDRESSES[n].get() else { return };
-
-    // SAFETY: addr contains the base of a mapped UART, passed in init().
-    let mut uart = unsafe { Uart::new(*addr) };
-
+pub fn ewriteln(format_args: Arguments) {
+    let mut uart = platform::emergency_uart();
     let _ = write(&mut uart, format_args);
     let _ = uart.write_str("\n");
 }
@@ -98,7 +56,7 @@
 /// use `eprintln!` instead.
 macro_rules! println {
     ($($arg:tt)*) => ({
-        $crate::console::console_writeln!($crate::console::DEFAULT_CONSOLE_INDEX, $($arg)*)
+        $crate::console::console_writeln!($crate::arch::platform::DEFAULT_CONSOLE_INDEX, $($arg)*)
     })
 }
 
@@ -111,6 +69,6 @@
 #[macro_export]
 macro_rules! eprintln {
     ($($arg:tt)*) => ({
-        $crate::console::ewriteln($crate::console::DEFAULT_EMERGENCY_CONSOLE_INDEX, format_args!($($arg)*))
+        $crate::console::ewriteln(format_args!($($arg)*))
     })
 }
diff --git a/libs/libvmbase/src/entry.rs b/libs/libvmbase/src/entry.rs
index 1bd38ca..5db0f25 100644
--- a/libs/libvmbase/src/entry.rs
+++ b/libs/libvmbase/src/entry.rs
@@ -15,53 +15,20 @@
 //! Rust entry point.
 
 use crate::{
-    arch::layout::{UART_ADDRESSES, UART_PAGE_ADDR},
-    bionic, console, heap, logger,
-    memory::{switch_to_dynamic_page_tables, PAGE_SIZE, SIZE_16KB, SIZE_4KB},
-    power::{reboot, shutdown},
+    arch::platform,
+    bionic, heap, logger,
+    memory::{switch_to_dynamic_page_tables, PAGE_SIZE, SIZE_4KB},
+    power::shutdown,
     rand,
 };
 use core::mem::size_of;
-use hypervisor_backends::{get_mmio_guard, Error};
-use static_assertions::const_assert_eq;
-
-fn try_console_init() -> Result<(), Error> {
-    if let Some(mmio_guard) = get_mmio_guard() {
-        mmio_guard.enroll()?;
-
-        // TODO(ptosi): Use MmioSharer::share() to properly track this MMIO_GUARD_MAP.
-        //
-        // The following call shares the UART but also anything else present in 0..granule.
-        //
-        // For 4KiB, that's only the UARTs. For 16KiB, it also covers the RTC and watchdog but, as
-        // neither is used by vmbase clients (and as both are outside of the UART page), they
-        // will never have valid stage-1 mappings to those devices. As a result, this
-        // MMIO_GUARD_MAP isn't affected by the granule size in any visible way. Larger granule
-        // sizes will need to be checked separately, if needed.
-        assert!({
-            let granule = mmio_guard.granule()?;
-            granule == SIZE_4KB || granule == SIZE_16KB
-        });
-        // Validate the assumption above by ensuring that the UART is not moved to another page:
-        const_assert_eq!(UART_PAGE_ADDR, 0);
-        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_ADDRESSES) };
-
-    Ok(())
-}
 
 /// 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) -> ! {
     heap::init();
-
-    if try_console_init().is_err() {
-        // Don't panic (or log) here to avoid accessing the console.
-        reboot()
-    }
+    // Initialize platform drivers
+    platform::init_console();
 
     logger::init().expect("Failed to initialize the logger");
     // We initialize the logger to Off (like the log crate) and clients should log::set_max_level.