Prevent boot from failing early via panic, and allow net boot to attempt boot

Author: jasoncouture

uefi/src/main.rs

@@ -13,8 +13,8 @@ use uefi::{
     prelude::{entry, Boot, Handle, Status, SystemTable},
     proto::{
         console::gop::{GraphicsOutput, PixelFormat},
-        device_path::DevicePath,
-        loaded_image::LoadedImage,
+        device_path::{self, DevicePath},
+        loaded_image::{self, LoadedImage},
         media::{
             file::{File, FileAttribute, FileInfo, FileMode},
             fs::SimpleFileSystem,
@@ -122,7 +122,7 @@ fn main_inner(image: Handle, mut st: SystemTable<Boot>) -> Status {
             rsdp.map(|entry| PhysAddr::new(entry.address as u64))
         },
         ramdisk_addr: None,
-        ramdisk_len: 0
+        ramdisk_len: 0,
     };
 
     bootloader_x86_64_common::load_and_switch_to_kernel(
@@ -144,59 +144,83 @@ fn load_kernel_file(image: Handle, st: &SystemTable<Boot>) -> Option<&'static mu
         .or_else(|| load_kernel_file_from_tftp_boot_server(image, st))
 }
 
-fn load_file_from_disk(name: &str, image: Handle, st: &SystemTable<Boot>) -> Option<&'static mut [u8]>{ 
+fn load_file_from_disk(
+    name: &str,
+    image: Handle,
+    st: &SystemTable<Boot>,
+) -> Option<&'static mut [u8]> {
     let file_system_raw = {
         let this = st.boot_services();
-        let loaded_image = this
-            .open_protocol::<LoadedImage>(
-                OpenProtocolParams {
-                    handle: image,
-                    agent: image,
-                    controller: None,
-                },
-                OpenProtocolAttributes::Exclusive,
-            )
-            .expect("Failed to retrieve `LoadedImage` protocol from handle");
+        let loaded_image = this.open_protocol::<LoadedImage>(
+            OpenProtocolParams {
+                handle: image,
+                agent: image,
+                controller: None,
+            },
+            OpenProtocolAttributes::Exclusive,
+        );
+
+        if loaded_image.is_err() {
+            log::error!("Failed to open protocol SimpleFileSystem while attempting to load {name}");
+            return None;
+        }
+        let loaded_image = loaded_image.unwrap();
         let loaded_image = unsafe { &*loaded_image.interface.get() };
 
         let device_handle = loaded_image.device();
 
-        let device_path = this
-            .open_protocol::<DevicePath>(
-                OpenProtocolParams {
-                    handle: device_handle,
-                    agent: image,
-                    controller: None,
-                },
-                OpenProtocolAttributes::Exclusive,
-            )
-            .expect("Failed to retrieve `DevicePath` protocol from image's device handle");
+        let device_path = this.open_protocol::<DevicePath>(
+            OpenProtocolParams {
+                handle: device_handle,
+                agent: image,
+                controller: None,
+            },
+            OpenProtocolAttributes::Exclusive,
+        );
+        if device_path.is_err() {
+            log::error!("Failed to open protocol DevicePath while attempting to load {name}");
+            return None;
+        }
+        let device_path = device_path.unwrap();
         let mut device_path = unsafe { &*device_path.interface.get() };
 
-        let fs_handle = this
-            .locate_device_path::<SimpleFileSystem>(&mut device_path)
-            .ok()?;
+        let fs_handle = this.locate_device_path::<SimpleFileSystem>(&mut device_path);
+        if fs_handle.is_err() {
+            log::error!("Failed to open device path while attempting to load {name}");
+            return None;
+        }
+
+        let fs_handle = fs_handle.unwrap();
 
-        this.open_protocol::<SimpleFileSystem>(
+        let opened_handle = this.open_protocol::<SimpleFileSystem>(
             OpenProtocolParams {
                 handle: fs_handle,
                 agent: image,
                 controller: None,
             },
             OpenProtocolAttributes::Exclusive,
-        )
+        );
+
+        if opened_handle.is_err() {
+            log::error!("Failed to open file system while attempting to load {name}");
+            return None;
+        }
+        Some(opened_handle.unwrap())
+    };
+
+    if file_system_raw.is_none() {
+        return None();
     }
-    .unwrap();
+    let file_system_raw = file_system_raw.unwrap();
     let file_system = unsafe { &mut *file_system_raw.interface.get() };
 
     let mut root = file_system.open_volume().unwrap();
-    if(name.len() > 255) {
+    if (name.len() > 255) {
         panic!("File name {}, exceeds maximum length!", name);
     }
     let mut buf = [0; 512];
     let filename = CStr16::from_str_with_buf(name, &mut buf).unwrap();
-    let kernel_file_handle_result = root
-        .open(filename, FileMode::Read, FileAttribute::empty());
+    let kernel_file_handle_result = root.open(filename, FileMode::Read, FileAttribute::empty());
 
     if kernel_file_handle_result.is_err() {
         return None;
@@ -229,7 +253,7 @@ fn load_file_from_disk(name: &str, image: Handle, st: &SystemTable<Boot>) -> Opt
 }
 
 fn load_kernel_file_from_disk(image: Handle, st: &SystemTable<Boot>) -> Option<&'static mut [u8]> {
-    Some(load_file_from_disk(KERNEL_FILENAME, image, st).expect("Failed to load kernel (expected file named `kernel-x86_64`)"))
+    load_file_from_disk(KERNEL_FILENAME, image, st)
 }
 
 /// Try to load a kernel from a TFTP boot server.