Refine relocations in x86-64 EFI boot

osdk/src/commands/build/bin.rs

@@ -167,7 +167,37 @@ fn install_setup_with_arch(
     let target_dir = std::fs::canonicalize(target_dir).unwrap();
 
     let mut cmd = Command::new("cargo");
-    cmd.env("RUSTFLAGS", "-Ccode-model=kernel -Crelocation-model=pie -Ctarget-feature=+crt-static -Zplt=yes -Zrelax-elf-relocations=yes -Crelro-level=full");
+    let mut rustflags = vec![
+        "-Cdebuginfo=2",
+        "-Ccode-model=kernel",
+        "-Crelocation-model=pie",
+        "-Zplt=yes",
+        "-Zrelax-elf-relocations=yes",
+        "-Crelro-level=full",
+    ];
+    let target_feature_args = match arch {
+        SetupInstallArch::X86_64 => {
+            concat!(
+                "-Ctarget-feature=",
+                "+crt-static",
+                ",-adx",
+                ",-aes",
+                ",-avx",
+                ",-avx2",
+                ",-fxsr",
+                ",-sse",
+                ",-sse2",
+                ",-sse3",
+                ",-sse4.1",
+                ",-sse4.2",
+                ",-ssse3",
+                ",-xsave",
+            )
+        }
+        SetupInstallArch::Other(_) => "-Ctarget-feature=+crt-static",
+    };
+    rustflags.push(target_feature_args);
+    cmd.env("RUSTFLAGS", rustflags.join(" "));
     cmd.arg("install").arg("linux-bzimage-setup");
     cmd.arg("--force");
     cmd.arg("--root").arg(install_dir.as_ref());

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

@@ -2,17 +2,17 @@
 
 use linux_boot_params::BootParams;
 use uefi::{
+    boot::{exit_boot_services, open_protocol_exclusive},
+    mem::memory_map::{MemoryMap, MemoryMapOwned},
     prelude::*,
     proto::loaded_image::LoadedImage,
-    boot::{exit_boot_services, open_protocol_exclusive},
-    mem::memory_map::{MemoryMapOwned, MemoryMap},
 };
 use uefi_raw::table::system::SystemTable;
 
 use super::{
     decoder::decode_payload,
     paging::{Ia32eFlags, PageNumber, PageTableCreator},
-    relocation::apply_rela_dyn_relocations,
+    relocation::apply_rela_relocations,
 };
 
 // Suppress warnings since using todo!.
@@ -21,10 +21,9 @@ use super::{
 #[allow(clippy::diverging_sub_expression)]
 #[export_name = "efi_stub_entry"]
 extern "sysv64" fn efi_stub_entry(handle: Handle, system_table: *const SystemTable) -> ! {
-    unsafe {
-        boot::set_image_handle(handle);
-        uefi::table::set_system_table(system_table);
-    }
+    // SAFETY: handle and system_table are valid pointers. It is only called once.
+    unsafe { system_init(handle, system_table) };
+
     uefi::helpers::init().unwrap();
 
     let boot_params = todo!("Use EFI boot services to fill boot params");
@@ -38,10 +37,9 @@ extern "sysv64" fn efi_handover_entry(
     system_table: *const SystemTable,
     boot_params_ptr: *mut BootParams,
 ) -> ! {
-    unsafe {
-        boot::set_image_handle(handle);
-        uefi::table::set_system_table(system_table);
-    }
+    // 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.
@@ -50,18 +48,36 @@ extern "sysv64" fn efi_handover_entry(
     efi_phase_boot(boot_params)
 }
 
-fn efi_phase_boot(boot_params: &mut BootParams) -> ! {
+/// 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: This is the right time to initialize the console and it is only
     // called once here before all console operations.
     unsafe {
         crate::console::init();
     }
 
-    // SAFETY: this is the right time to apply relocations.
-    unsafe { apply_rela_dyn_relocations() };
+    // SAFETY: This is the right time to apply relocations.
+    unsafe { apply_rela_relocations() };
 
+    // SAFETY: The handle and system_table are valid pointers. They are passed
+    // from the UEFI firmware. They are only called once.
+    unsafe {
+        boot::set_image_handle(handle);
+        uefi::table::set_system_table(system_table);
+    }
+}
+
+fn efi_phase_boot(boot_params: &mut BootParams) -> ! {
     uefi::println!("[EFI stub] Relocations applied.");
-    uefi::println!("[EFI stub] Stub loaded at {:#x?}", crate::x86::get_image_loaded_offset());
+    uefi::println!(
+        "[EFI stub] Stub loaded at {:#x?}",
+        crate::x86::get_image_loaded_offset()
+    );
 
     // Fill the boot params with the RSDP address if it is not provided.
     if boot_params.acpi_rsdp_addr == 0 {
@@ -84,9 +100,7 @@ fn efi_phase_boot(boot_params: &mut BootParams) -> ! {
     };
     // SAFETY: All allocations in the boot services phase are not used after
     // this point.
-    let memory_map = unsafe {
-        exit_boot_services(memory_type)
-    };
+    let memory_map = unsafe { exit_boot_services(memory_type) };
 
     efi_phase_runtime(memory_map, boot_params);
 }

ostd/libs/linux-bzimage/setup/src/x86/amd64_efi/linker.ld

@@ -41,14 +41,10 @@ SECTIONS
         *(.eh_frame_hdr .eh_frame_hdr.*)
     }
     
-    .rela.dyn               : {
-        PROVIDE(__rela_dyn_start = .);
-        *(.rela.dyn .rela.dyn.*)
-        PROVIDE(__rela_dyn_end = .);
-    }
-
-    .rela.plt               : {
-        *(.rela.plt .rela.plt.*)
+    .rela                   : {
+        PROVIDE(__rela_start = .);
+        *(.rela .rela.*)
+        PROVIDE(__rela_end = .);
     }
 
     .comment                : { *(.comment) }

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

@@ -18,6 +18,7 @@ const TABLE_ENTRY_COUNT: usize = 512;
 
 bitflags::bitflags! {
     #[derive(Clone, Copy)]
+    #[repr(C)]
     pub struct Ia32eFlags: u64 {
         const PRESENT =         1 << 0;
         const WRITABLE =        1 << 1;
@@ -32,9 +33,11 @@ bitflags::bitflags! {
     }
 }
 
+#[repr(C)]
 pub struct Ia32eEntry(u64);
 
 /// The table in the IA32E paging specification that occupies a physical page frame.
+#[repr(C)]
 pub struct Ia32eTable([Ia32eEntry; TABLE_ENTRY_COUNT]);
 
 /// A page number. It could be either a physical page number or a virtual page number.

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

@@ -2,62 +2,68 @@
 
 use crate::x86::get_image_loaded_offset;
 
-struct Elf64Rela {
-    r_offset: u64,
-    r_info: u64,
-    r_addend: i64,
-}
+/// Apply the relocations in the `.rela.*` sections.
+///
+/// The function will enable dyn Trait objects to work since they rely on
+/// vtable pointers. Vtable won't work without relocations.
+///
+/// We currently support R_X86_64_RELATIVE relocations only. And this type of
+/// relocation seems to be the only existing type if we compile Rust code to
+/// PIE ELF binaries.
+///
+/// # Safety
+///
+/// This function will modify the memory pointed by the relocations. And the
+/// Rust memory safety mechanisms are not aware of these kind of modification.
+/// Failure to do relocations will cause `dyn Trait` objects to break.
+pub unsafe fn apply_rela_relocations() {
+    use core::arch::asm;
+    let image_loaded_offset = get_image_loaded_offset();
 
-fn get_rela_array() -> &'static [Elf64Rela] {
-    extern "C" {
-        fn __rela_dyn_start();
-        fn __rela_dyn_end();
+    let mut start: usize;
+    let end: usize;
+
+    unsafe {
+        asm!(
+            "lea {}, [rip + __rela_start]",
+            out(reg) start,
+        );
+        asm!(
+            "lea {}, [rip + __rela_end]",
+            out(reg) end,
+        );
     }
-    let start = __rela_dyn_start as *const Elf64Rela;
-    let end = __rela_dyn_end as *const Elf64Rela;
-    let len = unsafe { end.offset_from(start) } as usize;
+
     #[cfg(feature = "debug_print")]
     unsafe {
         use crate::console::{print_hex, print_str};
-        print_str("[EFI stub debug] .rela.dyn section size = ");
-        print_hex(len as u64);
-        print_str("; __rela_dyn_start = ");
+        print_str("[EFI stub debug] loaded offset = ");
+        print_hex(image_loaded_offset as u64);
+        print_str("\n");
+        print_str("[EFI stub debug] .rela section start = ");
         print_hex(start as u64);
-        print_str(", __rela_dyn_end = ");
+        print_str(", end = ");
         print_hex(end as u64);
         print_str("\n");
     }
-    // SAFETY: the linker will ensure that the symbols are valid.
-    unsafe { core::slice::from_raw_parts(start, len) }
-}
 
-const R_X86_64_RELATIVE: u32 = 8;
+    #[cfg(feature = "debug_print")]
+    let mut count = 0;
 
-/// Apply the relocations in the `.rela.dyn` section.
-///
-/// The function will enable dyn Trait objects to work since they rely on vtable pointers. Vtable
-/// won't work without relocations.
-///
-/// We currently support R_X86_64_RELATIVE relocations only. And this type of relocation seems to
-/// be the only existing type if we compile Rust code to PIC ELF binaries.
-///
-/// # Safety
-/// This function will modify the memory pointed by the relocations. And the Rust memory safety
-/// mechanisms are not aware of these kind of modification. Failure to do relocations will cause
-/// dyn Trait objects to break.
-pub unsafe fn apply_rela_dyn_relocations() {
-    let image_loaded_offset = get_image_loaded_offset();
-    let relas = get_rela_array();
-    for rela in relas {
+    while start < end {
+        let rela = (start as *const Elf64Rela).read_volatile();
         let r_type = (rela.r_info & 0xffffffff) as u32;
         let _r_sym = (rela.r_info >> 32) as usize;
-        let r_addend = rela.r_addend;
-        let r_offset = rela.r_offset as usize;
-        let target = (image_loaded_offset + r_offset as isize) as usize;
+        let r_addend = rela.r_addend as isize;
+        let r_offset = rela.r_offset as isize;
+        let target = image_loaded_offset.wrapping_add(r_offset) as usize;
         #[cfg(feature = "debug_print")]
         unsafe {
             use crate::console::{print_hex, print_str};
-            print_str("[EFI stub debug] Applying relocation at offset ");
+            count += 1;
+            print_str("[EFI stub debug] Applying relocation #");
+            print_hex(count as u64);
+            print_str(" at offset ");
             print_hex(r_offset as u64);
             print_str(", type = ");
             print_hex(r_type as u64);
@@ -67,12 +73,23 @@ pub unsafe fn apply_rela_dyn_relocations() {
         }
         match r_type {
             R_X86_64_RELATIVE => {
-                let value = (image_loaded_offset as i64 + r_addend) as usize;
-                *(target as *mut usize) = value;
+                let value = image_loaded_offset.wrapping_add(r_addend) as usize;
+                (target as *mut usize).write(value);
             }
             _ => {
                 panic!("Unknown relocation type: {}", r_type);
             }
         }
+        start = start.wrapping_add(core::mem::size_of::<Elf64Rela>());
     }
 }
+
+const R_X86_64_RELATIVE: u32 = 8;
+
+#[derive(Copy, Clone)]
+#[repr(C)]
+struct Elf64Rela {
+    r_offset: u64,
+    r_info: u64,
+    r_addend: i64,
+}

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

@@ -16,10 +16,23 @@ const START_OF_SETUP32_VA: usize = 0x100000;
 /// The setup is a position-independent executable. We can get the loaded base
 /// address from the symbol.
 #[inline]
-#[allow(clippy::fn_to_numeric_cast)]
 pub fn get_image_loaded_offset() -> isize {
-    extern "C" {
-        fn start_of_setup32();
+    let address_of_start: usize;
+    #[cfg(target_arch = "x86_64")]
+    unsafe {
+        core::arch::asm!(
+            "lea {}, [rip + start_of_setup32]",
+            out(reg) address_of_start,
+            options(pure, nomem, nostack)
+        );
     }
-    start_of_setup32 as isize - START_OF_SETUP32_VA as isize
+    #[cfg(target_arch = "x86")]
+    unsafe {
+        core::arch::asm!(
+            "lea {}, [start_of_setup32]",
+            out(reg) address_of_start,
+            options(pure, nomem, nostack)
+        );
+    }
+    address_of_start as isize - START_OF_SETUP32_VA as isize
 }