Make sure ramdisk loads from the same source as the kernel

Author: jasoncouture

uefi/src/main.rs

@@ -76,17 +76,40 @@ fn main_inner(image: Handle, mut st: SystemTable<Boot>) -> Status {
     unsafe {
         *SYSTEM_TABLE.get() = Some(st.unsafe_clone());
     }
+    let stdout = unsafe { &mut *SYSTEM_TABLE.get() };
+    let stdout = stdout.as_mut().unwrap();
+    let stdout = stdout.stdout();
+    stdout.clear().unwrap();
+    writeln!(stdout, "UEFI bootloader started; trying to load kernel").unwrap();
+
+    let mut boot_mode = BootMode::Disk;
+    let mut kernel = load_kernel(image, &mut st, boot_mode);
+    if kernel.is_none() {
+        writeln!(
+            stdout,
+            "Failed to load kernel via {:?}, trying TFTP",
+            boot_mode
+        )
+        .unwrap();
+        // Try TFTP boot
+        boot_mode = BootMode::Tftp;
+        kernel = load_kernel(image, &mut st, boot_mode);
+    }
+    let kernel = kernel.expect("Failed to load kernel");
+    writeln!(stdout, "Trying to load ramdisk via {:?}", boot_mode).unwrap();
+    // Ramdisk must load from same source, or not at all.
+    let ramdisk = load_ramdisk(image, &mut st, boot_mode);
 
-    st.stdout().clear().unwrap();
     writeln!(
-        st.stdout(),
-        "UEFI bootloader started; trying to load kernel"
+        stdout,
+        "{}",
+        match ramdisk {
+            Some(_) => "Loaded ramdisk",
+            None => "Ramdisk not found.",
+        }
     )
     .unwrap();
 
-    let kernel = load_kernel(image, &st);
-    let ramdisk = load_ramdisk(image, &st);
-
     let framebuffer = init_logger(&st, kernel.config);
 
     // we no longer need the system table for printing panics
@@ -152,16 +175,39 @@ fn main_inner(image: Handle, mut st: SystemTable<Boot>) -> Status {
     );
 }
 
-fn load_ramdisk(image: Handle, st: &SystemTable<Boot>) -> Option<&'static mut [u8]> {
-    load_file_from_disk(RAMDISK_FILENAME, image, st)
-        .or_else(|| load_file_from_tftp_boot_server(RAMDISK_FILENAME, image, st))
+#[derive(Clone, Copy, Debug)]
+pub enum BootMode {
+    Disk,
+    Tftp,
+}
+
+fn load_ramdisk(
+    image: Handle,
+    st: &mut SystemTable<Boot>,
+    boot_mode: BootMode,
+) -> Option<&'static mut [u8]> {
+    load_file_from_boot_method(image, st, RAMDISK_FILENAME, boot_mode)
+}
+
+fn load_kernel(
+    image: Handle,
+    st: &mut SystemTable<Boot>,
+    boot_mode: BootMode,
+) -> Option<Kernel<'static>> {
+    let kernel_slice = load_file_from_boot_method(image, st, KERNEL_FILENAME, boot_mode)?;
+    Some(Kernel::parse(kernel_slice))
 }
 
-fn load_kernel(image: Handle, st: &SystemTable<Boot>) -> Kernel<'static> {
-    let kernel_slice = load_file_from_disk(KERNEL_FILENAME, image, st)
-        .or_else(|| load_file_from_tftp_boot_server(KERNEL_FILENAME, image, st))
-        .expect("couldn't find kernel");
-    Kernel::parse(kernel_slice)
+fn load_file_from_boot_method(
+    image: Handle,
+    st: &mut SystemTable<Boot>,
+    filename: &str,
+    boot_mode: BootMode,
+) -> Option<&'static mut [u8]> {
+    match boot_mode {
+        BootMode::Disk => load_file_from_disk(filename, image, st),
+        BootMode::Tftp => load_file_from_tftp_boot_server(filename, image, st),
+    }
 }
 
 fn open_device_path_protocol(