feat(uefi): add post-boot-services allocator

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

Cargo.toml

@@ -8,6 +8,7 @@ edition = "2021"
 
 [dependencies]
 align-address = "0.3"
+allocator-api2 = { version = "0.2", default-features = false }
 anstyle = { version = "1", default-features = false }
 cfg-if = "1"
 hermit-entry = { version = "0.10", features = ["loader"] }
@@ -33,11 +34,8 @@ aarch64-cpu = "9"
 hermit-dtb = { version = "0.1" }
 goblin = { version = "0.8", default-features = false, features = ["elf64"] }
 
-[target.'cfg(target_os = "none")'.dependencies]
-allocator-api2 = { version = "0.2", default-features = false }
-
 [target.'cfg(target_os = "uefi")'.dependencies]
-uefi = { version = "0.28", features = ["alloc", "global_allocator", "panic_handler", "qemu"] }
+uefi = { version = "0.28", features = ["alloc", "panic_handler", "qemu"] }
 qemu-exit = "3"
 
 [target.'cfg(target_arch = "riscv64")'.dependencies]

src/main.rs

@@ -14,6 +14,7 @@ use hermit_entry::boot_info::{BootInfo, RawBootInfo};
 mod macros;
 
 mod arch;
+mod bump_allocator;
 mod log;
 mod os;
 

src/os/none/allocator/mod.rs

@@ -1,7 +1,6 @@
 //! Implementation of the Hermit Allocator in the loader
 
 mod bootstrap;
-mod bump;
 
 use core::ptr;
 use core::ptr::NonNull;
@@ -10,7 +9,7 @@ use allocator_api2::alloc::{AllocError, Allocator, GlobalAlloc, Layout};
 use one_shot_mutex::OneShotMutex;
 
 use self::bootstrap::BootstrapAllocator;
-use self::bump::BumpAllocator;
+use crate::bump_allocator::BumpAllocator;
 
 /// The global system allocator for Hermit.
 struct GlobalAllocator {

src/os/uefi/allocator.rs

@@ -0,0 +1,59 @@
+use alloc::vec;
+use core::alloc::{GlobalAlloc, Layout};
+use core::mem::MaybeUninit;
+use core::ptr::{self, NonNull};
+
+use allocator_api2::alloc::Allocator;
+use one_shot_mutex::OneShotMutex;
+
+use crate::bump_allocator::BumpAllocator;
+
+pub enum GlobalAllocator {
+	Uefi,
+	Bump(BumpAllocator),
+}
+
+pub struct LockedAllocator(OneShotMutex<GlobalAllocator>);
+
+impl LockedAllocator {
+	pub const fn uefi() -> Self {
+		Self(OneShotMutex::new(GlobalAllocator::Uefi))
+	}
+}
+
+unsafe impl GlobalAlloc for LockedAllocator {
+	unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+		match &*self.0.lock() {
+			GlobalAllocator::Uefi => unsafe { uefi::allocator::Allocator.alloc(layout) },
+			GlobalAllocator::Bump(bump) => bump
+				.allocate(layout)
+				// FIXME: Use NonNull::as_mut_ptr once `slice_ptr_get` is stabilized
+				// https://github.com/rust-lang/rust/issues/74265
+				.map_or(ptr::null_mut(), |ptr| ptr.as_ptr().cast()),
+		}
+	}
+
+	unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+		match &*self.0.lock() {
+			GlobalAllocator::Uefi => unsafe { uefi::allocator::Allocator.dealloc(ptr, layout) },
+			GlobalAllocator::Bump(bump) => unsafe {
+				bump.deallocate(NonNull::new(ptr).unwrap(), layout)
+			},
+		}
+	}
+}
+
+#[global_allocator]
+static ALLOCATOR: LockedAllocator = LockedAllocator::uefi();
+
+pub fn exit_boot_services() {
+	assert!(matches!(*ALLOCATOR.0.lock(), GlobalAllocator::Uefi));
+
+	let mem = vec![MaybeUninit::uninit(); 4096].leak();
+
+	let bump = BumpAllocator::from(mem);
+
+	*ALLOCATOR.0.lock() = GlobalAllocator::Bump(bump);
+
+	uefi::allocator::exit_boot_services();
+}

src/os/uefi/mod.rs

@@ -1,3 +1,4 @@
+mod allocator;
 mod console;
 
 use alloc::string::String;
@@ -7,20 +8,31 @@ use log::info;
 use qemu_exit::QEMUExit;
 use uefi::fs::{FileSystem, Path};
 use uefi::prelude::*;
+use uefi::table::boot::MemoryType;
 
 pub use self::console::CONSOLE;
 
 // Entry Point of the Uefi Loader
 #[entry]
 fn loader_main(_handle: Handle, mut system_table: SystemTable<Boot>) -> Status {
 	uefi::helpers::init(&mut system_table).unwrap();
+	unsafe {
+		uefi::allocator::init(&mut system_table);
+	}
 	crate::log::init();
 
 	let app = read_app(system_table.boot_services());
 
 	let string = String::from_utf8(app).unwrap();
 	println!("{string}");
 
+	allocator::exit_boot_services();
+	let (_runtime_system_table, _memory_map) =
+		system_table.exit_boot_services(MemoryType::LOADER_DATA);
+
+	println!("Exited boot services!");
+	println!("Allocations still {}!", String::from("work"));
+
 	let custom_exit_success = 3;
 	let qemu_exit_handle = qemu_exit::X86::new(0xf4, custom_exit_success);
 	qemu_exit_handle.exit_success()