Make trojan compile in 64-bit mode

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

@@ -10,3 +10,7 @@ bytemuck = { version = "1.14.0", features = ["derive"] }
 bitflags = "1.3"
 serde = { version = "1.0.192", features = ["derive"] }
 xmas-elf = "0.9.1"
+
+[features]
+default = []
+trojan64 = []

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

@@ -2,13 +2,13 @@ mod mapping;
 mod pe_header;
 
 use std::{
-    error::Error,
+    ffi::OsStr,
     fs::File,
-    io::{Read, Seek, Write},
-    path::Path,
+    io::{Read, Write},
+    path::{Path, PathBuf},
 };
 
-use xmas_elf::program::{ProgramHeader, SegmentData};
+use xmas_elf::program::SegmentData;
 
 use mapping::{TrojanFileOffset, TrojanVA};
 
@@ -22,18 +22,15 @@ fn trojan_to_flat_binary(elf_file: &[u8]) -> Vec<u8> {
     let elf = xmas_elf::ElfFile::new(&elf_file).unwrap();
     let mut bin = Vec::<u8>::new();
 
-    for ph in elf.program_iter() {
-        let ProgramHeader::Ph32(program) = ph else {
-            panic!("Unexpected program header type");
-        };
+    for program in elf.program_iter() {
         if program.get_type().unwrap() == xmas_elf::program::Type::Load {
             let SegmentData::Undefined(header_data) = program.get_data(&elf).unwrap() else {
                 panic!("Unexpected segment data type");
             };
             let dst_file_offset = usize::from(TrojanFileOffset::from(TrojanVA::from(
-                program.virtual_addr as usize,
+                program.virtual_addr() as usize,
             )));
-            let dst_file_length = program.file_size as usize;
+            let dst_file_length = program.file_size() as usize;
             if bin.len() < dst_file_offset + dst_file_length {
                 bin.resize(dst_file_offset + dst_file_length, 0);
             }
@@ -64,7 +61,7 @@ fn fill_header_field(header: &mut [u8], offset: usize, value: &[u8]) {
 fn fill_legacy_header_fields(
     header: &mut [u8],
     kernel_len: usize,
-    header_len: usize,
+    trojan_len: usize,
     payload_offset: TrojanVA,
 ) {
     fill_header_field(
@@ -82,84 +79,98 @@ fn fill_legacy_header_fields(
     fill_header_field(
         header,
         0x260, /* init_size */
-        &((header_len + kernel_len) as u32).to_le_bytes(),
+        &((trojan_len + kernel_len) as u32).to_le_bytes(),
     );
 }
 
-pub fn make_bzimage(path: &Path, kernel_path: &Path, header_path: &Path) -> std::io::Result<()> {
-    let mut header_elf_file = Vec::new();
-    File::open(header_path)?.read_to_end(&mut header_elf_file)?;
-    let mut header = trojan_to_flat_binary(&header_elf_file);
-    // Pad the Linux boot header to let the payload starts with 8-byte alignment.
-    header.resize((header.len() + 7) & !7, 0x00);
+pub fn make_bzimage(path: &Path, kernel_path: &Path, trojan_src: &Path, trojan_out: &Path) {
+    #[cfg(feature = "trojan64")]
+    let trojan = build_trojan_with_arch(trojan_src, trojan_out, "x86_64-unknown-none".as_ref());
+
+    #[cfg(not(feature = "trojan64"))]
+    let trojan = {
+        let arch = trojan_src.join("x86_64-i386_pm-none.json");
+        build_trojan_with_arch(trojan_src, trojan_out, arch.as_os_str())
+    };
+
+    let mut trojan_elf = Vec::new();
+    File::open(trojan)
+        .unwrap()
+        .read_to_end(&mut trojan_elf)
+        .unwrap();
+    let mut trojan = trojan_to_flat_binary(&trojan_elf);
+    // Pad the header with 8-byte alignment.
+    trojan.resize((trojan.len() + 7) & !7, 0x00);
 
     let mut kernel = Vec::new();
-    File::open(kernel_path)?.read_to_end(&mut kernel)?;
+    File::open(kernel_path)
+        .unwrap()
+        .read_to_end(&mut kernel)
+        .unwrap();
     let payload = kernel;
 
-    let header_len = header.len();
+    let trojan_len = trojan.len();
     let payload_len = payload.len();
-    let payload_offset = TrojanFileOffset::from(header_len);
-    fill_legacy_header_fields(&mut header, payload_len, header_len, payload_offset.into());
+    let payload_offset = TrojanFileOffset::from(trojan_len);
+    fill_legacy_header_fields(&mut trojan, payload_len, trojan_len, payload_offset.into());
 
-    let mut kernel_image = File::create(path)?;
-    kernel_image.write_all(&header)?;
-    kernel_image.write_all(&payload)?;
+    let mut kernel_image = File::create(path).unwrap();
+    kernel_image.write_all(&trojan).unwrap();
+    kernel_image.write_all(&payload).unwrap();
 
-    let image_size = header_len + payload_len;
+    let image_size = trojan_len + payload_len;
 
     // Since the Linux boot header starts at 0x1f1, we can write the PE/COFF header directly to the
     // start of the file without overwriting the Linux boot header.
-    let pe_header = pe_header::make_pe_coff_header(&header_elf_file, image_size);
+    let pe_header = pe_header::make_pe_coff_header(&trojan_elf, image_size);
     assert!(
         pe_header.header_at_zero.len() <= 0x1f1,
         "PE/COFF header is too large"
     );
 
-    // FIXME: Oops, EFI hanover stucks, so I removed the pe header to let grub go through the legacy path.
-    kernel_image.seek(std::io::SeekFrom::Start(0))?;
-    // kernel_image.write_all(&pe_header.header_at_zero)?;
-    kernel_image.seek(std::io::SeekFrom::Start(
-        usize::from(pe_header.relocs.0) as u64
-    ))?;
-    // kernel_image.write_all(&pe_header.relocs.1)?;
-
-    Ok(())
+    #[cfg(feature = "trojan64")]
+    {
+        use std::io::{Seek, SeekFrom};
+        kernel_image.seek(SeekFrom::Start(0)).unwrap();
+        kernel_image.write_all(&pe_header.header_at_zero).unwrap();
+        kernel_image
+            .seek(SeekFrom::Start(usize::from(pe_header.relocs.0) as u64))
+            .unwrap();
+        kernel_image.write_all(&pe_header.relocs.1).unwrap();
+    }
 }
 
-pub fn build_linux_setup_header_from_trojan(
-    source_dir: &Path,
-    out_dir: &Path,
-) -> Result<(), Box<dyn Error + Send + Sync>> {
-    // Build the setup header to ELF.
-    let target_json = source_dir.join("x86_64-i386_protected_mode.json");
-
+fn build_trojan_with_arch(source_dir: &Path, out_dir: &Path, arch: &OsStr) -> PathBuf {
     let cargo = std::env::var("CARGO").unwrap();
     let mut cmd = std::process::Command::new(cargo);
-    cmd.arg("install").arg("aster-boot-trojan");
-    cmd.arg("--debug");
-    cmd.arg("--locked");
-    cmd.arg("--path").arg(source_dir.to_str().unwrap());
-    cmd.arg("--target").arg(target_json.as_os_str());
-    cmd.arg("-Zbuild-std=core,compiler_builtins");
+    cmd.arg("build");
+    cmd.arg("--package").arg("aster-boot-trojan");
+    cmd.arg("--manifest-path")
+        .arg(source_dir.join("Cargo.toml").as_os_str());
+    cmd.arg("--target").arg(arch);
+    cmd.arg("-Zbuild-std=core,alloc,compiler_builtins");
     cmd.arg("-Zbuild-std-features=compiler-builtins-mem");
-    // Specify the installation root.
-    cmd.arg("--root").arg(out_dir.as_os_str());
     // Specify the build target directory to avoid cargo running
     // into a deadlock reading the workspace files.
     cmd.arg("--target-dir").arg(out_dir.as_os_str());
     cmd.env_remove("RUSTFLAGS");
     cmd.env_remove("CARGO_ENCODED_RUSTFLAGS");
-    let output = cmd.output()?;
-    if !output.status.success() {
-        std::io::stdout().write_all(&output.stdout).unwrap();
-        std::io::stderr().write_all(&output.stderr).unwrap();
-        return Err(format!(
+    let mut child = cmd.spawn().unwrap();
+    let status = child.wait().unwrap();
+    if !status.success() {
+        panic!(
             "Failed to build linux x86 setup header:\n\tcommand `{:?}`\n\treturned {}",
-            cmd, output.status
-        )
-        .into());
+            cmd, status
+        );
     }
 
-    Ok(())
+    // If the arch is a builtin target rather than json, the path operation works as well.
+    let arch_name = Path::new(arch).file_stem().unwrap().to_str().unwrap();
+
+    let trojan_artifact = out_dir
+        .join(arch_name)
+        .join("debug")
+        .join("aster-boot-trojan");
+
+    trojan_artifact.to_owned()
 }

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

@@ -7,6 +7,7 @@ use std::{cmp::PartialOrd, convert::From, ops::Sub};
 // is page-aligned and the legacy setup section size would be 0x1000.
 pub const LEGACY_SETUP_SECS: usize = 7;
 pub const LEGACY_SETUP_SEC_SIZE: usize = 0x200 * (LEGACY_SETUP_SECS + 1);
+
 pub const SETUP32_LMA: usize = 0x100000;
 
 #[derive(PartialEq, Eq, PartialOrd, Ord, Debug, Clone, Copy)]

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

@@ -6,5 +6,6 @@ edition = "2021"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
+cfg-if = "1.0.0"
 uart_16550 = "0.3.0"
 xmas-elf = "0.8.0"

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

@@ -2,8 +2,17 @@ use std::path::PathBuf;
 
 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")
+    } else if target_arch == "x86_64-i386_pm-none" {
+        source_dir.join("src/arch/i386.linker.ld")
+    } else {
+        panic!("Unsupported target_arch: {}", target_arch);
+    };
+    println!("cargo:rerun-if-changed={}", linker_script.display());
     println!(
         "cargo:rustc-link-arg-bins=--script={}",
-        source_dir.join("linker.ld").display()
-    )
+        linker_script.display()
+    );
 }

framework/libs/boot-trojan/trojan/src/arch/header.S

@@ -3,8 +3,7 @@
 // more information on the Linux x86 Boot Protocol.
 
 // Some of the fields filled with a 0xab* values should be filled
-// by the runner, which is the only tool after building and can
-// access the info of the payload.
+// by the torjan builder.
 // Asterinas will use only a few of these fields, and some of them
 // are filled by the loader and will be read by Asterinas.
 
@@ -47,11 +46,11 @@ xloadflags:             .word 0b01111 # all handover protocols except kexec
 cmdline_size:           .long 4096-1
 hardware_subarch:       .long 0
 hardware_subarch_data:  .quad 0
-payload_offset:         .long 0xabababab # at 0x248/4, to be filled by the runner
-payload_length:         .long 0xabababab # at 0x24c/4, to be filled by the runner
+payload_offset:         .long 0xabababab # at 0x248/4, to be filled by the builder
+payload_length:         .long 0xabababab # at 0x24c/4, to be filled by the builder
 setup_data:             .quad 0
 pref_address:           .quad CODE32_START - 0x200 * (SETUP_SECTS + 1);
-init_size:              .long 0xabababab # at 0x260/4, to be filled by the runner
-handover_offset:        .long start_of_setup32
+init_size:              .long 0xabababab # at 0x260/4, to be filled by the builder
+handover_offset:        .long CODE32_START
 kernel_info_offset:     .long 0
 hdr_end:

framework/libs/boot-trojan/trojan/src/arch/i386.linker.ld

@@ -1,4 +1,4 @@
-ENTRY(start_of_setup64)
+ENTRY(start_of_setup32)
 OUTPUT_ARCH(i386:x86)
 OUTPUT_FORMAT(elf32-i386)
 
@@ -6,24 +6,18 @@ SETUP32_LMA = 0x100000;
 BOOTSECT_SIZE = 0x1000;
 BOOTSECT_START = SETUP32_LMA - BOOTSECT_SIZE;
 
+LEGACY_STACK_SIZE = 0x1000;
+
 SECTIONS
 {
     . = BOOTSECT_START;
     .header                 : { KEEP(*(.header)) }
 
     . = SETUP32_LMA;
-    .setup                  : {
-        PROVIDE(__setup_start = .);
-        KEEP(*(.header))
-        PROVIDE(__setup_end = .);
-    }
-    .stack                  : { KEEP(*(.stack)) }
+    .setup                  : { KEEP(*(.setup)) }
+    .stack                  : { *(.stack) }
 
-    .text                   : {
-        PROVIDE(__text_start = .);
-        *(.text .text.*)
-        PROVIDE(__text_end = .);
-    }
+    .text                   : { *(.text .text.*) }
     .rodata                 : { *(.rodata .rodata.*) }
 
     .data                   : { *(.data .data.*) }
@@ -39,14 +33,4 @@ SECTIONS
     .eh_frame_hdr           : {
         *(.eh_frame_hdr .eh_frame_hdr.*)
     }
-
-    .symtab                 : {
-        *(.symtab .symtab.*)
-    }
-    .strtab                 : {
-        *(.strtab .strtab.*)
-    }
-    .shstrtab               : {
-        *(.shstrtab .shstrtab.*)
-    }
 }

framework/libs/boot-trojan/trojan/src/arch/i386.rs

@@ -0,0 +1,18 @@
+use core::arch::{asm, global_asm};
+
+global_asm!(include_str!("header.S"));
+
+global_asm!(include_str!("setup.S"));
+
+#[no_mangle]
+extern "cdecl" fn _trojan_entry_32(boot_params_ptr: u32) -> ! {
+    crate::trojan_entry(boot_params_ptr);
+}
+
+pub 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!();
+}

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

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

framework/libs/boot-trojan/trojan/src/arch/setup.S

@@ -0,0 +1,28 @@
+// 32-bit setup code starts here, and will be loaded at CODE32_START.
+.section ".setup", "ax"
+.code32
+.global start_of_setup32
+start_of_setup32:
+    mov eax, offset __stack_top
+    mov esp, eax
+    mov eax, offset halt
+    push eax                        # the return address
+    mov ebp, esp
+    add ebp, -4
+    push ebp
+    mov ebp, esp
+
+.extern _trojan_entry_32
+    push esi                        # the boot_params pointer
+    call _trojan_entry_32
+
+    // Unreachable here.
+halt:
+    hlt
+    jmp halt
+
+// A small stack for the 32-bit code.
+.section ".stack", "aw"
+.align 8
+    .space 0x1000
+__stack_top:

framework/libs/boot-trojan/trojan/src/arch/setup64.S

@@ -0,0 +1,16 @@
+.section ".setup", "ax"
+.code64
+.org 0x200
+// start_of_setup64 should be at start_of_setup32 + 0x200
+.global start_of_setup64
+start_of_setup64:
+
+.extern _trojan_entry_64
+    lea rax, [rip + _trojan_entry_64]
+    push rsi
+    call rax
+    
+    // Unreachable here.
+halt:
+    hlt
+    jmp halt

framework/libs/boot-trojan/trojan/src/arch/x86_64.linker.ld

@@ -0,0 +1,39 @@
+ENTRY(start_of_setup64)
+OUTPUT_ARCH(i386:x86-64)
+OUTPUT_FORMAT(elf64-x86-64)
+
+SETUP32_LMA = 0x100000;
+BOOTSECT_SIZE = 0x1000;
+BOOTSECT_START = SETUP32_LMA - BOOTSECT_SIZE;
+
+LEGACY_STACK_SIZE = 0x1000;
+
+SECTIONS
+{
+    . = BOOTSECT_START;
+    .header                 : { KEEP(*(.header)) }
+
+    . = SETUP32_LMA;
+    .setup                  : { KEEP(*(.setup)) }
+
+    .text                   : { *(.text .text.*) }
+    .rodata                 : { *(.rodata .rodata.*) }
+
+    .data                   : { *(.data .data.*) }
+    .bss                    : {
+        PROVIDE(__bss_start = .);
+        *(.bss .bss.*) *(COMMON)
+        PROVIDE(__bss_end = .);
+    }
+
+    .eh_frame               : {
+        *(.eh_frame .eh_frame.*)
+    }
+    .eh_frame_hdr           : {
+        *(.eh_frame_hdr .eh_frame_hdr.*)
+    }
+
+    .rela.dyn               : {
+        *(.rela.dyn .rela.dyn.*)
+    }
+}

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

@@ -0,0 +1,18 @@
+use core::arch::{asm, global_asm};
+
+global_asm!(include_str!("header.S"));
+
+global_asm!(include_str!("setup64.S"));
+
+#[no_mangle]
+extern "cdecl" fn _trojan_entry_64(boot_params_ptr: u64) -> ! {
+    crate::trojan_entry(boot_params_ptr as u32);
+}
+
+pub unsafe fn call_aster_entrypoint(entrypoint: u64, boot_params_ptr: u64) -> ! {
+    asm!("mov rsi, {}", in(reg) boot_params_ptr as u64);
+    asm!("mov rax, {}", in(reg) entrypoint as u64);
+    asm!("jmp rax");
+
+    unreachable!();
+}

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

@@ -1,26 +1,12 @@
 #![no_std]
 #![no_main]
 
+mod arch;
 mod boot_params;
 mod console;
 mod loader;
 
-use core::arch::{asm, global_asm};
-
-global_asm!(include_str!("header.S"));
-
-global_asm!(include_str!("setup.S"));
-
-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!();
-}
-
-#[no_mangle]
-pub extern "cdecl" fn _rust_setup_entry(boot_params_ptr: u32) -> ! {
+fn trojan_entry(boot_params_ptr: u32) -> ! {
     // Safety: this init function is only called once.
     unsafe { console::init() };
     println!("[setup] boot_params_ptr: {:#x}", boot_params_ptr);
@@ -36,7 +22,7 @@ pub extern "cdecl" fn _rust_setup_entry(boot_params_ptr: u32) -> ! {
     println!("[setup] entrypoint: {:#x}", entrypoint);
 
     // Safety: the entrypoint and the ptr is valid.
-    unsafe { call_aster_entrypoint(entrypoint, boot_params_ptr) };
+    unsafe { arch::call_aster_entrypoint(entrypoint.into(), boot_params_ptr.into()) };
 }
 
 #[panic_handler]

framework/libs/boot-trojan/trojan/src/setup.S

@@ -1,39 +0,0 @@
-// 32-bit setup code starts here, and will be loaded at CODE32_START.
-.section ".setup", "ax"
-.code32
-.global start_of_setup32
-start_of_setup32:
-        mov eax, offset stack_bottom
-        mov esp, eax
-        mov eax, offset halt
-        push eax                        # the return address
-        mov ebp, esp
-        add ebp, -4
-        push ebp
-        mov ebp, esp
-
-        .extern _rust_setup_entry
-        push esi                        # the boot_params pointer
-        call _rust_setup_entry
-
-        // Unreachable here.
-halt:
-        hlt
-        jmp halt
-
-.code64
-.global start_of_setup64
-.org 0x200
-start_of_setup64:
-        // Unreachable here.
-halt64:
-        hlt
-        jmp halt64
-
-// A small stack for the legacy setup code.
-.section ".stack", "aw"
-SETUP_STACK_SIZE = 0x1000
-        .align 16
-        stack_top:
-        .skip SETUP_STACK_SIZE
-        stack_bottom:

runner/src/machine/qemu_grub_efi.rs

@@ -1,4 +1,4 @@
-use aster_boot_trojan_builder::{build_linux_setup_header_from_trojan, make_bzimage};
+use aster_boot_trojan_builder::make_bzimage;
 
 use std::{
     fs,
@@ -91,16 +91,17 @@ pub fn create_bootdev_image(
 
     let target_path = match protocol {
         BootProtocol::Linux => {
-            let trojan_install_dir = Path::new("target/-boot-trojan");
-            build_linux_setup_header_from_trojan(
-                Path::new("framework/libs/boot-trojan/trojan"),
-                trojan_install_dir,
-            )
-            .unwrap();
-            let header_path = trojan_install_dir.join("bin").join("aster-boot-trojan");
+            let trojan_src = Path::new("framework/libs/boot-trojan/trojan");
+            let trojan_out = Path::new("target/aster-boot-trojan");
             // Make the `bzImage`-compatible kernel image and place it in the boot directory.
             let target_path = iso_root.join("boot").join("asterinaz");
-            make_bzimage(&target_path, &aster_path.as_path(), &header_path.as_path()).unwrap();
+            println!("[aster-runner] Building bzImage.");
+            make_bzimage(
+                &target_path,
+                &aster_path.as_path(),
+                &trojan_src,
+                &trojan_out,
+            );
             target_path
         }
         BootProtocol::Multiboot | BootProtocol::Multiboot2 => {