Clean up Rust code in EFI stub

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

@@ -12,6 +12,7 @@ pub fn load_elf(file: &[u8]) {
                 "[setup] Unexpected program header type! Asterinas should be 64-bit ELF binary."
             );
         };
+
         if program.get_type().unwrap() == xmas_elf::program::Type::Load {
             load_segment(&elf, program);
         }
@@ -19,41 +20,25 @@ pub fn load_elf(file: &[u8]) {
 }
 
 fn load_segment(file: &xmas_elf::ElfFile, program: &xmas_elf::program::ProgramHeader64) {
-    let SegmentData::Undefined(header_data) = program.get_data(file).unwrap() else {
+    let SegmentData::Undefined(segment_data) = program.get_data(file).unwrap() else {
         panic!("[setup] Unexpected segment data type!");
     };
-    // SAFETY: the physical address from the ELF file is valid
+
+    // 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)
     };
+
     #[cfg(feature = "debug_print")]
     crate::println!(
         "[setup] Loading an ELF segment: addr={:#x}, size={:#x}",
         program.physical_addr,
         program.mem_size,
     );
-    // SAFETY: the ELF file is valid
-    // dst_slice[..program.file_size as usize].copy_from_slice(header_data);
-    unsafe {
-        memcpy(
-            dst_slice.as_mut_ptr(),
-            header_data.as_ptr(),
-            program.file_size as usize,
-        );
-    }
-    let zero_slice = &mut dst_slice[program.file_size as usize..];
-    zero_slice.fill(0);
-}
 
-/// TODO: remove this and use copy_from_slice instead
-///
-/// We use a custom memcpy because the standard library's compiler's builtin memcpy
-/// fails for some unknown reason. Sometimes that will result in "Unknown OPCode"
-/// machine error.
-unsafe fn memcpy(dst: *mut u8, src: *const u8, size: usize) {
-    let mut i = 0;
-    while i < size {
-        *dst.add(i) = *src.add(i);
-        i += 1;
-    }
+    let (left, right) = dst_slice.split_at_mut(program.file_size as usize);
+    left.copy_from_slice(segment_data);
+    right.fill(0);
 }

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

@@ -4,7 +4,7 @@
 
 extern crate alloc;
 
-pub use alloc::vec::Vec;
+use alloc::vec::Vec;
 use core::convert::TryFrom;
 
 use core2::io::Read;
@@ -16,20 +16,23 @@ enum MagicNumber {
     Zlib,
 }
 
+#[derive(Debug)]
+struct InvalidMagicNumber;
+
 impl TryFrom<&[u8]> for MagicNumber {
-    type Error = &'static str;
+    type Error = InvalidMagicNumber;
 
     fn try_from(slice: &[u8]) -> Result<Self, Self::Error> {
         match *slice {
             [0x7F, 0x45, 0x4C, 0x46, ..] => Ok(Self::Elf),
             [0x1F, 0x8B, ..] => Ok(Self::Gzip),
             [0x78, 0x9C, ..] => Ok(Self::Zlib),
-            _ => Err("Unsupported payload type"),
+            _ => Err(InvalidMagicNumber),
         }
     }
 }
 
-/// Checking the encoding format and matching decoding methods to decode payload.
+/// Detects the format used to encode the payload and decodes the payload accordingly.
 pub fn decode_payload(payload: &[u8]) -> Vec<u8> {
     let mut kernel = Vec::new();
     let magic = MagicNumber::try_from(payload).unwrap();

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

@@ -1,12 +1,7 @@
 // SPDX-License-Identifier: MPL-2.0
 
 use linux_boot_params::BootParams;
-use uefi::{
-    boot::{exit_boot_services, open_protocol_exclusive},
-    mem::memory_map::{MemoryMap, MemoryMapOwned},
-    prelude::*,
-    proto::loaded_image::LoadedImage,
-};
+use uefi::{boot::exit_boot_services, mem::memory_map::MemoryMap, prelude::*};
 use uefi_raw::table::system::SystemTable;
 
 use super::decoder::decode_payload;
@@ -19,33 +14,28 @@ extern "sysv64" fn main_efi_handover64(
     system_table: *const SystemTable,
     boot_params_ptr: *mut BootParams,
 ) -> ! {
-    // SAFETY: handle and system_table are valid pointers. It is only called once.
-    unsafe { system_init(handle, system_table) };
-
-    uefi::helpers::init().unwrap();
-
-    // SAFETY: boot_params is a valid pointer.
-    let boot_params = unsafe { &mut *boot_params_ptr };
-
-    efi_phase_boot(boot_params)
-}
-
-/// Initialize the system.
-///
-/// # Safety
-///
-/// This function should be called only once with valid parameters before all
-/// operations.
-unsafe fn system_init(handle: Handle, system_table: *const SystemTable) {
-    // SAFETY: The handle and system_table are valid pointers. They are passed
-    // from the UEFI firmware. They are only called once.
+    // SAFETY: We get `handle` and `system_table` from the UEFI firmware, so by contract the
+    // pointers are valid and correct.
     unsafe {
         boot::set_image_handle(handle);
         uefi::table::set_system_table(system_table);
     }
+
+    uefi::helpers::init().unwrap();
+
+    // SAFETY: We get boot parameters from the boot loader, so by contract the pointer is valid and
+    // the underlying memory is initialized. We are an exclusive owner of the memory region, so we
+    // can create a mutable reference of the plain-old-data type.
+    let boot_params = unsafe { &mut *boot_params_ptr };
+
+    efi_phase_boot(boot_params);
+
+    // SAFETY: We do not open boot service protocols or maintain references to boot service code
+    // and data.
+    unsafe { efi_phase_runtime(boot_params) };
 }
 
-fn efi_phase_boot(boot_params: &mut BootParams) -> ! {
+fn efi_phase_boot(boot_params: &mut BootParams) {
     uefi::println!(
         "[EFI stub] Loaded with offset {:#x}",
         crate::x86::image_load_offset(),
@@ -53,30 +43,40 @@ fn efi_phase_boot(boot_params: &mut BootParams) -> ! {
 
     // Fill the boot params with the RSDP address if it is not provided.
     if boot_params.acpi_rsdp_addr == 0 {
-        boot_params.acpi_rsdp_addr = get_rsdp_addr();
+        boot_params.acpi_rsdp_addr =
+            find_rsdp_addr().expect("ACPI RSDP address is not available") as usize as u64;
     }
 
-    // Load the kernel payload to memory.
+    // Decode the payload and load it as an ELF file.
+    uefi::println!("[EFI stub] Decoding the kernel payload");
     let kernel = decode_payload(crate::x86::payload());
-
     uefi::println!("[EFI stub] Loading the payload as an ELF file");
     crate::loader::load_elf(&kernel);
-
-    uefi::println!("[EFI stub] Exiting EFI boot services");
-    let memory_type = {
-        let Ok(loaded_image) = open_protocol_exclusive::<LoadedImage>(boot::image_handle()) else {
-            panic!("Failed to open LoadedImage protocol");
-        };
-        loaded_image.data_type()
-    };
-    // SAFETY: All allocations in the boot services phase are not used after
-    // this point.
-    let memory_map = unsafe { exit_boot_services(memory_type) };
-
-    efi_phase_runtime(memory_map, boot_params);
 }
 
-fn efi_phase_runtime(memory_map: MemoryMapOwned, boot_params: &mut BootParams) -> ! {
+fn find_rsdp_addr() -> Option<*const ()> {
+    use uefi::table::cfg::{ACPI2_GUID, ACPI_GUID};
+
+    // Prefer ACPI2 over ACPI.
+    for acpi_guid in [ACPI2_GUID, ACPI_GUID] {
+        if let Some(rsdp_addr) = uefi::system::with_config_table(|table| {
+            table
+                .iter()
+                .find(|entry| entry.guid == acpi_guid)
+                .map(|entry| entry.address.cast::<()>())
+        }) {
+            return Some(rsdp_addr);
+        }
+    }
+
+    None
+}
+
+unsafe fn efi_phase_runtime(boot_params: &mut BootParams) -> ! {
+    uefi::println!("[EFI stub] Exiting EFI boot services");
+    // SAFETY: The safety is upheld by the caller.
+    let memory_map = unsafe { exit_boot_services(uefi::table::boot::MemoryType::LOADER_DATA) };
+
     crate::println!(
         "[EFI stub] Processing {} memory map entries",
         memory_map.entries().len()
@@ -94,7 +94,7 @@ fn efi_phase_runtime(memory_map: MemoryMapOwned, boot_params: &mut BootParams) -
         })
     }
 
-    // Write memory map to e820 table in boot_params.
+    // Write the memory map to the E820 table in `boot_params`.
     let e820_table = &mut boot_params.e820_table;
     let mut e820_entries = 0usize;
     for md in memory_map.entries() {
@@ -131,29 +131,11 @@ fn efi_phase_runtime(memory_map: MemoryMapOwned, boot_params: &mut BootParams) -
     boot_params.e820_entries = e820_entries as u8;
 
     crate::println!(
-        "[EFI stub] Entering the Asterinas entry point at {:#x}",
+        "[EFI stub] Entering the Asterinas entry point at {:p}",
         super::ASTER_ENTRY_POINT,
     );
-    unsafe {
-        super::call_aster_entrypoint(
-            super::ASTER_ENTRY_POINT as u64,
-            boot_params as *const _ as u64,
-        )
-    }
-}
-
-fn get_rsdp_addr() -> u64 {
-    use uefi::table::cfg::{ACPI2_GUID, ACPI_GUID};
-    uefi::system::with_config_table(|table| {
-        for entry in table {
-            // Prefer ACPI2 over ACPI.
-            if entry.guid == ACPI2_GUID {
-                return entry.address as usize as u64;
-            }
-            if entry.guid == ACPI_GUID {
-                return entry.address as usize as u64;
-            }
-        }
-        panic!("ACPI RSDP not found");
-    })
+    // SAFETY:
+    // 1. The entry point address is correct and matches the kernel ELF file.
+    // 2. The boot parameter pointer is valid and points to the correct boot parameters.
+    unsafe { super::call_aster_entrypoint(super::ASTER_ENTRY_POINT, boot_params) }
 }

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

@@ -5,14 +5,21 @@ mod efi;
 
 use core::arch::{asm, global_asm};
 
+use linux_boot_params::BootParams;
+
 global_asm!(include_str!("setup.S"));
 
-pub const ASTER_ENTRY_POINT: u32 = 0x8001200;
+const ASTER_ENTRY_POINT: *const () = 0x8001200 as _;
 
-unsafe fn call_aster_entrypoint(entrypoint: u64, boot_params_ptr: u64) -> ! {
-    asm!("mov rsi, {}", in(reg) boot_params_ptr);
-    asm!("mov rax, {}", in(reg) entrypoint);
-    asm!("jmp rax");
-
-    unreachable!();
+unsafe fn call_aster_entrypoint(entrypoint: *const (), boot_params_ptr: *mut BootParams) -> ! {
+    unsafe {
+        asm!(
+            "mov rsi, {1}",
+            "mov rax, {0}",
+            "jmp rax",
+            in(reg) entrypoint,
+            in(reg) boot_params_ptr,
+            options(noreturn),
+        )
+    }
 }

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

@@ -2,12 +2,14 @@
 
 use core::arch::{asm, global_asm};
 
+use linux_boot_params::BootParams;
+
 global_asm!(include_str!("setup.S"));
 
-pub const ASTER_ENTRY_POINT: u32 = 0x8001000;
+const ASTER_ENTRY_POINT: *const () = 0x8001000 as _;
 
 #[export_name = "main_legacy32"]
-extern "cdecl" fn main_legacy32(boot_params_ptr: u32) -> ! {
+extern "cdecl" fn main_legacy32(boot_params_ptr: *mut BootParams) -> ! {
     crate::println!(
         "[setup] Loaded with offset {:#x}",
         crate::x86::image_load_offset(),
@@ -17,19 +19,26 @@ extern "cdecl" fn main_legacy32(boot_params_ptr: u32) -> ! {
     crate::loader::load_elf(crate::x86::payload());
 
     crate::println!(
-        "[setup] Entering the Asterinas entry point at {:#x}",
+        "[setup] Entering the Asterinas entry point at {:p}",
         ASTER_ENTRY_POINT,
     );
-    // SAFETY: the entrypoint and the ptr is valid.
-    unsafe { call_aster_entrypoint(ASTER_ENTRY_POINT, boot_params_ptr.try_into().unwrap()) };
+    // SAFETY:
+    // 1. The entry point address is correct and matches the kernel ELF file.
+    // 2. The boot parameter pointer is valid and points to the correct boot parameters.
+    unsafe { call_aster_entrypoint(ASTER_ENTRY_POINT, boot_params_ptr) };
 }
 
-unsafe fn call_aster_entrypoint(entrypoint: u32, boot_params_ptr: u32) -> ! {
-    asm!("mov esi, {}", in(reg) boot_params_ptr);
-    asm!("mov eax, {}", in(reg) entrypoint);
-    asm!("jmp eax");
-
-    unreachable!();
+unsafe fn call_aster_entrypoint(entrypoint: *const (), boot_params_ptr: *mut BootParams) -> ! {
+    unsafe {
+        asm!(
+            "mov esi, {1}",
+            "mov eax, {0}",
+            "jmp eax",
+            in(reg) entrypoint,
+            in(reg) boot_params_ptr,
+            options(noreturn),
+        )
+    }
 }
 
 #[panic_handler]