implement UEFI PXE

Author: Freax13

src/lib.rs

@@ -78,6 +78,7 @@ use std::{
 mod fat;
 mod gpt;
 mod mbr;
+mod pxe;
 
 const KERNEL_FILE_NAME: &str = "kernel-x86_64";
 
@@ -121,3 +122,12 @@ pub fn create_bios_disk_image(
 
     Ok(())
 }
+
+pub fn create_uefi_pxe_tftp_folder(kernel_binary: &Path, out_path: &Path) -> anyhow::Result<()> {
+    let bootloader_path = Path::new(env!("UEFI_BOOTLOADER_PATH"));
+
+    pxe::create_uefi_tftp_folder(bootloader_path, kernel_binary, out_path)
+        .context("failed to create UEFI PXE tftp folder")?;
+
+    Ok(())
+}

src/pxe.rs

@@ -0,0 +1,32 @@
+use std::path::Path;
+
+use anyhow::Context;
+
+pub fn create_uefi_tftp_folder(
+    bootloader_path: &Path,
+    kernel_binary: &Path,
+    out_path: &Path,
+) -> anyhow::Result<()> {
+    std::fs::create_dir_all(out_path)
+        .with_context(|| format!("failed to create out dir at {}", out_path.display()))?;
+
+    let to = out_path.join("bootloader");
+    std::fs::copy(bootloader_path, &to).with_context(|| {
+        format!(
+            "failed to copy bootloader from {} to {}",
+            bootloader_path.display(),
+            to.display()
+        )
+    })?;
+
+    let to = out_path.join("kernel-x86_64");
+    std::fs::copy(kernel_binary, &to).with_context(|| {
+        format!(
+            "failed to copy kernel from {} to {}",
+            kernel_binary.display(),
+            to.display()
+        )
+    })?;
+
+    Ok(())
+}

tests/runner/src/lib.rs

@@ -11,6 +11,12 @@ const QEMU_ARGS: &[&str] = &[
 ];
 
 pub fn run_test_kernel(kernel_binary_path: &str) {
+    run_test_kernel_on_uefi(kernel_binary_path);
+    run_test_kernel_on_uefi_pxe(kernel_binary_path);
+    // TODO: run tests with BIOS bootloader too
+}
+
+pub fn run_test_kernel_on_uefi(kernel_binary_path: &str) {
     let kernel_path = Path::new(kernel_binary_path);
     let out_fat_path = kernel_path.with_extension("fat");
     bootloader::create_boot_partition(kernel_path, &out_fat_path).unwrap();
@@ -19,8 +25,6 @@ pub fn run_test_kernel(kernel_binary_path: &str) {
     let out_mbr_path = kernel_path.with_extension("mbr");
     bootloader::create_bios_disk_image(&out_fat_path, &out_mbr_path).unwrap();
 
-    // TODO: run tests with BIOS bootloader too
-
     let mut run_cmd = Command::new("qemu-system-x86_64");
     run_cmd
         .arg("-drive")
@@ -42,3 +46,40 @@ pub fn run_test_kernel(kernel_binary_path: &str) {
         other => panic!("Test failed with unexpected exit code `{:?}`", other),
     }
 }
+
+pub fn run_test_kernel_on_uefi_pxe(kernel_binary_path: &str) {
+    let kernel_path = Path::new(kernel_binary_path);
+    let out_tftp_path = kernel_path.with_extension(".tftp");
+
+    bootloader::create_uefi_pxe_tftp_folder(kernel_path, &out_tftp_path).unwrap();
+
+    let out_fat_path = kernel_path.with_extension("fat");
+    bootloader::create_boot_partition(kernel_path, &out_fat_path).unwrap();
+    let out_gpt_path = kernel_path.with_extension("gpt");
+    bootloader::create_uefi_disk_image(&out_fat_path, &out_gpt_path).unwrap();
+    let out_mbr_path = kernel_path.with_extension("mbr");
+    bootloader::create_bios_disk_image(&out_fat_path, &out_mbr_path).unwrap();
+
+    let mut run_cmd = Command::new("qemu-system-x86_64");
+    run_cmd.arg("-netdev").arg(format!(
+        "user,id=net0,net=192.168.17.0/24,tftp={},bootfile=bootloader,id=net0",
+        out_tftp_path.display()
+    ));
+    run_cmd.arg("-device").arg("virtio-net-pci,netdev=net0");
+    run_cmd.args(QEMU_ARGS);
+    run_cmd.arg("-bios").arg(ovmf_prebuilt::ovmf_pure_efi());
+
+    let child_output = run_cmd.output().unwrap();
+    strip_ansi_escapes::Writer::new(std::io::stderr())
+        .write_all(&child_output.stderr)
+        .unwrap();
+    strip_ansi_escapes::Writer::new(std::io::stderr())
+        .write_all(&child_output.stdout)
+        .unwrap();
+
+    match child_output.status.code() {
+        Some(33) => {}                     // success
+        Some(35) => panic!("Test failed"), // success
+        other => panic!("Test failed with unexpected exit code `{:?}`", other),
+    }
+}

uefi/src/main.rs

@@ -17,9 +17,13 @@ use uefi::{
             file::{File, FileAttribute, FileInfo, FileMode},
             fs::SimpleFileSystem,
         },
+        network::{
+            pxe::{BaseCode, DhcpV4Packet},
+            IpAddress,
+        },
     },
     table::boot::{AllocateType, MemoryDescriptor, MemoryType},
-    CStr16,
+    CStr16, CStr8,
 };
 use x86_64::{
     structures::paging::{FrameAllocator, OffsetPageTable, PageTable, PhysFrame, Size4KiB},
@@ -122,6 +126,16 @@ fn main_inner(image: Handle, mut st: SystemTable<Boot>) -> Status {
 }
 
 fn load_kernel(image: Handle, st: &SystemTable<Boot>) -> Kernel<'static> {
+    let kernel_slice = load_kernel_file(image, st).expect("couldn't find kernel");
+    Kernel::parse(kernel_slice)
+}
+
+fn load_kernel_file(image: Handle, st: &SystemTable<Boot>) -> Option<&'static mut [u8]> {
+    load_kernel_file_from_disk(image, st)
+        .or_else(|| load_kernel_file_from_tftp_boot_server(image, st))
+}
+
+fn load_kernel_file_from_disk(image: Handle, st: &SystemTable<Boot>) -> Option<&'static mut [u8]> {
     let file_system_raw = {
         let ref this = st.boot_services();
         let loaded_image = this
@@ -138,7 +152,7 @@ fn load_kernel(image: Handle, st: &SystemTable<Boot>) -> Kernel<'static> {
 
         let device_handle = this
             .locate_device_path::<SimpleFileSystem>(&mut device_path)
-            .expect("Failed to locate `SimpleFileSystem` protocol on device path");
+            .ok()?;
 
         this.handle_protocol::<SimpleFileSystem>(device_handle)
     }
@@ -172,7 +186,64 @@ fn load_kernel(image: Handle, st: &SystemTable<Boot>) -> Kernel<'static> {
     let kernel_slice = unsafe { slice::from_raw_parts_mut(kernel_ptr, kernel_size) };
     kernel_file.read(kernel_slice).unwrap();
 
-    Kernel::parse(kernel_slice)
+    Some(kernel_slice)
+}
+
+fn load_kernel_file_from_tftp_boot_server(
+    image: Handle,
+    st: &SystemTable<Boot>,
+) -> Option<&'static mut [u8]> {
+    let ref this = st.boot_services();
+
+    let file_system_raw = {
+        let ref this = st.boot_services();
+        let loaded_image = this
+            .handle_protocol::<LoadedImage>(image)
+            .expect("Failed to retrieve `LoadedImage` protocol from handle");
+        let loaded_image = unsafe { &*loaded_image.get() };
+
+        let device_handle = loaded_image.device();
+
+        let device_path = this
+            .handle_protocol::<DevicePath>(device_handle)
+            .expect("Failed to retrieve `DevicePath` protocol from image's device handle");
+        let mut device_path = unsafe { &*device_path.get() };
+
+        let device_handle = this
+            .locate_device_path::<BaseCode>(&mut device_path)
+            .expect("Failed to locate `BaseCode` protocol on device path");
+
+        this.handle_protocol::<BaseCode>(device_handle)
+    }
+    .unwrap();
+    let base_code = unsafe { &mut *file_system_raw.get() };
+
+    let mode = base_code.mode();
+    assert!(mode.dhcp_ack_received);
+    let dhcpv4: &DhcpV4Packet = mode.dhcp_ack.as_ref();
+    let server_ip = IpAddress::new_v4(dhcpv4.bootp_si_addr);
+
+    let filename = CStr8::from_bytes_with_nul(b"kernel-x86_64\0").unwrap();
+    let file_size = base_code.tftp_get_file_size(&server_ip, filename).unwrap();
+
+    let kernel_size = usize::try_from(file_size).unwrap();
+
+    let kernel_ptr = st
+        .boot_services()
+        .allocate_pages(
+            AllocateType::AnyPages,
+            MemoryType::LOADER_DATA,
+            ((kernel_size - 1) / 4096) + 1,
+        )
+        .unwrap() as *mut u8;
+    unsafe { ptr::write_bytes(kernel_ptr, 0, kernel_size) };
+    let kernel_slice = unsafe { slice::from_raw_parts_mut(kernel_ptr, kernel_size) };
+
+    base_code
+        .tftp_read_file(&server_ip, filename, Some(kernel_slice))
+        .unwrap();
+
+    Some(kernel_slice)
 }
 
 /// Creates page table abstraction types for both the bootloader and kernel page tables.