Update test-runner

Author: nicholasbishop

uefi-test-runner/src/bin/shell_launcher.rs

@@ -13,6 +13,7 @@ extern crate alloc;
 
 use alloc::vec::Vec;
 use log::info;
+use uefi::boot;
 use uefi::prelude::*;
 use uefi::proto::device_path::build::{self, DevicePathBuilder};
 use uefi::proto::device_path::{DevicePath, DeviceSubType, DeviceType, LoadedImageDevicePath};
@@ -21,13 +22,10 @@ use uefi::table::boot::LoadImageSource;
 
 /// Get the device path of the shell app. This is the same as the
 /// currently-loaded image's device path, but with the file path part changed.
-fn get_shell_app_device_path<'a>(
-    boot_services: &BootServices,
-    storage: &'a mut Vec<u8>,
-) -> &'a DevicePath {
-    let loaded_image_device_path = boot_services
-        .open_protocol_exclusive::<LoadedImageDevicePath>(boot_services.image_handle())
-        .expect("failed to open LoadedImageDevicePath protocol");
+fn get_shell_app_device_path(storage: &mut Vec<u8>) -> &DevicePath {
+    let loaded_image_device_path =
+        boot::open_protocol_exclusive::<LoadedImageDevicePath>(boot::image_handle())
+            .expect("failed to open LoadedImageDevicePath protocol");
 
     let mut builder = DevicePathBuilder::with_vec(storage);
     for node in loaded_image_device_path.node_iter() {
@@ -47,26 +45,23 @@ fn get_shell_app_device_path<'a>(
 #[entry]
 fn efi_main(image: Handle, st: SystemTable<Boot>) -> Status {
     uefi::helpers::init().unwrap();
-    let boot_services = st.boot_services();
 
     let mut storage = Vec::new();
-    let shell_image_path = get_shell_app_device_path(boot_services, &mut storage);
+    let shell_image_path = get_shell_app_device_path(&mut storage);
 
     // Load the shell app.
-    let shell_image_handle = boot_services
-        .load_image(
-            image,
-            LoadImageSource::FromDevicePath {
-                device_path: shell_image_path,
-                from_boot_manager: false,
-            },
-        )
-        .expect("failed to load shell app");
+    let shell_image_handle = boot::load_image(
+        boot::image_handle(),
+        LoadImageSource::FromDevicePath {
+            device_path: shell_image_path,
+            from_boot_manager: false,
+        },
+    )
+    .expect("failed to load shell app");
 
     // Set the command line passed to the shell app so that it will run the
     // test-runner app. This automatically turns off the five-second delay.
-    let mut shell_loaded_image = boot_services
-        .open_protocol_exclusive::<LoadedImage>(shell_image_handle)
+    let mut shell_loaded_image = boot::open_protocol_exclusive::<LoadedImage>(shell_image_handle)
         .expect("failed to open LoadedImage protocol");
     let load_options = cstr16!(r"shell.efi test_runner.efi arg1 arg2");
     unsafe {
@@ -77,9 +72,7 @@ fn efi_main(image: Handle, st: SystemTable<Boot>) -> Status {
     }
 
     info!("launching the shell app");
-    boot_services
-        .start_image(shell_image_handle)
-        .expect("failed to launch the shell app");
+    boot::start_image(shell_image_handle).expect("failed to launch the shell app");
 
     Status::SUCCESS
 }

uefi-test-runner/src/boot/memory.rs

@@ -1,25 +1,24 @@
-use uefi::table::boot::{AllocateType, BootServices, MemoryMap, MemoryMapMut, MemoryType};
+use uefi::boot;
+use uefi::table::boot::{AllocateType, MemoryMap, MemoryMapMut, MemoryType};
 
 use alloc::vec::Vec;
 
-pub fn test(bt: &BootServices) {
+pub fn test() {
     info!("Testing memory functions");
 
-    allocate_pages(bt);
+    allocate_pages();
     vec_alloc();
     alloc_alignment();
 
-    memory_map(bt);
+    memory_map();
 }
 
-fn allocate_pages(bt: &BootServices) {
+fn allocate_pages() {
     info!("Allocating some pages of memory");
 
     let ty = AllocateType::AnyPages;
     let mem_ty = MemoryType::LOADER_DATA;
-    let pgs = bt
-        .allocate_pages(ty, mem_ty, 1)
-        .expect("Failed to allocate a page of memory");
+    let pgs = boot::allocate_pages(ty, mem_ty, 1).expect("Failed to allocate a page of memory");
 
     assert_eq!(pgs % 4096, 0, "Page pointer is not page-aligned");
 
@@ -31,7 +30,7 @@ fn allocate_pages(bt: &BootServices) {
     buf[4095] = 0x23;
 
     // Clean up to avoid memory leaks.
-    unsafe { bt.free_pages(pgs, 1) }.unwrap();
+    unsafe { boot::free_pages(pgs, 1) }.unwrap();
 }
 
 // Simple test to ensure our custom allocator works with the `alloc` crate.
@@ -60,15 +59,14 @@ fn alloc_alignment() {
     assert_eq!(value.as_ptr() as usize % 0x100, 0, "Wrong alignment");
 }
 
-fn memory_map(bt: &BootServices) {
+fn memory_map() {
     info!("Testing memory map functions");
 
     // Ensure that the memory map is freed after each iteration (on drop).
     // Otherwise, we will have an OOM.
     for _ in 0..200000 {
-        let mut memory_map = bt
-            .memory_map(MemoryType::LOADER_DATA)
-            .expect("Failed to retrieve UEFI memory map");
+        let mut memory_map =
+            boot::memory_map(MemoryType::LOADER_DATA).expect("Failed to retrieve UEFI memory map");
 
         memory_map.sort();
 

uefi-test-runner/src/boot/misc.rs

@@ -4,51 +4,49 @@ use core::ptr::{self, NonNull};
 use core::mem;
 use uefi::proto::unsafe_protocol;
 use uefi::table::boot::{
-    BootServices, EventType, MemoryType, OpenProtocolAttributes, OpenProtocolParams, SearchType,
-    TimerTrigger, Tpl,
+    EventType, MemoryType, OpenProtocolAttributes, OpenProtocolParams, SearchType, TimerTrigger,
+    Tpl,
 };
-use uefi::table::{Boot, SystemTable};
+use uefi::{boot, system};
 use uefi::{guid, Event, Guid, Identify};
 
-pub fn test(st: &SystemTable<Boot>) {
-    let bt = st.boot_services();
+pub fn test() {
     info!("Testing timer...");
-    test_timer(bt);
+    test_timer();
     info!("Testing events...");
-    test_event_callback(bt);
-    test_callback_with_ctx(bt);
+    test_event_callback();
+    test_callback_with_ctx();
     info!("Testing watchdog...");
-    test_watchdog(bt);
+    test_watchdog();
     info!("Testing protocol handler services...");
-    test_register_protocol_notify(bt);
-    test_install_protocol_interface(bt);
-    test_reinstall_protocol_interface(bt);
-    test_uninstall_protocol_interface(bt);
-    test_install_configuration_table(st);
+    test_register_protocol_notify();
+    test_install_protocol_interface();
+    test_reinstall_protocol_interface();
+    test_uninstall_protocol_interface();
+    test_install_configuration_table();
 }
 
-fn test_timer(bt: &BootServices) {
-    let timer_event = unsafe { bt.create_event(EventType::TIMER, Tpl::APPLICATION, None, None) }
+fn test_timer() {
+    let timer_event = unsafe { boot::create_event(EventType::TIMER, Tpl::APPLICATION, None, None) }
         .expect("Failed to create TIMER event");
     let mut events = unsafe { [timer_event.unsafe_clone()] };
-    bt.set_timer(&timer_event, TimerTrigger::Relative(5_0 /*00 ns */))
+    boot::set_timer(&timer_event, TimerTrigger::Relative(5_0 /*00 ns */))
         .expect("Failed to set timer");
-    bt.wait_for_event(&mut events)
-        .expect("Wait for event failed");
+    boot::wait_for_event(&mut events).expect("Wait for event failed");
 }
 
-fn test_event_callback(bt: &BootServices) {
+fn test_event_callback() {
     extern "efiapi" fn callback(_event: Event, _ctx: Option<NonNull<c_void>>) {
         info!("Inside the event callback");
     }
 
     let event =
-        unsafe { bt.create_event(EventType::NOTIFY_WAIT, Tpl::CALLBACK, Some(callback), None) }
+        unsafe { boot::create_event(EventType::NOTIFY_WAIT, Tpl::CALLBACK, Some(callback), None) }
             .expect("Failed to create custom event");
-    bt.check_event(event).expect("Failed to check event");
+    boot::check_event(event).expect("Failed to check event");
 }
 
-fn test_callback_with_ctx(bt: &BootServices) {
+fn test_callback_with_ctx() {
     let mut data = 123u32;
 
     extern "efiapi" fn callback(_event: Event, ctx: Option<NonNull<c_void>>) {
@@ -65,7 +63,7 @@ fn test_callback_with_ctx(bt: &BootServices) {
     let ctx = NonNull::new(ctx.cast::<c_void>()).unwrap();
 
     let event = unsafe {
-        bt.create_event(
+        boot::create_event(
             EventType::NOTIFY_WAIT,
             Tpl::CALLBACK,
             Some(callback),
@@ -74,16 +72,15 @@ fn test_callback_with_ctx(bt: &BootServices) {
         .expect("Failed to create event with context")
     };
 
-    bt.check_event(event).expect("Failed to check event");
+    boot::check_event(event).expect("Failed to check event");
 
     // Check that `data` was updated inside the event callback.
     assert_eq!(data, 456);
 }
 
-fn test_watchdog(bt: &BootServices) {
+fn test_watchdog() {
     // Disable the UEFI watchdog timer
-    bt.set_watchdog_timer(0, 0x10000, None)
-        .expect("Could not set watchdog timer");
+    boot::set_watchdog_timer(0, 0x10000, None).expect("Could not set watchdog timer");
 }
 
 /// Dummy protocol for tests
@@ -96,10 +93,10 @@ unsafe extern "efiapi" fn _test_notify(_event: Event, _context: Option<NonNull<c
     info!("Protocol was (re)installed and this function notified.")
 }
 
-fn test_register_protocol_notify(bt: &BootServices) {
+fn test_register_protocol_notify() {
     let protocol = &TestProtocol::GUID;
     let event = unsafe {
-        bt.create_event(
+        boot::create_event(
             EventType::NOTIFY_SIGNAL,
             Tpl::NOTIFY,
             Some(_test_notify),
@@ -108,41 +105,37 @@ fn test_register_protocol_notify(bt: &BootServices) {
         .expect("Failed to create an event")
     };
 
-    bt.register_protocol_notify(protocol, event)
-        .expect("Failed to register protocol notify fn");
+    boot::register_protocol_notify(protocol, event).expect("Failed to register protocol notify fn");
 }
 
-fn test_install_protocol_interface(bt: &BootServices) {
+fn test_install_protocol_interface() {
     info!("Installing TestProtocol");
 
-    let alloc: *mut TestProtocol = bt
-        .allocate_pool(
-            MemoryType::BOOT_SERVICES_DATA,
-            mem::size_of::<TestProtocol>(),
-        )
-        .unwrap()
-        .cast()
-        .as_ptr();
+    let alloc: *mut TestProtocol = boot::allocate_pool(
+        MemoryType::BOOT_SERVICES_DATA,
+        mem::size_of::<TestProtocol>(),
+    )
+    .unwrap()
+    .as_ptr()
+    .cast();
     unsafe { alloc.write(TestProtocol { data: 123 }) };
 
     let _ = unsafe {
-        bt.install_protocol_interface(None, &TestProtocol::GUID, alloc.cast())
+        boot::install_protocol_interface(None, &TestProtocol::GUID, alloc.cast())
             .expect("Failed to install protocol interface")
     };
 
-    let _ = bt
-        .locate_handle_buffer(SearchType::from_proto::<TestProtocol>())
+    let _ = boot::locate_handle_buffer(SearchType::from_proto::<TestProtocol>())
         .expect("Failed to find protocol after it was installed");
 }
 
-fn test_reinstall_protocol_interface(bt: &BootServices) {
+fn test_reinstall_protocol_interface() {
     info!("Reinstalling TestProtocol");
-    let handle = bt
-        .locate_handle_buffer(SearchType::from_proto::<TestProtocol>())
+    let handle = boot::locate_handle_buffer(SearchType::from_proto::<TestProtocol>())
         .expect("Failed to find protocol to uninstall")[0];
 
     unsafe {
-        let _ = bt.reinstall_protocol_interface(
+        let _ = boot::reinstall_protocol_interface(
             handle,
             &TestProtocol::GUID,
             ptr::null_mut(),
@@ -151,61 +144,57 @@ fn test_reinstall_protocol_interface(bt: &BootServices) {
     }
 }
 
-fn test_uninstall_protocol_interface(bt: &BootServices) {
+fn test_uninstall_protocol_interface() {
     info!("Uninstalling TestProtocol");
 
-    let handle = bt
-        .locate_handle_buffer(SearchType::from_proto::<TestProtocol>())
+    let handle = boot::locate_handle_buffer(SearchType::from_proto::<TestProtocol>())
         .expect("Failed to find protocol to uninstall")[0];
 
     unsafe {
         // Uninstalling a protocol interface requires knowing the interface
         // pointer. Open the protocol to get that pointer, making sure to drop
         // the `ScopedProtocol` _before_ uninstalling the protocol interface.
         let interface_ptr: *mut TestProtocol = {
-            let mut sp = bt
-                .open_protocol::<TestProtocol>(
-                    OpenProtocolParams {
-                        handle,
-                        agent: bt.image_handle(),
-                        controller: None,
-                    },
-                    OpenProtocolAttributes::GetProtocol,
-                )
-                .unwrap();
+            let mut sp = boot::open_protocol::<TestProtocol>(
+                OpenProtocolParams {
+                    handle,
+                    agent: boot::image_handle(),
+                    controller: None,
+                },
+                OpenProtocolAttributes::GetProtocol,
+            )
+            .unwrap();
             assert_eq!(sp.data, 123);
             &mut *sp
         };
 
-        bt.uninstall_protocol_interface(handle, &TestProtocol::GUID, interface_ptr.cast())
+        boot::uninstall_protocol_interface(handle, &TestProtocol::GUID, interface_ptr.cast())
             .expect("Failed to uninstall protocol interface");
 
-        bt.free_pool(interface_ptr.cast()).unwrap();
+        boot::free_pool(interface_ptr.cast()).unwrap();
     }
 }
 
-fn test_install_configuration_table(st: &SystemTable<Boot>) {
-    let config = st
-        .boot_services()
-        .allocate_pool(MemoryType::ACPI_RECLAIM, 1)
+fn test_install_configuration_table() {
+    let config = boot::allocate_pool(MemoryType::ACPI_RECLAIM, 1)
         .expect("Failed to allocate config table")
         .as_ptr();
     unsafe { config.write(42) };
 
-    let count = st.config_table().len();
+    let count = system::with_config_table(|t| t.len());
     const ID: Guid = guid!("3bdb3089-5662-42df-840e-3922ed6467c9");
 
     unsafe {
-        st.boot_services()
-            .install_configuration_table(&ID, config.cast())
+        boot::install_configuration_table(&ID, config.cast())
             .expect("Failed to install configuration table");
     }
 
-    assert_eq!(count + 1, st.config_table().len());
-    let config_entry = st
-        .config_table()
-        .iter()
-        .find(|ct| ct.guid == ID)
-        .expect("Failed to find test config table");
-    assert_eq!(unsafe { *(config_entry.address as *const u8) }, 42);
+    assert_eq!(count + 1, system::with_config_table(|t| t.len()));
+    let entry_addr = system::with_config_table(|t| {
+        t.iter()
+            .find(|ct| ct.guid == ID)
+            .expect("Failed to find test config table")
+            .address
+    });
+    assert_eq!(unsafe { *(entry_addr as *const u8) }, 42);
 }