From c3b0e2624a0c7ad5987712c01eaf79b7ada967e4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Martin=20Kr=C3=B6ning?=
 <martin.kroening@eonerc.rwth-aachen.de>
Date: Tue, 9 Apr 2024 14:52:48 +0200
Subject: [PATCH] feat(uefi): allow printing after exiting boot services
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Martin Kröning <martin.kroening@eonerc.rwth-aachen.de>
---
 src/arch/mod.rs        |  2 +-
 src/arch/x86_64/mod.rs |  6 +++++
 src/os/uefi/console.rs | 58 +++++++++++++++++++++++++++++++++++++++---
 3 files changed, 62 insertions(+), 4 deletions(-)

diff --git a/src/arch/mod.rs b/src/arch/mod.rs
index 67f9ebed..7bf01322 100644
--- a/src/arch/mod.rs
+++ b/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::*;
 	}
diff --git a/src/arch/x86_64/mod.rs b/src/arch/x86_64/mod.rs
index dad04b73..4d891459 100644
--- a/src/arch/x86_64/mod.rs
+++ b/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};
diff --git a/src/os/uefi/console.rs b/src/os/uefi/console.rs
index cf159fa6..36ddf26e 100644
--- a/src/os/uefi/console.rs
+++ b/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());