Avoid segment overlapping in EFI stub

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

@@ -1,5 +1,7 @@
 // SPDX-License-Identifier: MPL-2.0
 
+use core::mem::MaybeUninit;
+
 use xmas_elf::program::{ProgramHeader, SegmentData};
 
 /// Load the kernel ELF payload to memory.
@@ -24,12 +26,7 @@ fn load_segment(file: &xmas_elf::ElfFile, program: &xmas_elf::program::ProgramHe
         panic!("[setup] Unexpected segment data type!");
     };
 
-    // FIXME: This can be unsafe if the memory region overlaps with allocated memory or memory
-    // reserved by the firmware. We need to query UEFI or check the memory map ourselves to prevent
-    // this from happening.
-    let dst_slice = unsafe {
-        core::slice::from_raw_parts_mut(program.physical_addr as *mut u8, program.mem_size as usize)
-    };
+    let dst_slice = crate::x86::alloc_at(program.physical_addr as usize, program.mem_size as usize);
 
     #[cfg(feature = "debug_print")]
     crate::println!(
@@ -39,6 +36,6 @@ fn load_segment(file: &xmas_elf::ElfFile, program: &xmas_elf::program::ProgramHe
     );
 
     let (left, right) = dst_slice.split_at_mut(program.file_size as usize);
-    left.copy_from_slice(segment_data);
-    right.fill(0);
+    left.write_copy_of_slice(segment_data);
+    MaybeUninit::fill(right, 0);
 }

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

@@ -21,6 +21,8 @@
 
 #![no_std]
 #![no_main]
+#![feature(maybe_uninit_fill)]
+#![feature(maybe_uninit_write_slice)]
 
 mod console;
 mod loader;

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

@@ -0,0 +1,34 @@
+// SPDX-License-Identifier: MPL-2.0
+
+use core::mem::MaybeUninit;
+
+pub fn alloc_at(addr: usize, size: usize) -> &'static mut [MaybeUninit<u8>] {
+    assert_ne!(addr, 0, "the address to allocate is zero");
+    assert!(
+        size <= isize::MAX as usize,
+        "the size to allocate exceeds `isize::MAX`"
+    );
+
+    addr.checked_add(size)
+        .expect("the range to allocate overflows");
+
+    let allocated = uefi::boot::allocate_pages(
+        uefi::boot::AllocateType::Address(addr as u64),
+        uefi::boot::MemoryType::LOADER_DATA,
+        size.div_ceil(super::efi::PAGE_SIZE as usize),
+    )
+    .expect("the UEFI allocation fails");
+    assert_eq!(
+        allocated.as_ptr() as usize,
+        addr,
+        "the allocated address is not the request address"
+    );
+
+    // SAFETY:
+    // 1. The address is not zero and the size is reasonable (there are less the `isize::MAX` bytes
+    //    and the range won't overflow the address space), as asserted above.
+    // 2. The memory region is allocated via the UEFI firmware, so it is valid for reading and
+    //    writing. We will not deallocate it, so it live for `'static`.
+    // 3. The type alignment is 1 and the type can contain uninitialized data.
+    unsafe { core::slice::from_raw_parts_mut(addr as *mut MaybeUninit<u8>, size) }
+}

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

@@ -6,7 +6,7 @@ use uefi_raw::table::system::SystemTable;
 
 use super::decoder::decode_payload;
 
-const PAGE_SIZE: u64 = 4096;
+pub(super) const PAGE_SIZE: u64 = 4096;
 
 #[export_name = "main_efi_handover64"]
 extern "sysv64" fn main_efi_handover64(

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

@@ -1,5 +1,6 @@
 // SPDX-License-Identifier: MPL-2.0
 
+pub(super) mod alloc;
 mod decoder;
 mod efi;
 

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

@@ -0,0 +1,89 @@
+// SPDX-License-Identifier: MPL-2.0
+
+use core::{mem::MaybeUninit, ops::Range};
+
+use crate::sync::Mutex;
+
+const NUM_USED_RANGES: usize = 16;
+
+struct State {
+    e820: &'static [linux_boot_params::BootE820Entry],
+    used: [Range<usize>; NUM_USED_RANGES],
+}
+
+static STATE: Mutex<Option<State>> = Mutex::new(None);
+
+/// # Safety
+///
+/// The caller must ensure that the E820 entries in the boot parameters correctly represent the
+/// current memory map.
+pub(super) unsafe fn init(boot_params: &'static linux_boot_params::BootParams) {
+    let mut state = STATE.lock();
+
+    assert!(state.is_none());
+
+    extern "C" {
+        fn __executable_start();
+        fn __executable_end();
+    }
+
+    let mut used = core::array::from_fn(|_| 0..0);
+    used[0] = (__executable_start as usize)..(__executable_end as usize);
+    used[1] = {
+        let params_addr = core::ptr::from_ref(boot_params).addr();
+        params_addr..(params_addr + core::mem::size_of::<linux_boot_params::BootParams>())
+    };
+
+    *state = Some(State {
+        e820: &boot_params.e820_table[..(boot_params.e820_entries as usize)],
+        used,
+    });
+}
+
+pub fn alloc_at(addr: usize, size: usize) -> &'static mut [MaybeUninit<u8>] {
+    let mut state = STATE.lock();
+    let state = state.as_mut().unwrap();
+
+    assert_ne!(addr, 0, "the address to allocate is zero");
+    assert!(
+        size <= isize::MAX as usize,
+        "the size to allocate exceeds `isize::MAX`"
+    );
+
+    let range = addr..addr
+        .checked_add(size)
+        .expect("the range to allocate overflows");
+
+    assert!(
+        state.e820.iter().any(|entry| {
+            let typ = entry.typ;
+            typ == linux_boot_params::E820Type::Ram
+                && entry.addr as usize <= range.start
+                && range.end <= (entry.addr + entry.size) as usize
+        }),
+        "the range to allocate is not usable"
+    );
+
+    let overlapped = state
+        .used
+        .iter()
+        .find(|used| used.start < range.end && range.start < used.end);
+    assert!(overlapped.is_none(), "the range to allocate is used");
+
+    let empty = state
+        .used
+        .iter_mut()
+        // Use fully qualified syntax to avoid collisions with the unstable method
+        // `ExactSizeIterator::is_empty`. See <https://github.com/rust-lang/rust/issues/86682>.
+        .find(|used| Range::is_empty(used))
+        .expect("the allocated ranges are full");
+    *empty = range;
+
+    // SAFETY:
+    // 1. The address is not zero and the size is reasonable (there are less the `isize::MAX` bytes
+    //    and the range won't overflow the address space), as asserted above.
+    // 2. The memory region is usable and just allocated, so it is valid for reading and writing.
+    //    We will not deallocate it, so it live for `'static`.
+    // 3. The type alignment is 1 and the type can contain uninitialized data.
+    unsafe { core::slice::from_raw_parts_mut(addr as *mut MaybeUninit<u8>, size) }
+}

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

@@ -4,6 +4,8 @@ use core::arch::{asm, global_asm};
 
 use linux_boot_params::BootParams;
 
+pub(super) mod alloc;
+
 global_asm!(include_str!("setup.S"));
 
 const ASTER_ENTRY_POINT: *const () = 0x8001000 as _;
@@ -15,6 +17,13 @@ extern "cdecl" fn main_legacy32(boot_params_ptr: *mut BootParams) -> ! {
         crate::x86::image_load_offset(),
     );
 
+    // SAFETY: We get boot parameters from the boot loader, so by contract the pointer is valid and
+    // the underlying memory is initialized. We never mutate the boot parameters, so we can create
+    // an immutable reference of the plain-old-data type.
+    let boot_params = unsafe { &*boot_params_ptr };
+    // SAFETY: We get boot parameters from the boot loader. By contract they are correct.
+    unsafe { alloc::init(boot_params) };
+
     crate::println!("[setup] Loading the payload as an ELF file");
     crate::loader::load_elf(crate::x86::payload());
 

ostd/libs/linux-bzimage/setup/src/x86/legacy_i386/setup.S

@@ -24,6 +24,19 @@ load_address:
     // Set up the stack.
     mov esp, offset __stack_top
 
+    // Load the GDT.
+    push 8        // 32-bit code
+    mov eax, offset gdt_loaded
+    push eax
+    lgdt [gdtr]
+    retf
+gdt_loaded:
+    mov ax, 16    // 32-bit data
+    mov ds, ax
+    mov es, ax
+    mov fs, ax
+    mov gs, ax
+
     // Call the Rust main routine.
     push esi
     call main_legacy32
@@ -39,6 +52,20 @@ halt:
     hlt
     jmp halt
 
+// GDT. We define GDT ourselves to ensure that the GDT page will not be
+// accidentally overwritten by the allocated memory.
+.rodata
+.align 16
+gdtr:
+    .word gdt_end - gdt - 1
+    .long gdt
+.align 16
+gdt:
+    .quad 0x0000000000000000 // 0:  null descriptor
+    .quad 0x00cf9a000000ffff // 8:  32-bit code segment
+    .quad 0x00cf92000000ffff // 16: 32-bit data segment
+gdt_end:
+
 // A small stack for the setup code.
 .bss
 .align 8

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

@@ -6,10 +6,14 @@ cfg_if::cfg_if! {
     if #[cfg(target_arch = "x86_64")] {
         mod amd64_efi;
 
+        pub use amd64_efi::alloc::alloc_at;
+
         const CFG_TARGET_ARCH_X86_64: usize = 1;
     } else if #[cfg(target_arch = "x86")] {
         mod legacy_i386;
 
+        pub use legacy_i386::alloc::alloc_at;
+
         const CFG_TARGET_ARCH_X86_64: usize = 0;
     } else {
         compile_error!("unsupported target architecture");