Merge pull request #348 from hermit-os/uefi-alloc

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

.github/workflows/ci.yml

@@ -42,20 +42,15 @@ jobs:
     name: Run
     strategy:
       matrix:
-        target: [x86_64, x86_64-uefi, aarch64, riscv64]
+        target: [x86_64, aarch64, riscv64]
         os: [ubuntu-latest, macos-latest, windows-latest]
-        exclude:
-          - target: x86_64-uefi
-            os: macos-latest
-          - target: x86_64-uefi
-            os: windows-latest
     runs-on: ${{ matrix.os }}
     steps:
       - name: Install QEMU (ubuntu)
         if: matrix.os == 'ubuntu-latest'
         run: |
           sudo apt-get update
-          sudo apt-get install qemu-system-x86 qemu-system-arm qemu-system-misc ovmf
+          sudo apt-get install qemu-system-x86 qemu-system-arm qemu-system-misc
       - name: Install QEMU (macos)
         if: matrix.os == 'macos-latest'
         run: |
@@ -76,10 +71,29 @@ jobs:
         run: |
           gh release download v1.4 --repo riscv-software-src/opensbi --pattern 'opensbi-*-rv-bin.tar.xz'
           tar -xvf opensbi-*-rv-bin.tar.xz opensbi-1.4-rv-bin/share/opensbi/lp64/generic/firmware/fw_jump.bin
+      - name: Download OVMF (prebuilt)
+        run: |
+          gh release download edk2-stable202402-r1 --repo rust-osdev/ovmf-prebuilt --pattern 'edk2-stable*-bin.tar.xz'
+          tar -xvf edk2-stable*-bin.tar.xz edk2-stable202402-r1-bin/x64
+        if: matrix.os != 'ubuntu-latest'
+      - name: Download OVMF (ubuntu)
+        run: |
+          sudo apt-get update
+          sudo apt-get install ovmf
+          mkdir -p edk2-stable202402-r1-bin/x64
+          cp /usr/share/OVMF/OVMF_CODE.fd edk2-stable202402-r1-bin/x64/code.fd
+          cp /usr/share/OVMF/OVMF_VARS.fd edk2-stable202402-r1-bin/x64/vars.fd
+        if: matrix.os == 'ubuntu-latest'
       - name: Run VM (hello_world, dev)
         run: cargo xtask ci qemu --target ${{ matrix.target }}
       - name: Run VM (hello_world, release)
         run: cargo xtask ci qemu --target ${{ matrix.target }} --release
+      - name: Run VM (hello_world, uefi, dev)
+        run: cargo xtask ci qemu --target ${{ matrix.target }}-uefi
+        if: matrix.target == 'x86_64'
+      - name: Run VM (hello_world, uefi, release)
+        run: cargo xtask ci qemu --target ${{ matrix.target }}-uefi --release
+        if: matrix.target == 'x86_64'
       - name: Run VM (hello_world-microvm, dev)
         if: matrix.target == 'x86_64' && matrix.os == 'ubuntu-latest'
         run: cargo xtask ci qemu --target ${{ matrix.target }} --image hello_world-microvm --microvm
@@ -100,7 +114,7 @@ jobs:
       - name: Install QEMU
         run: |
           sudo apt-get update
-          sudo apt-get install -y --no-install-recommends qemu-system-x86 ovmf
+          sudo apt-get install -y --no-install-recommends qemu-system-x86
       - uses: actions/checkout@v4
         with:
             lfs: true
@@ -118,6 +132,10 @@ jobs:
           echo "$PWD" >> $GITHUB_PATH
   
           ./firecracker --version
+      - name: Download OVMF
+        run: |
+          gh release download edk2-stable202402-r1 --repo rust-osdev/ovmf-prebuilt --pattern 'edk2-stable*-bin.tar.xz'
+          tar -xvf edk2-stable*-bin.tar.xz edk2-stable202402-r1-bin/x64
       - uses: dtolnay/rust-toolchain@stable
       - name: Run QEMU (hello_world, dev)
         run: cargo xtask ci qemu --target x86_64 --accel

Cargo.toml

@@ -8,14 +8,15 @@ 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"] }
 log = "0.4"
 one-shot-mutex = "0.1"
 sptr = "0.3"
 take-static = "0.1"
 vm-fdt = { version = "0.3", default-features = false, features = ["alloc"] }
-anstyle = { version = "1", default-features = false }
 
 [features]
 default = []
@@ -33,13 +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 }
-exclusive_cell = "0.1"
-spinning_top = "0.3"
-
 [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/bootstrap.rs

@@ -33,7 +33,7 @@ use core::ops::Range;
 use core::ptr::NonNull;
 
 use allocator_api2::alloc::{AllocError, Allocator, Layout};
-use exclusive_cell::ExclusiveCell;
+use take_static::take_static;
 
 use self::ptr_range::PtrRange;
 
@@ -57,8 +57,10 @@ where
 		let mem = {
 			const SIZE: usize = 4 * 1024;
 			const BYTE: MaybeUninit<u8> = MaybeUninit::uninit();
-			/// The actual memory of the boostrap allocator.
-			static MEM: ExclusiveCell<[MaybeUninit<u8>; SIZE]> = ExclusiveCell::new([BYTE; SIZE]);
+			take_static! {
+				/// The actual memory of the boostrap allocator.
+				static MEM: [MaybeUninit<u8>; SIZE] = [BYTE; SIZE];
+			}
 			MEM.take().unwrap()
 		};
 

src/os/none/allocator/mod.rs

@@ -1,16 +1,15 @@
 //! Implementation of the Hermit Allocator in the loader
 
 mod bootstrap;
-mod bump;
 
 use core::ptr;
 use core::ptr::NonNull;
 
 use allocator_api2::alloc::{AllocError, Allocator, GlobalAlloc, Layout};
-use spinning_top::Spinlock;
+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 {
@@ -53,12 +52,12 @@ impl GlobalAllocator {
 	}
 }
 
-pub struct LockedAllocator(Spinlock<GlobalAllocator>);
+pub struct LockedAllocator(OneShotMutex<GlobalAllocator>);
 
 impl LockedAllocator {
 	/// Creates an empty allocator. All allocate calls will return `None`.
 	pub const fn empty() -> LockedAllocator {
-		LockedAllocator(Spinlock::new(GlobalAllocator::empty()))
+		LockedAllocator(OneShotMutex::new(GlobalAllocator::empty()))
 	}
 }
 

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()

xtask/src/ci/qemu.rs

@@ -42,11 +42,6 @@ impl Qemu {
 		let sh = crate::sh()?;
 
 		if self.build.target() == Target::X86_64Uefi {
-			if super::in_ci() {
-				sh.copy_file("/usr/share/OVMF/OVMF_CODE.fd", "OVMF_CODE.fd")?;
-				sh.copy_file("/usr/share/OVMF/OVMF_VARS.fd", "OVMF_VARS.fd")?;
-			}
-
 			sh.create_dir("target/esp/efi/boot")?;
 			sh.copy_file(self.build.dist_object(), "target/esp/efi/boot/bootx64.efi")?;
 			sh.write_file("target/esp/efi/boot/hermit-app", "Hello, UEFI!\n")?;
@@ -139,10 +134,10 @@ impl Qemu {
 					Target::X86_64Uefi => {
 						cpu_args.push("-drive".to_string());
 						cpu_args
-							.push("if=pflash,format=raw,readonly=on,file=OVMF_CODE.fd".to_string());
+							.push("if=pflash,format=raw,readonly=on,file=edk2-stable202402-r1-bin/x64/code.fd".to_string());
 						cpu_args.push("-drive".to_string());
 						cpu_args
-							.push("if=pflash,format=raw,readonly=on,file=OVMF_VARS.fd".to_string());
+							.push("if=pflash,format=raw,readonly=on,file=edk2-stable202402-r1-bin/x64/vars.fd".to_string());
 						cpu_args.push("-drive".to_string());
 						cpu_args.push("format=raw,file=fat:rw:target/esp".to_string());
 					}