Revise console implementation in EFI stub

ostd/libs/linux-bzimage/setup/src/console.rs

@@ -1,27 +1,23 @@
 // SPDX-License-Identifier: MPL-2.0
 
+//! A serial console.
+
 use core::fmt::{self, Write};
 
 use uart_16550::SerialPort;
 
+use crate::sync::Mutex;
+
 struct Stdout {
     serial_port: SerialPort,
 }
 
-static mut STDOUT: Stdout = Stdout {
-    serial_port: unsafe { SerialPort::new(0x0) },
-};
-
-/// SAFETY: this function must only be called once
-pub unsafe fn init() {
-    STDOUT = Stdout::init();
-}
-
 impl Stdout {
-    /// SAFETY: this function must only be called once
-    pub unsafe fn init() -> Self {
+    fn new() -> Self {
+        // FIXME: Is it safe to assume that the serial port always exists?
         let mut serial_port = unsafe { SerialPort::new(0x3F8) };
         serial_port.init();
+
         Self { serial_port }
     }
 }
@@ -33,76 +29,28 @@ impl Write for Stdout {
     }
 }
 
-/// Print a string to the console.
-///
-/// This is used when dyn Trait is not supported or fmt::Arguments is fragile to use in PIE.
-///
-/// # Safety
-///
-/// [`init()`] must be called before it and there should be no race conditions.
-pub unsafe fn print_str(s: &str) {
-    #[expect(static_mut_refs)]
-    STDOUT.write_str(s).unwrap();
-}
+static STDOUT: Mutex<Option<Stdout>> = Mutex::new(None);
 
-/// Print a single character to the console.
-///
-/// This is used when dyn Trait is not supported or fmt::Arguments is fragile to use in PIE.
-///
-/// # Safety
-///
-/// [`init()`] must be called before it and there should be no race conditions.
-unsafe fn print_char(c: char) {
-    #[expect(static_mut_refs)]
-    STDOUT.serial_port.send(c as u8);
-}
+/// Prints a format string and its arguments to the standard output.
+pub fn print_fmt(args: fmt::Arguments) {
+    let mut stdout = STDOUT.lock();
 
-/// Print a hexadecimal number to the console.
-///
-/// This is used when dyn Trait is not supported or fmt::Arguments is fragile to use in PIE.
-///
-/// # Safety
-///
-/// [`init()`] must be called before it and there should be no race conditions.
-pub unsafe fn print_hex(n: u64) {
-    print_str("0x");
-    for i in (0..16).rev() {
-        let digit = (n >> (i * 4)) & 0xf;
-        if digit < 10 {
-            print_char((b'0' + digit as u8) as char);
-        } else {
-            print_char((b'A' + (digit - 10) as u8) as char);
-        }
-    }
-}
-
-// TODO: Figure out why fmt::Arguments wont work even if relocations are applied.
-// We just settle on simple print functions for now.
-/*--------------------------------------------------------------------------------------------------
-
-/// Glue code for print!() and println!() macros.
-///
-/// SAFETY: init() must be called before print_fmt() and there should be no race conditions.
-pub unsafe fn print_fmt(args: fmt::Arguments) {
-    STDOUT.write_fmt(args).unwrap();
-}
-
-#[macro_export]
-macro_rules! print {
-    ($fmt: literal $(, $($arg: tt)+)?) => {
-        unsafe {
-            $crate::console::print_fmt(format_args!($fmt $(, $($arg)+)?))
-        }
+    // Fast path: The standard output has been initialized.
+    if let Some(inner) = stdout.as_mut() {
+        inner.write_fmt(args).unwrap();
+        return;
     }
+
+    // Initialize the standard output and print the string.
+    let mut inner = Stdout::new();
+    inner.write_fmt(args).unwrap();
+    *stdout = Some(inner);
 }
 
+/// Prints to the standard output, with a newline.
 #[macro_export]
 macro_rules! println {
-    ($fmt: literal $(, $($arg: tt)+)?) => {
-        unsafe {
-            $crate::console::print_fmt(format_args!(concat!($fmt, "\n") $(, $($arg)+)?))
-        }
+    ($fmt:literal $($arg:tt)*) => {
+        $crate::console::print_fmt(format_args!(concat!($fmt, "\n") $($arg)*))
     }
 }
-
- *------------------------------------------------------------------------------------------------*/

ostd/libs/linux-bzimage/setup/src/loader.rs

@@ -26,20 +26,12 @@ fn load_segment(file: &xmas_elf::ElfFile, program: &xmas_elf::program::ProgramHe
     let dst_slice = unsafe {
         core::slice::from_raw_parts_mut(program.physical_addr as *mut u8, program.mem_size as usize)
     };
-    /* crate::println!(
-        "[setup loader debug] loading ELF segment at {:#x}, size = {:#x}",
+    #[cfg(feature = "debug_print")]
+    crate::println!(
+        "[setup] Loading an ELF segment: addr={:#x}, size={:#x}",
         program.physical_addr,
         program.mem_size,
-    ); */
-    #[cfg(feature = "debug_print")]
-    unsafe {
-        use crate::console::{print_hex, print_str};
-        print_str("[setup loader debug] loading ELF segment at ");
-        print_hex(program.physical_addr as u64);
-        print_str(", size = ");
-        print_hex(program.mem_size as u64);
-        print_str("\n");
-    }
+    );
     // SAFETY: the ELF file is valid
     // dst_slice[..program.file_size as usize].copy_from_slice(header_data);
     unsafe {

ostd/libs/linux-bzimage/setup/src/main.rs

@@ -24,6 +24,7 @@
 
 mod console;
 mod loader;
+mod sync;
 
 // The entry points are defined in `x86/*/setup.S`.
 mod x86;

ostd/libs/linux-bzimage/setup/src/sync.rs

@@ -0,0 +1,27 @@
+// SPDX-License-Identifier: MPL-2.0
+
+//! Synchronization primitives.
+
+use core::cell::{RefCell, RefMut};
+
+/// A mutex.
+pub struct Mutex<T>(RefCell<T>);
+
+// SAFETY: We're single-threaded.
+unsafe impl<T: Send> Send for Mutex<T> {}
+unsafe impl<T: Sync> Sync for Mutex<T> {}
+
+/// A mutex guard.
+type MutexGuard<'a, T> = RefMut<'a, T>;
+
+impl<T> Mutex<T> {
+    /// Creates a new mutex.
+    pub const fn new(data: T) -> Self {
+        Self(RefCell::new(data))
+    }
+
+    /// Locks the mutex.
+    pub fn lock(&self) -> MutexGuard<T> {
+        self.0.borrow_mut()
+    }
+}

ostd/libs/linux-bzimage/setup/src/x86/amd64_efi/efi.rs

@@ -37,12 +37,6 @@ extern "sysv64" fn main_efi_handover64(
 /// This function should be called only once with valid parameters before all
 /// operations.
 unsafe fn system_init(handle: Handle, system_table: *const SystemTable) {
-    // SAFETY: This is the right time to initialize the console and it is only
-    // called once here before all console operations.
-    unsafe {
-        crate::console::init();
-    }
-
     // SAFETY: The handle and system_table are valid pointers. They are passed
     // from the UEFI firmware. They are only called once.
     unsafe {
@@ -53,8 +47,8 @@ unsafe fn system_init(handle: Handle, system_table: *const SystemTable) {
 
 fn efi_phase_boot(boot_params: &mut BootParams) -> ! {
     uefi::println!(
-        "[EFI stub] Stub loaded at {:#x?}",
-        crate::x86::image_load_offset()
+        "[EFI stub] Loaded with offset {:#x}",
+        crate::x86::image_load_offset(),
     );
 
     // Fill the boot params with the RSDP address if it is not provided.
@@ -65,10 +59,10 @@ fn efi_phase_boot(boot_params: &mut BootParams) -> ! {
     // Load the kernel payload to memory.
     let kernel = decode_payload(crate::x86::payload());
 
-    uefi::println!("[EFI stub] Loading payload.");
+    uefi::println!("[EFI stub] Loading the payload as an ELF file");
     crate::loader::load_elf(&kernel);
 
-    uefi::println!("[EFI stub] Exiting EFI boot services.");
+    uefi::println!("[EFI stub] Exiting EFI boot services");
     let memory_type = {
         let Ok(loaded_image) = open_protocol_exclusive::<LoadedImage>(boot::image_handle()) else {
             panic!("Failed to open LoadedImage protocol");
@@ -83,27 +77,21 @@ fn efi_phase_boot(boot_params: &mut BootParams) -> ! {
 }
 
 fn efi_phase_runtime(memory_map: MemoryMapOwned, boot_params: &mut BootParams) -> ! {
-    unsafe {
-        crate::console::print_str("[EFI stub] Entered runtime services.\n");
-    }
-
+    crate::println!(
+        "[EFI stub] Processing {} memory map entries",
+        memory_map.entries().len()
+    );
     #[cfg(feature = "debug_print")]
-    unsafe {
-        use crate::console::{print_hex, print_str};
-        print_str("[EFI stub debug] EFI Memory map:\n");
-        for md in memory_map.entries() {
-            // crate::println!("    [{:#x}] {:#x} ({:#x})", md.ty.0, md.phys_start, md.page_count);
-            print_str("    [");
-            print_hex(md.ty.0 as u64);
-            print_str("]");
-            print_hex(md.phys_start);
-            print_str("(size=");
-            print_hex(md.page_count);
-            print_str(")");
-            print_str("{flags=");
-            print_hex(md.att.bits());
-            print_str("}\n");
-        }
+    {
+        memory_map.entries().for_each(|entry| {
+            crate::println!(
+                "    [{:#x}] {:#x} (size={:#x}) {{flags={:#x}}}",
+                entry.ty.0,
+                entry.phys_start,
+                entry.page_count,
+                entry.att.bits()
+            );
+        })
     }
 
     // Write memory map to e820 table in boot_params.
@@ -111,11 +99,7 @@ fn efi_phase_runtime(memory_map: MemoryMapOwned, boot_params: &mut BootParams) -
     let mut e820_entries = 0usize;
     for md in memory_map.entries() {
         if e820_entries >= e820_table.len() || e820_entries >= 127 {
-            unsafe {
-                crate::console::print_str(
-                    "[EFI stub] Warning: number of E820 entries exceeded 128!\n",
-                );
-            }
+            crate::println!("[EFI stub] Warning: The number of E820 entries exceeded 128!");
             break;
         }
         e820_table[e820_entries] = linux_boot_params::BootE820Entry {
@@ -131,7 +115,7 @@ fn efi_phase_runtime(memory_map: MemoryMapOwned, boot_params: &mut BootParams) -
                 #[cfg(feature = "cvm_guest")]
                 uefi::table::boot::MemoryType::UNACCEPTED => {
                     unsafe {
-                        crate::console::print_str("[EFI stub] Accepting pending pages...\n");
+                        crate::println!("[EFI stub] Accepting pending pages");
                         for page_idx in 0..md.page_count {
                             tdx_guest::tdcall::accept_page(0, md.phys_start + page_idx * PAGE_SIZE)
                                 .unwrap();
@@ -146,13 +130,10 @@ fn efi_phase_runtime(memory_map: MemoryMapOwned, boot_params: &mut BootParams) -
     }
     boot_params.e820_entries = e820_entries as u8;
 
-    unsafe {
-        use crate::console::{print_hex, print_str};
-        print_str("[EFI stub] Entering Asterinas entrypoint at ");
-        print_hex(super::ASTER_ENTRY_POINT as u64);
-        print_str("\n");
-    }
-
+    crate::println!(
+        "[EFI stub] Entering the Asterinas entry point at {:#x}",
+        super::ASTER_ENTRY_POINT,
+    );
     unsafe {
         super::call_aster_entrypoint(
             super::ASTER_ENTRY_POINT as u64,

ostd/libs/linux-bzimage/setup/src/x86/legacy_i386/mod.rs

@@ -4,24 +4,22 @@ use core::arch::{asm, global_asm};
 
 global_asm!(include_str!("setup.S"));
 
-use crate::console::{print_hex, print_str};
-
 pub const ASTER_ENTRY_POINT: u32 = 0x8001000;
 
 #[export_name = "main_legacy32"]
 extern "cdecl" fn main_legacy32(boot_params_ptr: u32) -> ! {
-    // SAFETY: this init function is only called once.
-    unsafe { crate::console::init() };
-
-    // println!("[setup] bzImage loaded at {:#x}", x86::relocation::image_load_offset());
-    unsafe {
-        print_str("[setup] bzImage loaded offset: ");
-        print_hex(crate::x86::image_load_offset() as u64);
-        print_str("\n");
-    }
+    crate::println!(
+        "[setup] Loaded with offset {:#x}",
+        crate::x86::image_load_offset(),
+    );
 
+    crate::println!("[setup] Loading the payload as an ELF file");
     crate::loader::load_elf(crate::x86::payload());
 
+    crate::println!(
+        "[setup] Entering the Asterinas entry point at {:#x}",
+        ASTER_ENTRY_POINT,
+    );
     // SAFETY: the entrypoint and the ptr is valid.
     unsafe { call_aster_entrypoint(ASTER_ENTRY_POINT, boot_params_ptr.try_into().unwrap()) };
 }
@@ -35,6 +33,11 @@ unsafe fn call_aster_entrypoint(entrypoint: u32, boot_params_ptr: u32) -> ! {
 }
 
 #[panic_handler]
-fn panic(_info: &core::panic::PanicInfo) -> ! {
-    loop {}
+fn panic(info: &core::panic::PanicInfo) -> ! {
+    crate::println!("[PANIC]: {}", info);
+
+    loop {
+        // SAFETY: `hlt` has no effect other than to stop the CPU and wait for another interrupt.
+        unsafe { asm!("hlt", options(nomem, nostack, preserves_flags)) };
+    }
 }