feat(uefi): allow printing after exiting boot services

Signed-off-by: Martin Kröning <[email protected]>

src/arch/mod.rs

@@ -5,7 +5,7 @@ cfg_if::cfg_if! {
 	} else if #[cfg(target_arch = "riscv64")] {
 		mod riscv64;
 		pub use self::riscv64::*;
-	} else if #[cfg(all(target_arch = "x86_64", target_os = "none"))] {
+	} else if #[cfg(all(target_arch = "x86_64"))] {
 		mod x86_64;
 		pub use self::x86_64::*;
 	}

src/arch/x86_64/mod.rs

@@ -9,14 +9,19 @@ cfg_if::cfg_if! {
 }
 
 mod console;
+#[cfg(target_os = "none")]
 mod paging;
+#[cfg(target_os = "none")]
 mod physicalmem;
 
 pub use console::Console;
 
+#[cfg(target_os = "none")]
 const KERNEL_STACK_SIZE: u64 = 32_768;
+#[cfg(target_os = "none")]
 const SERIAL_IO_PORT: u16 = 0x3F8;
 
+#[cfg(target_os = "none")]
 unsafe fn map_memory(address: usize, memory_size: usize) -> usize {
 	use align_address::Align;
 	use x86_64::structures::paging::{PageSize, PageTableFlags, Size2MiB};
@@ -29,6 +34,7 @@ unsafe fn map_memory(address: usize, memory_size: usize) -> usize {
 	address
 }
 
+#[cfg(target_os = "none")]
 pub unsafe fn get_memory(memory_size: u64) -> u64 {
 	use align_address::Align;
 	use x86_64::structures::paging::{PageSize, Size2MiB};

src/os/uefi/console.rs

@@ -1,14 +1,66 @@
+use core::ffi::c_void;
 use core::fmt;
+use core::ptr::NonNull;
+use core::sync::atomic::{AtomicBool, Ordering};
 
 use one_shot_mutex::OneShotMutex;
+use uefi::table::boot::{EventType, Tpl};
+use uefi::table::{Boot, SystemTable};
+use uefi::Event;
 
-pub struct Console(());
+use crate::arch;
+
+pub enum Console {
+	None,
+	BootServices,
+	Native { console: arch::Console },
+}
+
+impl Console {
+	const fn new() -> Self {
+		Self::None
+	}
+
+	fn exit_boot_services(&mut self) {
+		assert!(matches!(self, Self::BootServices { .. }));
+		*self = Self::Native {
+			console: arch::Console::default(),
+		};
+	}
+
+	fn init(&mut self) {
+		assert!(matches!(self, Console::None));
+		unsafe {
+			uefi_services::system_table()
+				.boot_services()
+				.create_event(
+					EventType::SIGNAL_EXIT_BOOT_SERVICES,
+					Tpl::NOTIFY,
+					Some(exit_boot_services),
+					None,
+				)
+				.unwrap();
+		}
+		*self = Console::BootServices;
+	}
+}
 
 impl fmt::Write for Console {
 	fn write_str(&mut self, s: &str) -> fmt::Result {
-		uefi_services::system_table().stdout().write_str(s)?;
+		match self {
+			Console::None => {
+				self.init();
+				self.write_str(s)?;
+			}
+			Console::BootServices => uefi_services::system_table().stdout().write_str(s)?,
+			Console::Native { console } => console.write_bytes(s.as_bytes()),
+		}
 		Ok(())
 	}
 }
 
-pub static CONSOLE: OneShotMutex<Console> = OneShotMutex::new(Console(()));
+unsafe extern "efiapi" fn exit_boot_services(_e: Event, _ctx: Option<NonNull<c_void>>) {
+	CONSOLE.lock().exit_boot_services();
+}
+
+pub static CONSOLE: OneShotMutex<Console> = OneShotMutex::new(Console::new());