Apply relocations to enable vtable

framework/libs/boot-trojan/builder/src/lib.rs

@@ -157,12 +157,16 @@ fn build_trojan_with_arch(source_dir: &Path, out_dir: &Path, arch: &TrojanBuildA
     }
     let out_dir = std::fs::canonicalize(out_dir).unwrap();
 
+    // Relocations are fewer in release mode. But with release mode it crashes.
+    let profile = "debug";
+
     let cargo = std::env::var("CARGO").unwrap();
     let mut cmd = std::process::Command::new(cargo);
     cmd.current_dir(source_dir);
     cmd.arg("build");
-    // Relocations are fewer in release mode, saving header real-estate.
-    cmd.arg("--release");
+    if profile == "release" {
+        cmd.arg("--release");
+    }
     cmd.arg("--package").arg("aster-boot-trojan");
     cmd.arg("--target").arg(match arch {
         TrojanBuildArch::X86_64 => "x86_64-unknown-none",
@@ -193,7 +197,7 @@ fn build_trojan_with_arch(source_dir: &Path, out_dir: &Path, arch: &TrojanBuildA
 
     let trojan_artifact = out_dir
         .join(arch_name)
-        .join("release")
+        .join(profile)
         .join("aster-boot-trojan");
 
     trojan_artifact.to_owned()

framework/libs/boot-trojan/trojan/Cargo.toml

@@ -15,3 +15,7 @@ xmas-elf = "0.8.0"
 log = "0.4.20"
 uefi = "0.26.0"
 uefi-services = "0.23.0"
+
+[features]
+default = ["debug_print"]
+debug_print = []

framework/libs/boot-trojan/trojan/build.rs

@@ -4,9 +4,9 @@ fn main() {
     let source_dir = PathBuf::from(std::env::var("CARGO_MANIFEST_DIR").unwrap());
     let target_arch = std::env::var("TARGET").unwrap();
     let linker_script = if target_arch == "x86_64-unknown-none" {
-        source_dir.join("src/arch/x86_64/linker.ld")
+        source_dir.join("src/x86/amd64_efi/linker.ld")
     } else if target_arch == "x86_64-i386_pm-none" {
-        source_dir.join("src/arch/i386/linker.ld")
+        source_dir.join("src/arch/legacy_i386/linker.ld")
     } else {
         panic!("Unsupported target_arch: {}", target_arch);
     };

framework/libs/boot-trojan/trojan/src/arch/x86_64/efi.rs

@@ -1,66 +0,0 @@
-use core::fmt::Write;
-use uefi::{
-    data_types::Handle,
-    proto::loaded_image::LoadedImage,
-    table::{Boot, SystemTable},
-};
-
-use linux_boot_params::BootParams;
-
-#[no_mangle]
-extern "sysv64" fn efi_stub_entry(handle: Handle, mut system_table: SystemTable<Boot>) -> ! {
-    unsafe {
-        system_table.boot_services().set_image_handle(handle);
-    }
-    uefi_services::init(&mut system_table).unwrap();
-
-    // Suppress TODO warning.
-    #[allow(unreachable_code)]
-    efi_entry(
-        handle,
-        system_table,
-        todo!("Use EFI services to fill boot params"),
-    );
-}
-
-#[no_mangle]
-extern "sysv64" fn efi_handover_entry(
-    handle: Handle,
-    mut system_table: SystemTable<Boot>,
-    boot_params: *mut BootParams,
-) -> ! {
-    unsafe {
-        system_table.boot_services().set_image_handle(handle);
-    }
-    uefi_services::init(&mut system_table).unwrap();
-
-    efi_entry(handle, system_table, boot_params)
-}
-
-fn efi_entry(
-    handle: Handle,
-    mut system_table: SystemTable<Boot>,
-    boot_params: *mut BootParams,
-) -> ! {
-    system_table
-        .stdout()
-        .write_str("[EFI stub] Exiting EFI boot services.\n")
-        .unwrap();
-    let memory_type = {
-        let boot_services = system_table.boot_services();
-        let Ok(loaded_image) = boot_services.open_protocol_exclusive::<LoadedImage>(handle) else {
-            panic!("Failed to open LoadedImage protocol");
-        };
-        loaded_image.data_type().clone()
-    };
-    let _ = system_table.exit_boot_services(memory_type);
-
-    let loaded_base = {
-        extern "C" {
-            fn start_of_setup32();
-        }
-        start_of_setup32 as usize
-    };
-
-    crate::trojan_entry(loaded_base, boot_params as usize);
-}

framework/libs/boot-trojan/trojan/src/console.rs

@@ -1,4 +1,4 @@
-use core::fmt::Write;
+use core::fmt::{self, Write};
 
 use uart_16550::SerialPort;
 
@@ -17,45 +17,73 @@ pub unsafe fn init() {
 
 impl Stdout {
     /// safety: this function must only be called once
-    unsafe fn init() -> Self {
+    pub unsafe fn init() -> Self {
         let mut serial_port = unsafe { SerialPort::new(0x3F8) };
         serial_port.init();
         Self { serial_port }
     }
 }
 
-impl Stdout {
-    fn write_str(&mut self, s: &str) {
+impl Write for Stdout {
+    fn write_str(&mut self, s: &str) -> fmt::Result {
         self.serial_port.write_str(s).unwrap();
-    }
-
-    fn write_char(&mut self, c: char) {
-        self.serial_port.send(c as u8);
+        Ok(())
     }
 }
 
-/// Safety: init() must be called before print() and there should be no race condition
-pub unsafe fn print(s: &str) {
-    STDOUT.write_str(s);
+/// This is used when dyn Trait is not supported or fmt::Arguments is fragile to use in PIE.
+///
+/// Safety: init() must be called before print_str() and there should be no race conditions.
+pub unsafe fn print_str(s: &str) {
+    STDOUT.write_str(s).unwrap();
 }
 
-/// Safety: init() must be called before print_char() and there should be no race condition
-pub unsafe fn print_char(c: char) {
-    STDOUT.write_char(c);
+unsafe fn print_char(c: char) {
+    STDOUT.serial_port.send(c as u8);
 }
 
-// Safety: init() must be called before print_hex() and there should be no race condition
-pub unsafe fn print_hex(n: usize) {
-    print("0x");
-    let mut n = n;
-    for _ in 0..16 {
-        let digit = (n & 0xf) as u8;
-        n >>= 4;
-        let c = if digit < 10 {
-            (b'0' + digit) as char
+/// This is used when dyn Trait is not supported or fmt::Arguments is fragile to use in PIE.
+///
+/// Safety: init() must be called before print_hex() and there should be no race conditions.
+pub unsafe fn print_hex(n: u64) {
+    print_str("0x");
+    for i in (0..16).rev() {
+        let digit = (n >> (i * 4)) & 0xf;
+        if digit < 10 {
+            print_char((b'0' + digit as u8) as char);
         } else {
-            (b'a' + digit - 10) as char
-        };
-        print_char(c);
+            print_char((b'A' + (digit - 10) as u8) as char);
+        }
     }
 }
+
+// TODO: Figure out why fmt::Arguments wont work even if relocations are applied.
+// We just settle on simple print functions for now.
+/*--------------------------------------------------------------------------------------------------
+
+/// Glue code for print!() and println!() macros.
+///
+/// Safety: init() must be called before print_fmt() and there should be no race conditions.
+pub unsafe fn print_fmt(args: fmt::Arguments) {
+    STDOUT.write_fmt(args).unwrap();
+}
+
+#[macro_export]
+macro_rules! print {
+    ($fmt: literal $(, $($arg: tt)+)?) => {
+        unsafe {
+            $crate::console::print_fmt(format_args!($fmt $(, $($arg)+)?))
+        }
+    }
+}
+
+#[macro_export]
+macro_rules! println {
+    ($fmt: literal $(, $($arg: tt)+)?) => {
+        unsafe {
+            $crate::console::print_fmt(format_args!(concat!($fmt, "\n") $(, $($arg)+)?))
+        }
+    }
+}
+
+ *------------------------------------------------------------------------------------------------*/

framework/libs/boot-trojan/trojan/src/loader.rs

@@ -18,6 +18,20 @@ pub fn load_elf(file: &[u8]) -> u32 {
                     program.mem_size as usize,
                 )
             };
+            /* crate::println!(
+                "[setup loader debug] loading ELF segment at {:#x}, size = {:#x}",
+                program.physical_addr,
+                program.mem_size,
+            ); */
+            #[cfg(feature = "debug_print")]
+            unsafe {
+                use crate::console::{print_hex, print_str};
+                print_str("[setup loader debug] loading ELF segment at ");
+                print_hex(program.physical_addr as u64);
+                print_str(", size = ");
+                print_hex(program.mem_size as u64);
+                print_str("\n");
+            }
             dst_slice[..program.file_size as usize].copy_from_slice(header_data);
             let zero_slice = &mut dst_slice[program.file_size as usize..];
             zero_slice.fill(0);
@@ -25,5 +39,5 @@ pub fn load_elf(file: &[u8]) -> u32 {
     }
 
     // Return the Linux Boot Protocol entry point defined by Asterinas.
-    crate::arch::ASTER_ENTRY_POINT
+    crate::x86::ASTER_ENTRY_POINT
 }

framework/libs/boot-trojan/trojan/src/main.rs

@@ -3,46 +3,49 @@
 
 use linux_boot_params::BootParams;
 
-mod arch;
 mod console;
 mod loader;
+mod x86;
 
-use console::{print, print_hex};
+use console::{print_hex, print_str};
 
 /// The entrypoint of the trojan. The architecture-specific entrypoint will call this function.
-///
-/// The loaded address of the CODE32_START should be passed in as `loaded_base`, since the trojan
-/// may be loaded at any address, and offsets in the header are not position-independent.
-fn trojan_entry(loaded_base: usize, boot_params_ptr: usize) -> ! {
-    // Safety: this init function is only called once.
-    unsafe { console::init() };
+fn trojan_entry(boot_params_ptr: usize) -> ! {
+    // println!("[setup] bzImage loaded at {:#x}", x86::relocation::get_image_loaded_offset());
     unsafe {
-        print("[setup] bzImage loaded at ");
-        print_hex(loaded_base);
-        print("\n");
+        print_str("[setup] bzImage loaded at ");
+        print_hex(x86::relocation::get_image_loaded_offset() as u64);
+        print_str("\n");
     }
 
     // Safety: the boot_params_ptr is a valid pointer to be borrowed.
     let boot_params = unsafe { &*(boot_params_ptr as *const BootParams) };
-    let hdr = &boot_params.hdr;
-    let payload_offset = loaded_base + hdr.payload_offset as usize;
-    let payload_length = hdr.payload_length as usize;
+    // Safety: the payload_offset and payload_length is valid.
     let payload = unsafe {
+        let hdr = &boot_params.hdr;
+        // The payload_offset field is not recorded in the relocation table, so we need to
+        // calculate the loaded offset manually.
+        let loaded_offset = x86::relocation::get_image_loaded_offset();
+        let payload_offset = (loaded_offset + hdr.payload_offset as isize) as usize;
+        let payload_length = hdr.payload_length as usize;
         core::slice::from_raw_parts_mut(payload_offset as *mut u8, payload_length as usize)
     };
 
+    // println!("[setup] loading ELF payload at {:#x}", payload as *const _ as *const u8 as usize);
     unsafe {
-        print("[setup] loading ELF payload at ");
-        print_hex(payload_offset);
-        print("...\n");
+        print_str("[setup] loading ELF payload at ");
+        print_hex(payload as *const _ as *const u8 as u64);
+        print_str("\n");
     }
     let entrypoint = loader::load_elf(payload);
 
+    // println!("[setup] jumping to payload entrypoint at {:#x}", entrypoint);
     unsafe {
-        print("[setup] jumping to payload entrypoint at ");
-        print_hex(entrypoint as usize);
-        print("...\n");
+        print_str("[setup] jumping to payload entrypoint at ");
+        print_hex(entrypoint as u64);
+        print_str("\n");
     }
+
     // Safety: the entrypoint and the ptr is valid.
-    unsafe { arch::call_aster_entrypoint(entrypoint.into(), boot_params_ptr.try_into().unwrap()) };
+    unsafe { x86::call_aster_entrypoint(entrypoint.into(), boot_params_ptr.try_into().unwrap()) };
 }

framework/libs/boot-trojan/trojan/src/x86/amd64_efi/efi.rs

@@ -0,0 +1,94 @@
+use core::fmt::Write;
+use uefi::{
+    data_types::Handle,
+    proto::loaded_image::LoadedImage,
+    table::{boot::MemoryMap, Boot, Runtime, SystemTable},
+};
+
+use linux_boot_params::BootParams;
+
+#[no_mangle]
+extern "sysv64" fn efi_stub_entry(handle: Handle, mut system_table: SystemTable<Boot>) -> ! {
+    unsafe {
+        system_table.boot_services().set_image_handle(handle);
+    }
+    uefi_services::init(&mut system_table).unwrap();
+
+    // Suppress TODO warning.
+    #[allow(unreachable_code)]
+    efi_phase_boot(
+        handle,
+        system_table,
+        todo!("Use EFI boot services to fill boot params"),
+    );
+}
+
+#[no_mangle]
+extern "sysv64" fn efi_handover_entry(
+    handle: Handle,
+    mut system_table: SystemTable<Boot>,
+    boot_params: *mut BootParams,
+) -> ! {
+    unsafe {
+        system_table.boot_services().set_image_handle(handle);
+    }
+    uefi_services::init(&mut system_table).unwrap();
+
+    efi_phase_boot(handle, system_table, boot_params)
+}
+
+fn efi_phase_boot(
+    handle: Handle,
+    mut system_table: SystemTable<Boot>,
+    boot_params: *mut BootParams,
+) -> ! {
+    // Safety: this init function is only called once.
+    unsafe { crate::console::init() };
+    
+    // Safety: this is a right place to call this function.
+    unsafe { crate::x86::relocation::apply_rela_dyn_relocations() };
+
+    uefi_services::println!("[EFI stub] Relocations applied.");
+    uefi_services::println!("[EFI stub] Exiting EFI boot services.");
+    let memory_type = {
+        let boot_services = system_table.boot_services();
+        let Ok(loaded_image) = boot_services.open_protocol_exclusive::<LoadedImage>(handle) else {
+            panic!("Failed to open LoadedImage protocol");
+        };
+        loaded_image.data_type().clone()
+    };
+    let (system_table, memory_map) = system_table.exit_boot_services(memory_type);
+
+    efi_phase_runtime(system_table, memory_map, boot_params);
+}
+
+fn efi_phase_runtime(
+    mut system_table: SystemTable<Runtime>,
+    memory_map: MemoryMap<'static>,
+    boot_params: *mut BootParams,
+) -> ! {
+    unsafe {
+        crate::console::print_str("[EFI stub] Entered runtime services.\n");
+    }
+
+    #[cfg(feature = "debug_print")]
+    unsafe {
+        use crate::console::{print_hex, print_str};
+        print_str("[EFI stub debug] Memory map:\n");
+        for md in memory_map.entries() {
+            // crate::println!("    [{:#x}] {:#x} ({:#x})", md.ty.0, md.phys_start, md.page_count);
+            print_str("    [");
+            print_hex(md.ty.0 as u64);
+            print_str("]");
+            print_hex(md.phys_start);
+            print_str("(size=");
+            print_hex(md.page_count);
+            print_str(")");
+            print_str("{flags=");
+            print_hex(md.att.bits());
+            print_str("}\n");
+        }
+    }
+
+    crate::trojan_entry(boot_params as usize);
+}

framework/libs/boot-trojan/trojan/src/x86/amd64_efi/linker.ld

@@ -42,7 +42,9 @@ SECTIONS
     }
 
     .rela.dyn               : {
+        PROVIDE(__rela_dyn_start = .);
         *(.rela.dyn .rela.dyn.*)
+        PROVIDE(__rela_dyn_end = .);
     }
 
     .rela.plt               : {

framework/libs/boot-trojan/trojan/src/x86/legacy_i386/mod.rs

@@ -1,5 +1,3 @@
-use crate::println;
-
 use core::arch::{asm, global_asm};
 
 global_asm!(include_str!("header.S"));
@@ -8,7 +6,9 @@ global_asm!(include_str!("setup.S"));
 
 #[no_mangle]
 extern "cdecl" fn _trojan_entry_32(boot_params_ptr: u32) -> ! {
-    crate::trojan_entry(0x100000, boot_params_ptr.try_into().unwrap());
+    // Safety: this init function is only called once.
+    unsafe { crate::console::init() };
+    crate::trojan_entry(0, boot_params_ptr.try_into().unwrap());
 }
 
 pub const ASTER_ENTRY_POINT: u32 = 0x8001000;
@@ -23,6 +23,5 @@ pub unsafe fn call_aster_entrypoint(entrypoint: u32, boot_params_ptr: u32) -> !
 
 #[panic_handler]
 fn panic(info: &core::panic::PanicInfo) -> ! {
-    println!("panic: {:?}", info);
     loop {}
 }

framework/libs/boot-trojan/trojan/src/x86/mod.rs

@@ -1,11 +1,13 @@
 cfg_if::cfg_if! {
     if #[cfg(target_arch = "x86_64")] {
-        mod x86_64;
-        pub use x86_64::*;
+        mod amd64_efi;
+        pub use amd64_efi::*;
     } else if #[cfg(target_arch = "x86")] {
-        mod i386;
-        pub use i386::*;
+        mod legacy_i386;
+        pub use legacy_i386::*;
     } else {
         compile_error!("Unsupported target_arch");
     }
 }
+
+pub mod relocation;

framework/libs/boot-trojan/trojan/src/x86/relocation.rs

@@ -0,0 +1,92 @@
+// This is enforced in the linker script.
+const START_OF_SETUP32_VA: usize = 0x100000;
+
+/// The trojan is a position-independent executable. We can get the loaded base
+/// address from the symbol.
+#[inline]
+pub fn get_image_loaded_offset() -> isize {
+    extern "C" {
+        fn start_of_setup32();
+    }
+    start_of_setup32 as isize - START_OF_SETUP32_VA as isize
+}
+
+struct Elf64Rela {
+    r_offset: u64,
+    r_info: u64,
+    r_addend: i64,
+}
+
+fn get_rela_array() -> &'static [Elf64Rela] {
+    extern "C" {
+        fn __rela_dyn_start();
+        fn __rela_dyn_end();
+    }
+    let start = __rela_dyn_start as *const Elf64Rela;
+    let end = __rela_dyn_end as *const Elf64Rela;
+    // FIXME: 2023/11/29
+    // There should be a Rust compiler bug that makes the calculation of len incorrect.
+    // The current implementation only works in debug mode.
+    // To do it in release mode, a workaround is to use inline asm. But that crashes
+    // in debug mode. QaQ
+    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_hex(start as u64);
+        print_str(", __rela_dyn_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 as *const Elf64Rela, len) }
+}
+
+const R_X86_64_RELATIVE: u32 = 8;
+
+/// 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 {
+        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;
+        #[cfg(feature = "debug_print")]
+        unsafe {
+            use crate::console::{print_hex, print_str};
+            print_str("[EFI stub debug] Applying relocation at offset ");
+            print_hex(r_offset as u64);
+            print_str(", type = ");
+            print_hex(r_type as u64);
+            print_str(", addend = ");
+            print_hex(r_addend as u64);
+            print_str("\n");
+        }
+        match r_type {
+            R_X86_64_RELATIVE => {
+                let value = (image_loaded_offset as i64 + r_addend) as usize;
+                *(target as *mut usize) = value;
+            }
+            _ => {
+                panic!("Unknown relocation type: {}", r_type);
+            }
+        }
+    }
+}