From e6c83089b6206cecf9d62c36e3a8bbaf5c4ede0c Mon Sep 17 00:00:00 2001
From: Laurin-Luis Lehning <65224843+e820@users.noreply.github.com>
Date: Thu, 14 Jul 2022 15:06:03 +0200
Subject: [PATCH 01/11] Add  and  protocols

---
 src/proto/media/disk.rs | 200 ++++++++++++++++++++++++++++++++++++++++
 src/proto/media/mod.rs  |   1 +
 2 files changed, 201 insertions(+)
 create mode 100644 src/proto/media/disk.rs

diff --git a/src/proto/media/disk.rs b/src/proto/media/disk.rs
new file mode 100644
index 000000000..bc142ef3e
--- /dev/null
+++ b/src/proto/media/disk.rs
@@ -0,0 +1,200 @@
+//! Disk I/O protocols.
+
+use crate::proto::Protocol;
+use crate::{unsafe_guid, Event, Result, Status};
+
+/// The disk I/O protocol.
+///
+/// This protocol is used to abstract the block accesses of the block I/O
+/// protocol to a more general offset-length protocol. Firmware is
+/// reponsible for adding this protocol to any block I/O interface that
+/// appears in the system that does not already have a disk I/O protocol.
+#[repr(C)]
+#[unsafe_guid("ce345171-ba0b-11d2-4f8e-00a0c969723b")]
+#[derive(Protocol)]
+pub struct DiskIo {
+    revision: u64,
+    read_disk: extern "efiapi" fn(
+        this: &DiskIo,
+        media_id: u32,
+        offset: u64,
+        len: usize,
+        buffer: *mut u8,
+    ) -> Status,
+    write_disk: extern "efiapi" fn(
+        this: &mut DiskIo,
+        media_id: u32,
+        offset: u64,
+        len: usize,
+        buffer: *const u8,
+    ) -> Status,
+}
+
+impl DiskIo {
+    /// Reads bytes from the disk device.
+    ///
+    /// # Arguments:
+    /// * `media_id` - ID of the medium to be read.
+    /// * `offset` - Starting byte offset on the logical block I/O device to read from.
+    /// * `buffer` - Pointer to a buffer to read into.
+    ///
+    /// # Errors:
+    /// * `uefi::status::INVALID_PARAMETER` The read request contains device addresses that
+    ///                                     are not valid for the device.
+    /// * `uefi::status::DEVICE_ERROR`      The device reported an error while performing
+    ///                                     the read operation.
+    /// * `uefi::status::NO_MEDIA`          There is no medium in the device.
+    /// * `uefi::status::MEDIA_CHANGED`     `media_id` is not for the current medium.
+    pub fn read_disk(&self, media_id: u32, offset: u64, buffer: &mut [u8]) -> Result {
+        (self.read_disk)(self, media_id, offset, buffer.len(), buffer.as_mut_ptr()).into()
+    }
+
+    /// Writes bytes to the disk device.
+    ///
+    /// # Arguments:
+    /// * `media_id` - ID of the medium to be written.
+    /// * `offset` - Starting byte offset on the logical block I/O device to write to.
+    /// * `buffer` - Pointer to a buffer to write from.
+    ///
+    /// # Errors:
+    /// * `uefi::status::INVALID_PARAMETER` The write request contains device addresses that
+    ///                                     are not valid for the device.
+    /// * `uefi::status::DEVICE_ERROR`      The device reported an error while performing
+    ///                                     the write operation.
+    /// * `uefi::status::NO_MEDIA`          There is no medium in the device.
+    /// * `uefi::status::MEDIA_CHANGED`     `media_id` is not for the current medium.
+    /// * `uefi::status::WRITE_PROTECTED`   The device cannot be written to.
+    pub fn write_disk(&mut self, media_id: u32, offset: u64, buffer: &[u8]) -> Result {
+        (self.write_disk)(self, media_id, offset, buffer.len(), buffer.as_ptr()).into()
+    }
+}
+
+/// Asynchronous transaction token for disk I/O 2 operations.
+#[repr(C)]
+pub struct DiskIo2Token {
+    /// Event to be signalled when an asynchronous disk I/O operation completes.
+    event: Option<Event>,
+    /// Transaction status code.
+    transaction_status: Status,
+}
+
+/// The disk I/O 2 protocol.
+///
+/// This protocol provides an extension to the disk I/O protocol to enable
+/// non-blocking / asynchronous byte-oriented disk operation.
+#[repr(C)]
+#[unsafe_guid("151c8eae-7f2c-472c-549e-9828194f6a88")]
+#[derive(Protocol)]
+pub struct DiskIo2 {
+    revision: u64,
+    cancel: extern "efiapi" fn(this: &mut DiskIo2) -> Status,
+    read_disk_ex: extern "efiapi" fn(
+        this: &DiskIo2,
+        media_id: u32,
+        offset: u64,
+        token: &mut DiskIo2Token,
+        len: usize,
+        buffer: *mut u8,
+    ) -> Status,
+    write_disk_ex: extern "efiapi" fn(
+        this: &mut DiskIo2,
+        media_id: u32,
+        offset: u64,
+        token: &mut DiskIo2Token,
+        len: usize,
+        buffer: *const u8,
+    ) -> Status,
+    flush_disk_ex: extern "efiapi" fn(this: &mut DiskIo2, token: &mut DiskIo2Token) -> Status,
+}
+
+impl DiskIo2 {
+    /// Terminates outstanding asynchronous requests to the device.
+    ///
+    /// # Errors:
+    /// * `uefi::status::DEVICE_ERROR`  The device reported an error while performing
+    ///                                 the cancel operation.
+    pub fn cancel(&mut self) -> Result {
+        (self.cancel)(self).into()
+    }
+
+    /// Reads bytes from the disk device.
+    ///
+    /// # Arguments:
+    /// * `media_id` - ID of the medium to be read from.
+    /// * `offset` - Starting byte offset on the logical block I/O device to read from.
+    /// * `token` - Transaction token for asynchronous read.
+    /// * `buffer` - Buffer to read into.
+    ///
+    /// # Errors:
+    /// * `uefi::status::INVALID_PARAMETER` The read request contains device addresses
+    ///                                     that are not valid for the device.
+    /// * `uefi::status::OUT_OF_RESOURCES`  The request could not be completed due to
+    ///                                     a lack of resources.
+    /// * `uefi::status::MEDIA_CHANGED`     `media_id` is not for the current medium.
+    /// * `uefi::status::NO_MEDIA`          There is no medium in the device.
+    /// * `uefi::status::DEVICE_ERROR`      The device reported an error while performing
+    ///                                     the read operation.
+    pub fn read_disk(
+        &self,
+        media_id: u32,
+        offset: u64,
+        token: &mut DiskIo2Token,
+        buffer: &mut [u8],
+    ) -> Result {
+        (self.read_disk_ex)(
+            self,
+            media_id,
+            offset,
+            token,
+            buffer.len(),
+            buffer.as_mut_ptr(),
+        )
+        .into()
+    }
+
+    /// Writes bytes to the disk device.
+    ///
+    /// # Arguments:
+    /// * `media_id` - ID of the medium to write to.
+    /// * `offset` - Starting byte offset on the logical block I/O device to write to.
+    /// * `token` - Transaction token for asynchronous write.
+    /// * `buffer` - Buffer to write from.
+    ///
+    /// # Errors:
+    /// * `uefi::status::INVALID_PARAMETER` The write request contains device addresses
+    ///                                     that are not valid for the device.
+    /// * `uefi::status::OUT_OF_RESOURCES`  The request could not be completed due to
+    ///                                     a lack of resources.
+    /// * `uefi::status::MEDIA_CHANGED`     `media_id` is not for the current medium.
+    /// * `uefi::status::NO_MEDIA`          There is no medium in the device.
+    /// * `uefi::status::DEVICE_ERROR`      The device reported an error while performing
+    ///                                     the write operation.
+    /// * `uefi::status::WRITE_PROTECTED`   The device cannot be written to.
+    pub fn write_disk(
+        &mut self,
+        media_id: u32,
+        offset: u64,
+        token: &mut DiskIo2Token,
+        buffer: &[u8],
+    ) -> Result {
+        (self.write_disk_ex)(self, media_id, offset, token, buffer.len(), buffer.as_ptr()).into()
+    }
+
+    /// Flushes all modified data to the physical device.
+    ///
+    /// # Arguments:
+    /// * `token` - Transaction token for the asynchronous flush.
+    ///
+    /// # Errors:
+    /// * `uefi::status::OUT_OF_RESOURCES`  The request could not be completed due to
+    ///                                     a lack of resources.
+    /// * `uefi::status::MEDIA_CHANGED`     The medium in the device has changed since
+    ///                                     the last access.
+    /// * `uefi::status::NO_MEDIA`          There is no medium in the device.
+    /// * `uefi::status::DEVICE_ERROR`      The device reported an error while performing
+    ///                                     the flush operation.
+    /// * `uefi::status::WRITE_PROTECTED`   The device cannot be written to.
+    pub fn flush_disk(&mut self, token: &mut DiskIo2Token) -> Result {
+        (self.flush_disk_ex)(self, token).into()
+    }
+}
diff --git a/src/proto/media/mod.rs b/src/proto/media/mod.rs
index 19c13486e..6750875a6 100644
--- a/src/proto/media/mod.rs
+++ b/src/proto/media/mod.rs
@@ -7,5 +7,6 @@
 pub mod file;
 
 pub mod block;
+pub mod disk;
 pub mod fs;
 pub mod partition;

From f80c9b447bea7b861e803d8ac7e792f47e591a2e Mon Sep 17 00:00:00 2001
From: Laurin-Luis Lehning <65224843+e820@users.noreply.github.com>
Date: Thu, 14 Jul 2022 15:10:21 +0200
Subject: [PATCH 02/11] Fix DiskIo and DiskIo2 GUIDs

---
 src/proto/media/disk.rs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/proto/media/disk.rs b/src/proto/media/disk.rs
index bc142ef3e..651c94bfe 100644
--- a/src/proto/media/disk.rs
+++ b/src/proto/media/disk.rs
@@ -10,7 +10,7 @@ use crate::{unsafe_guid, Event, Result, Status};
 /// reponsible for adding this protocol to any block I/O interface that
 /// appears in the system that does not already have a disk I/O protocol.
 #[repr(C)]
-#[unsafe_guid("ce345171-ba0b-11d2-4f8e-00a0c969723b")]
+#[unsafe_guid("ce345171-ba0b-11d2-8e4f-00a0c969723b")]
 #[derive(Protocol)]
 pub struct DiskIo {
     revision: u64,
@@ -83,7 +83,7 @@ pub struct DiskIo2Token {
 /// This protocol provides an extension to the disk I/O protocol to enable
 /// non-blocking / asynchronous byte-oriented disk operation.
 #[repr(C)]
-#[unsafe_guid("151c8eae-7f2c-472c-549e-9828194f6a88")]
+#[unsafe_guid("151c8eae-7f2c-472c-9e54-9828194f6a88")]
 #[derive(Protocol)]
 pub struct DiskIo2 {
     revision: u64,

From 2c9fdeee23a38e367032d45269c879a18e1221eb Mon Sep 17 00:00:00 2001
From: Laurin-Luis Lehning <65224843+e820@users.noreply.github.com>
Date: Thu, 14 Jul 2022 15:35:27 +0200
Subject: [PATCH 03/11] Make DiskIo2Token fields pub

---
 src/proto/media/disk.rs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/proto/media/disk.rs b/src/proto/media/disk.rs
index 651c94bfe..3a25c04e1 100644
--- a/src/proto/media/disk.rs
+++ b/src/proto/media/disk.rs
@@ -73,9 +73,9 @@ impl DiskIo {
 #[repr(C)]
 pub struct DiskIo2Token {
     /// Event to be signalled when an asynchronous disk I/O operation completes.
-    event: Option<Event>,
+    pub event: Option<Event>,
     /// Transaction status code.
-    transaction_status: Status,
+    pub transaction_status: Status,
 }
 
 /// The disk I/O 2 protocol.

From 468c4aff9325a5873622b5bf7232421104d150e9 Mon Sep 17 00:00:00 2001
From: Laurin-Luis Lehning <65224843+e820@users.noreply.github.com>
Date: Sat, 30 Jul 2022 22:52:47 +0200
Subject: [PATCH 04/11] Make DiskIo2 I/O unsafe & add tests

---
 src/proto/media/disk.rs                       |  28 ++---
 .../src/proto/media/known_disk.rs             | 107 +++++++++++++++++-
 2 files changed, 118 insertions(+), 17 deletions(-)

diff --git a/src/proto/media/disk.rs b/src/proto/media/disk.rs
index 3a25c04e1..323407627 100644
--- a/src/proto/media/disk.rs
+++ b/src/proto/media/disk.rs
@@ -123,6 +123,7 @@ impl DiskIo2 {
     /// * `media_id` - ID of the medium to be read from.
     /// * `offset` - Starting byte offset on the logical block I/O device to read from.
     /// * `token` - Transaction token for asynchronous read.
+    /// * `len` - Buffer size.
     /// * `buffer` - Buffer to read into.
     ///
     /// # Errors:
@@ -134,22 +135,15 @@ impl DiskIo2 {
     /// * `uefi::status::NO_MEDIA`          There is no medium in the device.
     /// * `uefi::status::DEVICE_ERROR`      The device reported an error while performing
     ///                                     the read operation.
-    pub fn read_disk(
+    pub unsafe fn read_disk_raw(
         &self,
         media_id: u32,
         offset: u64,
-        token: &mut DiskIo2Token,
-        buffer: &mut [u8],
+        token: *mut DiskIo2Token,
+        len: usize,
+        buffer: *mut u8,
     ) -> Result {
-        (self.read_disk_ex)(
-            self,
-            media_id,
-            offset,
-            token,
-            buffer.len(),
-            buffer.as_mut_ptr(),
-        )
-        .into()
+        (self.read_disk_ex)(self, media_id, offset, &mut *token, len, buffer).into()
     }
 
     /// Writes bytes to the disk device.
@@ -158,6 +152,7 @@ impl DiskIo2 {
     /// * `media_id` - ID of the medium to write to.
     /// * `offset` - Starting byte offset on the logical block I/O device to write to.
     /// * `token` - Transaction token for asynchronous write.
+    /// * `len` - Buffer size.
     /// * `buffer` - Buffer to write from.
     ///
     /// # Errors:
@@ -170,14 +165,15 @@ impl DiskIo2 {
     /// * `uefi::status::DEVICE_ERROR`      The device reported an error while performing
     ///                                     the write operation.
     /// * `uefi::status::WRITE_PROTECTED`   The device cannot be written to.
-    pub fn write_disk(
+    pub unsafe fn write_disk_raw(
         &mut self,
         media_id: u32,
         offset: u64,
-        token: &mut DiskIo2Token,
-        buffer: &[u8],
+        token: *mut DiskIo2Token,
+        len: usize,
+        buffer: *const u8,
     ) -> Result {
-        (self.write_disk_ex)(self, media_id, offset, token, buffer.len(), buffer.as_ptr()).into()
+        (self.write_disk_ex)(self, media_id, offset, &mut *token, len, buffer).into()
     }
 
     /// Flushes all modified data to the physical device.
diff --git a/uefi-test-runner/src/proto/media/known_disk.rs b/uefi-test-runner/src/proto/media/known_disk.rs
index f1d6744ee..6f6de0658 100644
--- a/uefi-test-runner/src/proto/media/known_disk.rs
+++ b/uefi-test-runner/src/proto/media/known_disk.rs
@@ -1,11 +1,15 @@
 use alloc::string::ToString;
+use core::ffi::c_void;
+use core::ptr::NonNull;
 use uefi::prelude::*;
+use uefi::proto::media::disk::{DiskIo, DiskIo2, DiskIo2Token};
 use uefi::proto::media::file::{
     Directory, File, FileAttribute, FileInfo, FileMode, FileSystemInfo,
 };
 use uefi::proto::media::fs::SimpleFileSystem;
-use uefi::table::boot::{OpenProtocolAttributes, OpenProtocolParams};
+use uefi::table::boot::{EventType, MemoryType, OpenProtocolAttributes, OpenProtocolParams, Tpl};
 use uefi::table::runtime::{Daylight, Time, TimeParams};
+use uefi::Event;
 
 /// Test directory entry iteration.
 fn test_existing_dir(directory: &mut Directory) {
@@ -138,6 +142,103 @@ fn test_create_file(directory: &mut Directory) {
     file.write(b"test output data").unwrap();
 }
 
+/// Tests raw disk I/O.
+fn test_raw_disk_io(handle: Handle, image: Handle, bt: &BootServices) {
+    info!("Testing raw disk I/O");
+
+    // Open the disk I/O protocol on the input handle
+    let disk_io = bt
+        .open_protocol::<DiskIo>(
+            OpenProtocolParams {
+                handle,
+                agent: image,
+                controller: None,
+            },
+            OpenProtocolAttributes::GetProtocol,
+        )
+        .expect("Failed to get disk I/O protocol");
+
+    // Allocate a temporary buffer to read into
+    const SIZE: usize = 512;
+    let buf = bt
+        .allocate_pool(MemoryType::LOADER_DATA, SIZE)
+        .expect("Failed to allocate temporary buffer");
+
+    // SAFETY: A valid buffer of `SIZE` bytes was allocated above
+    let slice = unsafe { core::slice::from_raw_parts_mut(buf, SIZE) };
+
+    // Read from the first sector of the disk into the buffer
+    disk_io
+        .read_disk(0, 0, slice)
+        .expect("Failed to read from disk");
+
+    // Verify that the disk's MBR signature is correct
+    assert_eq!(slice[510], 0x55);
+    assert_eq!(slice[511], 0xaa);
+
+    info!("Raw disk I/O succeeded");
+    bt.free_pool(buf).unwrap();
+}
+
+/// Asynchronous disk I/O 2 transaction callback
+unsafe extern "efiapi" fn disk_io2_callback(_event: Event, ctx: Option<NonNull<c_void>>) {
+    let ptr = ctx.unwrap().as_ptr() as *const u8;
+
+    // Verify that the disk's MBR signature is correct
+    assert_eq!(*ptr.offset(510), 0x55);
+    assert_eq!(*ptr.offset(511), 0xaa);
+}
+
+/// Tests raw disk I/O through the DiskIo2 protocol.
+fn test_raw_disk_io2(handle: Handle, image: Handle, bt: &BootServices) {
+    info!("Testing raw disk I/O 2");
+
+    // Open the disk I/O protocol on the input handle
+    let disk_io2 = bt
+        .open_protocol::<DiskIo2>(
+            OpenProtocolParams {
+                handle,
+                agent: image,
+                controller: None,
+            },
+            OpenProtocolAttributes::GetProtocol,
+        )
+        .expect("Failed to get disk I/O 2 protocol");
+
+    // Allocate a temporary buffer to read into
+    const SIZE: usize = 512;
+    let buf = bt
+        .allocate_pool(MemoryType::LOADER_DATA, SIZE)
+        .expect("Failed to allocate temporary buffer");
+
+    // Create an event callback for the disk read completion
+    let event = unsafe {
+        bt.create_event(
+            EventType::NOTIFY_SIGNAL,
+            Tpl::NOTIFY,
+            Some(disk_io2_callback),
+            NonNull::new(buf as *mut c_void),
+        )
+        .expect("Failed to create event for disk I/O 2 transaction")
+    };
+
+    // Read from the first sector of the disk into the buffer
+    // SAFETY: The cloned `event` is only used for this transaction
+    unsafe {
+        let mut token = DiskIo2Token {
+            event: Some(event.unsafe_clone()),
+            transaction_status: uefi::Status::SUCCESS,
+        };
+        disk_io2
+            .read_disk_raw(0, 0, &mut token, SIZE, buf)
+            .expect("Failed to read from disk");
+    }
+
+    info!("Raw disk I/O 2 succeeded");
+    bt.close_event(event).unwrap();
+    bt.free_pool(buf).unwrap();
+}
+
 /// Run various tests on a special test disk. The disk is created by
 /// xtask/src/disk.rs.
 pub fn test_known_disk(image: Handle, bt: &BootServices) {
@@ -154,6 +255,10 @@ pub fn test_known_disk(image: Handle, bt: &BootServices) {
 
     let mut found_test_disk = false;
     for handle in handles {
+        // Test raw disk I/O first
+        test_raw_disk_io(handle, image, bt);
+        test_raw_disk_io2(handle, image, bt);
+
         let mut sfs = bt
             .open_protocol::<SimpleFileSystem>(
                 OpenProtocolParams {

From 7151e7fc517e0f63837ae1b2ccbb220abb4b60a3 Mon Sep 17 00:00:00 2001
From: Laurin-Luis Lehning <65224843+e820@users.noreply.github.com>
Date: Sat, 30 Jul 2022 23:11:23 +0200
Subject: [PATCH 05/11] Move close_event() to callback

---
 uefi-test-runner/src/proto/media/known_disk.rs | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/uefi-test-runner/src/proto/media/known_disk.rs b/uefi-test-runner/src/proto/media/known_disk.rs
index 6f6de0658..5374e9d66 100644
--- a/uefi-test-runner/src/proto/media/known_disk.rs
+++ b/uefi-test-runner/src/proto/media/known_disk.rs
@@ -10,6 +10,7 @@ use uefi::proto::media::fs::SimpleFileSystem;
 use uefi::table::boot::{EventType, MemoryType, OpenProtocolAttributes, OpenProtocolParams, Tpl};
 use uefi::table::runtime::{Daylight, Time, TimeParams};
 use uefi::Event;
+use uefi_services::system_table;
 
 /// Test directory entry iteration.
 fn test_existing_dir(directory: &mut Directory) {
@@ -181,12 +182,14 @@ fn test_raw_disk_io(handle: Handle, image: Handle, bt: &BootServices) {
 }
 
 /// Asynchronous disk I/O 2 transaction callback
-unsafe extern "efiapi" fn disk_io2_callback(_event: Event, ctx: Option<NonNull<c_void>>) {
+unsafe extern "efiapi" fn disk_io2_callback(event: Event, ctx: Option<NonNull<c_void>>) {
     let ptr = ctx.unwrap().as_ptr() as *const u8;
 
     // Verify that the disk's MBR signature is correct
     assert_eq!(*ptr.offset(510), 0x55);
     assert_eq!(*ptr.offset(511), 0xaa);
+
+    system_table().as_ref().boot_services().close_event(event).unwrap();
 }
 
 /// Tests raw disk I/O through the DiskIo2 protocol.
@@ -233,9 +236,8 @@ fn test_raw_disk_io2(handle: Handle, image: Handle, bt: &BootServices) {
             .read_disk_raw(0, 0, &mut token, SIZE, buf)
             .expect("Failed to read from disk");
     }
-
+ 
     info!("Raw disk I/O 2 succeeded");
-    bt.close_event(event).unwrap();
     bt.free_pool(buf).unwrap();
 }
 

From 7b7a7dad573ae302b028a690a658cfa4b7136e1f Mon Sep 17 00:00:00 2001
From: Laurin-Luis Lehning <65224843+e820@users.noreply.github.com>
Date: Sun, 31 Jul 2022 00:07:36 +0200
Subject: [PATCH 06/11] Format and make DiskIo2 test optional

---
 .../src/proto/media/known_disk.rs             | 31 ++++++++++++-------
 1 file changed, 19 insertions(+), 12 deletions(-)

diff --git a/uefi-test-runner/src/proto/media/known_disk.rs b/uefi-test-runner/src/proto/media/known_disk.rs
index 5374e9d66..beba5c9f4 100644
--- a/uefi-test-runner/src/proto/media/known_disk.rs
+++ b/uefi-test-runner/src/proto/media/known_disk.rs
@@ -189,7 +189,11 @@ unsafe extern "efiapi" fn disk_io2_callback(event: Event, ctx: Option<NonNull<c_
     assert_eq!(*ptr.offset(510), 0x55);
     assert_eq!(*ptr.offset(511), 0xaa);
 
-    system_table().as_ref().boot_services().close_event(event).unwrap();
+    system_table()
+        .as_ref()
+        .boot_services()
+        .close_event(event)
+        .unwrap();
 }
 
 /// Tests raw disk I/O through the DiskIo2 protocol.
@@ -197,16 +201,19 @@ fn test_raw_disk_io2(handle: Handle, image: Handle, bt: &BootServices) {
     info!("Testing raw disk I/O 2");
 
     // Open the disk I/O protocol on the input handle
-    let disk_io2 = bt
-        .open_protocol::<DiskIo2>(
-            OpenProtocolParams {
-                handle,
-                agent: image,
-                controller: None,
-            },
-            OpenProtocolAttributes::GetProtocol,
-        )
-        .expect("Failed to get disk I/O 2 protocol");
+    let disk_io2 = bt.open_protocol::<DiskIo2>(
+        OpenProtocolParams {
+            handle,
+            agent: image,
+            controller: None,
+        },
+        OpenProtocolAttributes::GetProtocol,
+    );
+    if disk_io2.is_err() {
+        return;
+    }
+
+    let disk_io2 = disk_io2.unwrap();
 
     // Allocate a temporary buffer to read into
     const SIZE: usize = 512;
@@ -236,7 +243,7 @@ fn test_raw_disk_io2(handle: Handle, image: Handle, bt: &BootServices) {
             .read_disk_raw(0, 0, &mut token, SIZE, buf)
             .expect("Failed to read from disk");
     }
- 
+
     info!("Raw disk I/O 2 succeeded");
     bt.free_pool(buf).unwrap();
 }

From 62fb39f26c975fab16c2ed12a0bf601aa66cb09a Mon Sep 17 00:00:00 2001
From: Laurin-Luis Lehning <65224843+e820@users.noreply.github.com>
Date: Sun, 31 Jul 2022 00:54:01 +0200
Subject: [PATCH 07/11] Add safety docs & rework test for DiskIo2

---
 src/proto/media/disk.rs                       |  10 ++
 .../src/proto/media/known_disk.rs             | 100 ++++++++++--------
 2 files changed, 64 insertions(+), 46 deletions(-)

diff --git a/src/proto/media/disk.rs b/src/proto/media/disk.rs
index 323407627..92dec4000 100644
--- a/src/proto/media/disk.rs
+++ b/src/proto/media/disk.rs
@@ -126,6 +126,11 @@ impl DiskIo2 {
     /// * `len` - Buffer size.
     /// * `buffer` - Buffer to read into.
     ///
+    /// # Safety:
+    ///
+    /// Because of the asynchronous nature of the disk transaction, manual lifetime
+    /// tracking is required.
+    ///
     /// # Errors:
     /// * `uefi::status::INVALID_PARAMETER` The read request contains device addresses
     ///                                     that are not valid for the device.
@@ -155,6 +160,11 @@ impl DiskIo2 {
     /// * `len` - Buffer size.
     /// * `buffer` - Buffer to write from.
     ///
+    /// # Safety:
+    ///
+    /// Because of the asynchronous nature of the disk transaction, manual lifetime
+    /// tracking is required.
+    ///
     /// # Errors:
     /// * `uefi::status::INVALID_PARAMETER` The write request contains device addresses
     ///                                     that are not valid for the device.
diff --git a/uefi-test-runner/src/proto/media/known_disk.rs b/uefi-test-runner/src/proto/media/known_disk.rs
index beba5c9f4..7577ceb72 100644
--- a/uefi-test-runner/src/proto/media/known_disk.rs
+++ b/uefi-test-runner/src/proto/media/known_disk.rs
@@ -181,19 +181,31 @@ fn test_raw_disk_io(handle: Handle, image: Handle, bt: &BootServices) {
     bt.free_pool(buf).unwrap();
 }
 
+/// Asynchronous disk I/O task context
+#[repr(C)]
+struct DiskIoTask {
+    /// Token for the transaction
+    token: DiskIo2Token,
+    /// Pointer to a buffer holding the read data
+    buffer: *mut u8,
+}
+
 /// Asynchronous disk I/O 2 transaction callback
 unsafe extern "efiapi" fn disk_io2_callback(event: Event, ctx: Option<NonNull<c_void>>) {
-    let ptr = ctx.unwrap().as_ptr() as *const u8;
+    let task = ctx.unwrap().as_ptr() as *mut DiskIoTask;
 
     // Verify that the disk's MBR signature is correct
-    assert_eq!(*ptr.offset(510), 0x55);
-    assert_eq!(*ptr.offset(511), 0xaa);
-
-    system_table()
-        .as_ref()
-        .boot_services()
-        .close_event(event)
-        .unwrap();
+    assert_eq!((*task).token.transaction_status, uefi::Status::SUCCESS);
+    assert_eq!(*(*task).buffer.offset(510), 0x55);
+    assert_eq!(*(*task).buffer.offset(511), 0xaa);
+
+    // Close the completion event
+    let bt = system_table().as_ref().boot_services();
+    bt.close_event(event).unwrap();
+
+    // Free the disk data buffer & task context
+    bt.free_pool((*task).buffer).unwrap();
+    bt.free_pool(task as *mut u8).unwrap();
 }
 
 /// Tests raw disk I/O through the DiskIo2 protocol.
@@ -201,51 +213,47 @@ fn test_raw_disk_io2(handle: Handle, image: Handle, bt: &BootServices) {
     info!("Testing raw disk I/O 2");
 
     // Open the disk I/O protocol on the input handle
-    let disk_io2 = bt.open_protocol::<DiskIo2>(
+    if let Ok(disk_io2) = bt.open_protocol::<DiskIo2>(
         OpenProtocolParams {
             handle,
             agent: image,
             controller: None,
         },
         OpenProtocolAttributes::GetProtocol,
-    );
-    if disk_io2.is_err() {
-        return;
-    }
-
-    let disk_io2 = disk_io2.unwrap();
-
-    // Allocate a temporary buffer to read into
-    const SIZE: usize = 512;
-    let buf = bt
-        .allocate_pool(MemoryType::LOADER_DATA, SIZE)
-        .expect("Failed to allocate temporary buffer");
-
-    // Create an event callback for the disk read completion
-    let event = unsafe {
-        bt.create_event(
-            EventType::NOTIFY_SIGNAL,
-            Tpl::NOTIFY,
-            Some(disk_io2_callback),
-            NonNull::new(buf as *mut c_void),
-        )
-        .expect("Failed to create event for disk I/O 2 transaction")
-    };
+    ) {
+        // Allocate the task context structure
+        let task = bt
+            .allocate_pool(MemoryType::LOADER_DATA, core::mem::size_of::<DiskIoTask>())
+            .expect("Failed to allocate task struct for disk I/O")
+            as *mut DiskIoTask;
+
+        // SAFETY: `task` refers to a valid allocation of at least `size_of::<DiskIoTask>()`
+        unsafe {
+            // Create the completion event as part of the disk I/O token
+            (*task).token.event = Some(
+                bt.create_event(
+                    EventType::NOTIFY_SIGNAL,
+                    Tpl::NOTIFY,
+                    Some(disk_io2_callback),
+                    NonNull::new(task as *mut c_void),
+                )
+                .expect("Failed to create disk I/O completion event"),
+            );
+
+            // Allocate a buffer for the transaction to read into
+            const SIZE_TO_READ: usize = 512;
+            (*task).buffer = bt
+                .allocate_pool(MemoryType::LOADER_DATA, SIZE_TO_READ)
+                .expect("Failed to allocate buffer for disk I/O");
+
+            // Initiate the asynchronous read operation
+            disk_io2
+                .read_disk_raw(0, 0, &mut (*task).token, SIZE_TO_READ, (*task).buffer)
+                .expect("Failed to initiate asynchronous disk I/O read");
+        }
 
-    // Read from the first sector of the disk into the buffer
-    // SAFETY: The cloned `event` is only used for this transaction
-    unsafe {
-        let mut token = DiskIo2Token {
-            event: Some(event.unsafe_clone()),
-            transaction_status: uefi::Status::SUCCESS,
-        };
-        disk_io2
-            .read_disk_raw(0, 0, &mut token, SIZE, buf)
-            .expect("Failed to read from disk");
+        info!("Raw disk I/O 2 succeeded");
     }
-
-    info!("Raw disk I/O 2 succeeded");
-    bt.free_pool(buf).unwrap();
 }
 
 /// Run various tests on a special test disk. The disk is created by

From bdf8374e8dedb1f1b724b1caa007919192077cf7 Mon Sep 17 00:00:00 2001
From: Laurin-Luis Lehning <65224843+e820@users.noreply.github.com>
Date: Sun, 31 Jul 2022 01:19:02 +0200
Subject: [PATCH 08/11] Remove colon after Safety doc section

---
 src/proto/media/disk.rs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/proto/media/disk.rs b/src/proto/media/disk.rs
index 92dec4000..237442e7a 100644
--- a/src/proto/media/disk.rs
+++ b/src/proto/media/disk.rs
@@ -126,7 +126,7 @@ impl DiskIo2 {
     /// * `len` - Buffer size.
     /// * `buffer` - Buffer to read into.
     ///
-    /// # Safety:
+    /// # Safety
     ///
     /// Because of the asynchronous nature of the disk transaction, manual lifetime
     /// tracking is required.
@@ -160,7 +160,7 @@ impl DiskIo2 {
     /// * `len` - Buffer size.
     /// * `buffer` - Buffer to write from.
     ///
-    /// # Safety:
+    /// # Safety
     ///
     /// Because of the asynchronous nature of the disk transaction, manual lifetime
     /// tracking is required.

From cc9970b04aab38136a1d48bd96ecc7f4872bd502 Mon Sep 17 00:00:00 2001
From: Laurin-Luis Lehning <65224843+e820@users.noreply.github.com>
Date: Sun, 31 Jul 2022 17:54:55 +0200
Subject: [PATCH 09/11] Don't hardcode media ID during DiskIo & DiskIo2 tests

---
 .../src/proto/media/known_disk.rs             | 58 +++++++++++++------
 1 file changed, 40 insertions(+), 18 deletions(-)

diff --git a/uefi-test-runner/src/proto/media/known_disk.rs b/uefi-test-runner/src/proto/media/known_disk.rs
index 7577ceb72..8c0dec85c 100644
--- a/uefi-test-runner/src/proto/media/known_disk.rs
+++ b/uefi-test-runner/src/proto/media/known_disk.rs
@@ -2,6 +2,7 @@ use alloc::string::ToString;
 use core::ffi::c_void;
 use core::ptr::NonNull;
 use uefi::prelude::*;
+use uefi::proto::media::block::BlockIO;
 use uefi::proto::media::disk::{DiskIo, DiskIo2, DiskIo2Token};
 use uefi::proto::media::file::{
     Directory, File, FileAttribute, FileInfo, FileMode, FileSystemInfo,
@@ -147,6 +148,18 @@ fn test_create_file(directory: &mut Directory) {
 fn test_raw_disk_io(handle: Handle, image: Handle, bt: &BootServices) {
     info!("Testing raw disk I/O");
 
+    // Open the block I/O protocol on the handle
+    let block_io = bt
+        .open_protocol::<BlockIO>(
+            OpenProtocolParams {
+                handle,
+                agent: image,
+                controller: None,
+            },
+            OpenProtocolAttributes::GetProtocol,
+        )
+        .expect("Failed to get block I/O protocol");
+
     // Open the disk I/O protocol on the input handle
     let disk_io = bt
         .open_protocol::<DiskIo>(
@@ -159,26 +172,17 @@ fn test_raw_disk_io(handle: Handle, image: Handle, bt: &BootServices) {
         )
         .expect("Failed to get disk I/O protocol");
 
-    // Allocate a temporary buffer to read into
-    const SIZE: usize = 512;
-    let buf = bt
-        .allocate_pool(MemoryType::LOADER_DATA, SIZE)
-        .expect("Failed to allocate temporary buffer");
-
-    // SAFETY: A valid buffer of `SIZE` bytes was allocated above
-    let slice = unsafe { core::slice::from_raw_parts_mut(buf, SIZE) };
-
     // Read from the first sector of the disk into the buffer
+    let mut buf = vec![0; 512];
     disk_io
-        .read_disk(0, 0, slice)
+        .read_disk(block_io.media().media_id(), 0, &mut buf)
         .expect("Failed to read from disk");
 
     // Verify that the disk's MBR signature is correct
-    assert_eq!(slice[510], 0x55);
-    assert_eq!(slice[511], 0xaa);
+    assert_eq!(buf[510], 0x55);
+    assert_eq!(buf[511], 0xaa);
 
     info!("Raw disk I/O succeeded");
-    bt.free_pool(buf).unwrap();
 }
 
 /// Asynchronous disk I/O task context
@@ -221,6 +225,18 @@ fn test_raw_disk_io2(handle: Handle, image: Handle, bt: &BootServices) {
         },
         OpenProtocolAttributes::GetProtocol,
     ) {
+        // Open the block I/O protocol on the handle
+        let block_io = bt
+            .open_protocol::<BlockIO>(
+                OpenProtocolParams {
+                    handle,
+                    agent: image,
+                    controller: None,
+                },
+                OpenProtocolAttributes::GetProtocol,
+            )
+            .expect("Failed to get block I/O protocol");
+
         // Allocate the task context structure
         let task = bt
             .allocate_pool(MemoryType::LOADER_DATA, core::mem::size_of::<DiskIoTask>())
@@ -248,7 +264,13 @@ fn test_raw_disk_io2(handle: Handle, image: Handle, bt: &BootServices) {
 
             // Initiate the asynchronous read operation
             disk_io2
-                .read_disk_raw(0, 0, &mut (*task).token, SIZE_TO_READ, (*task).buffer)
+                .read_disk_raw(
+                    block_io.media().media_id(),
+                    0,
+                    &mut (*task).token,
+                    SIZE_TO_READ,
+                    (*task).buffer,
+                )
                 .expect("Failed to initiate asynchronous disk I/O read");
         }
 
@@ -272,10 +294,6 @@ pub fn test_known_disk(image: Handle, bt: &BootServices) {
 
     let mut found_test_disk = false;
     for handle in handles {
-        // Test raw disk I/O first
-        test_raw_disk_io(handle, image, bt);
-        test_raw_disk_io2(handle, image, bt);
-
         let mut sfs = bt
             .open_protocol::<SimpleFileSystem>(
                 OpenProtocolParams {
@@ -300,6 +318,10 @@ pub fn test_known_disk(image: Handle, bt: &BootServices) {
             continue;
         }
 
+        // Test raw disk I/O first
+        test_raw_disk_io(handle, image, bt);
+        test_raw_disk_io2(handle, image, bt);
+
         assert!(!fs_info.read_only());
         assert_eq!(fs_info.volume_size(), 512 * 1192);
         assert_eq!(fs_info.free_space(), 512 * 1190);

From c326d5175051cf32157ef70ba20f402aed63d0cc Mon Sep 17 00:00:00 2001
From: Laurin-Luis Lehning <65224843+e820@users.noreply.github.com>
Date: Sun, 31 Jul 2022 22:13:58 +0200
Subject: [PATCH 10/11] Rework disk I/O API and test

---
 src/proto/media/disk.rs                       | 20 +++--
 .../src/proto/media/known_disk.rs             | 87 ++++++++-----------
 2 files changed, 49 insertions(+), 58 deletions(-)

diff --git a/src/proto/media/disk.rs b/src/proto/media/disk.rs
index 237442e7a..9cf7926a7 100644
--- a/src/proto/media/disk.rs
+++ b/src/proto/media/disk.rs
@@ -2,6 +2,7 @@
 
 use crate::proto::Protocol;
 use crate::{unsafe_guid, Event, Result, Status};
+use core::ptr::NonNull;
 
 /// The disk I/O protocol.
 ///
@@ -73,7 +74,7 @@ impl DiskIo {
 #[repr(C)]
 pub struct DiskIo2Token {
     /// Event to be signalled when an asynchronous disk I/O operation completes.
-    pub event: Option<Event>,
+    pub event: Event,
     /// Transaction status code.
     pub transaction_status: Status,
 }
@@ -92,7 +93,7 @@ pub struct DiskIo2 {
         this: &DiskIo2,
         media_id: u32,
         offset: u64,
-        token: &mut DiskIo2Token,
+        token: Option<NonNull<DiskIo2Token>>,
         len: usize,
         buffer: *mut u8,
     ) -> Status,
@@ -100,11 +101,12 @@ pub struct DiskIo2 {
         this: &mut DiskIo2,
         media_id: u32,
         offset: u64,
-        token: &mut DiskIo2Token,
+        token: Option<NonNull<DiskIo2Token>>,
         len: usize,
         buffer: *const u8,
     ) -> Status,
-    flush_disk_ex: extern "efiapi" fn(this: &mut DiskIo2, token: &mut DiskIo2Token) -> Status,
+    flush_disk_ex:
+        extern "efiapi" fn(this: &mut DiskIo2, token: Option<NonNull<DiskIo2Token>>) -> Status,
 }
 
 impl DiskIo2 {
@@ -144,11 +146,11 @@ impl DiskIo2 {
         &self,
         media_id: u32,
         offset: u64,
-        token: *mut DiskIo2Token,
+        token: Option<NonNull<DiskIo2Token>>,
         len: usize,
         buffer: *mut u8,
     ) -> Result {
-        (self.read_disk_ex)(self, media_id, offset, &mut *token, len, buffer).into()
+        (self.read_disk_ex)(self, media_id, offset, token, len, buffer).into()
     }
 
     /// Writes bytes to the disk device.
@@ -179,11 +181,11 @@ impl DiskIo2 {
         &mut self,
         media_id: u32,
         offset: u64,
-        token: *mut DiskIo2Token,
+        token: Option<NonNull<DiskIo2Token>>,
         len: usize,
         buffer: *const u8,
     ) -> Result {
-        (self.write_disk_ex)(self, media_id, offset, &mut *token, len, buffer).into()
+        (self.write_disk_ex)(self, media_id, offset, token, len, buffer).into()
     }
 
     /// Flushes all modified data to the physical device.
@@ -200,7 +202,7 @@ impl DiskIo2 {
     /// * `uefi::status::DEVICE_ERROR`      The device reported an error while performing
     ///                                     the flush operation.
     /// * `uefi::status::WRITE_PROTECTED`   The device cannot be written to.
-    pub fn flush_disk(&mut self, token: &mut DiskIo2Token) -> Result {
+    pub fn flush_disk(&mut self, token: Option<NonNull<DiskIo2Token>>) -> Result {
         (self.flush_disk_ex)(self, token).into()
     }
 }
diff --git a/uefi-test-runner/src/proto/media/known_disk.rs b/uefi-test-runner/src/proto/media/known_disk.rs
index 8c0dec85c..558a76ebb 100644
--- a/uefi-test-runner/src/proto/media/known_disk.rs
+++ b/uefi-test-runner/src/proto/media/known_disk.rs
@@ -1,5 +1,4 @@
 use alloc::string::ToString;
-use core::ffi::c_void;
 use core::ptr::NonNull;
 use uefi::prelude::*;
 use uefi::proto::media::block::BlockIO;
@@ -8,10 +7,8 @@ use uefi::proto::media::file::{
     Directory, File, FileAttribute, FileInfo, FileMode, FileSystemInfo,
 };
 use uefi::proto::media::fs::SimpleFileSystem;
-use uefi::table::boot::{EventType, MemoryType, OpenProtocolAttributes, OpenProtocolParams, Tpl};
+use uefi::table::boot::{EventType, OpenProtocolAttributes, OpenProtocolParams, Tpl};
 use uefi::table::runtime::{Daylight, Time, TimeParams};
-use uefi::Event;
-use uefi_services::system_table;
 
 /// Test directory entry iteration.
 fn test_existing_dir(directory: &mut Directory) {
@@ -190,26 +187,8 @@ fn test_raw_disk_io(handle: Handle, image: Handle, bt: &BootServices) {
 struct DiskIoTask {
     /// Token for the transaction
     token: DiskIo2Token,
-    /// Pointer to a buffer holding the read data
-    buffer: *mut u8,
-}
-
-/// Asynchronous disk I/O 2 transaction callback
-unsafe extern "efiapi" fn disk_io2_callback(event: Event, ctx: Option<NonNull<c_void>>) {
-    let task = ctx.unwrap().as_ptr() as *mut DiskIoTask;
-
-    // Verify that the disk's MBR signature is correct
-    assert_eq!((*task).token.transaction_status, uefi::Status::SUCCESS);
-    assert_eq!(*(*task).buffer.offset(510), 0x55);
-    assert_eq!(*(*task).buffer.offset(511), 0xaa);
-
-    // Close the completion event
-    let bt = system_table().as_ref().boot_services();
-    bt.close_event(event).unwrap();
-
-    // Free the disk data buffer & task context
-    bt.free_pool((*task).buffer).unwrap();
-    bt.free_pool(task as *mut u8).unwrap();
+    /// Buffer holding the read data
+    buffer: [u8; 512],
 }
 
 /// Tests raw disk I/O through the DiskIo2 protocol.
@@ -237,44 +216,54 @@ fn test_raw_disk_io2(handle: Handle, image: Handle, bt: &BootServices) {
             )
             .expect("Failed to get block I/O protocol");
 
-        // Allocate the task context structure
-        let task = bt
-            .allocate_pool(MemoryType::LOADER_DATA, core::mem::size_of::<DiskIoTask>())
-            .expect("Failed to allocate task struct for disk I/O")
-            as *mut DiskIoTask;
-
-        // SAFETY: `task` refers to a valid allocation of at least `size_of::<DiskIoTask>()`
         unsafe {
-            // Create the completion event as part of the disk I/O token
-            (*task).token.event = Some(
-                bt.create_event(
-                    EventType::NOTIFY_SIGNAL,
+            // Create the task context structure
+            let mut task = core::mem::MaybeUninit::uninit();
+
+            // Create the completion event
+            let mut event = bt
+                .create_event(
+                    EventType::empty(),
                     Tpl::NOTIFY,
-                    Some(disk_io2_callback),
-                    NonNull::new(task as *mut c_void),
+                    None,
+                    NonNull::new(task.as_mut_ptr() as _),
                 )
-                .expect("Failed to create disk I/O completion event"),
-            );
+                .expect("Failed to create disk I/O completion event");
+
+            // Initialise the task context
+            task.write(DiskIoTask {
+                token: DiskIo2Token {
+                    event: event.unsafe_clone(),
+                    transaction_status: uefi::Status::NOT_READY,
+                },
+                buffer: [0; 512],
+            });
 
-            // Allocate a buffer for the transaction to read into
-            const SIZE_TO_READ: usize = 512;
-            (*task).buffer = bt
-                .allocate_pool(MemoryType::LOADER_DATA, SIZE_TO_READ)
-                .expect("Failed to allocate buffer for disk I/O");
+            // Get a mutable reference to the task to not move it
+            let task_ref = task.assume_init_mut();
 
             // Initiate the asynchronous read operation
             disk_io2
                 .read_disk_raw(
                     block_io.media().media_id(),
                     0,
-                    &mut (*task).token,
-                    SIZE_TO_READ,
-                    (*task).buffer,
+                    NonNull::new(&mut task_ref.token as _),
+                    task_ref.buffer.len(),
+                    task_ref.buffer.as_mut_ptr(),
                 )
                 .expect("Failed to initiate asynchronous disk I/O read");
-        }
 
-        info!("Raw disk I/O 2 succeeded");
+            // Wait for the transaction to complete
+            bt.wait_for_event(core::slice::from_mut(&mut event))
+                .expect("Failed to wait on completion event");
+
+            // Verify that the disk's MBR signature is correct
+            assert_eq!(task_ref.token.transaction_status, uefi::Status::SUCCESS);
+            assert_eq!(task_ref.buffer[510], 0x55);
+            assert_eq!(task_ref.buffer[511], 0xaa);
+
+            info!("Raw disk I/O 2 succeeded");
+        }
     }
 }
 

From 5f67272a6bf670a19e8b36aaa8c298c7d1ca7305 Mon Sep 17 00:00:00 2001
From: Laurin-Luis Lehning <65224843+e820@users.noreply.github.com>
Date: Sun, 31 Jul 2022 22:25:33 +0200
Subject: [PATCH 11/11] Simplify disk I/O 2 test

---
 .../src/proto/media/known_disk.rs             | 29 ++++++-------------
 1 file changed, 9 insertions(+), 20 deletions(-)

diff --git a/uefi-test-runner/src/proto/media/known_disk.rs b/uefi-test-runner/src/proto/media/known_disk.rs
index 558a76ebb..6b69390c6 100644
--- a/uefi-test-runner/src/proto/media/known_disk.rs
+++ b/uefi-test-runner/src/proto/media/known_disk.rs
@@ -217,39 +217,28 @@ fn test_raw_disk_io2(handle: Handle, image: Handle, bt: &BootServices) {
             .expect("Failed to get block I/O protocol");
 
         unsafe {
-            // Create the task context structure
-            let mut task = core::mem::MaybeUninit::uninit();
-
             // Create the completion event
             let mut event = bt
-                .create_event(
-                    EventType::empty(),
-                    Tpl::NOTIFY,
-                    None,
-                    NonNull::new(task.as_mut_ptr() as _),
-                )
+                .create_event(EventType::empty(), Tpl::NOTIFY, None, None)
                 .expect("Failed to create disk I/O completion event");
 
             // Initialise the task context
-            task.write(DiskIoTask {
+            let mut task = DiskIoTask {
                 token: DiskIo2Token {
                     event: event.unsafe_clone(),
                     transaction_status: uefi::Status::NOT_READY,
                 },
                 buffer: [0; 512],
-            });
-
-            // Get a mutable reference to the task to not move it
-            let task_ref = task.assume_init_mut();
+            };
 
             // Initiate the asynchronous read operation
             disk_io2
                 .read_disk_raw(
                     block_io.media().media_id(),
                     0,
-                    NonNull::new(&mut task_ref.token as _),
-                    task_ref.buffer.len(),
-                    task_ref.buffer.as_mut_ptr(),
+                    NonNull::new(&mut task.token as _),
+                    task.buffer.len(),
+                    task.buffer.as_mut_ptr(),
                 )
                 .expect("Failed to initiate asynchronous disk I/O read");
 
@@ -258,9 +247,9 @@ fn test_raw_disk_io2(handle: Handle, image: Handle, bt: &BootServices) {
                 .expect("Failed to wait on completion event");
 
             // Verify that the disk's MBR signature is correct
-            assert_eq!(task_ref.token.transaction_status, uefi::Status::SUCCESS);
-            assert_eq!(task_ref.buffer[510], 0x55);
-            assert_eq!(task_ref.buffer[511], 0xaa);
+            assert_eq!(task.token.transaction_status, uefi::Status::SUCCESS);
+            assert_eq!(task.buffer[510], 0x55);
+            assert_eq!(task.buffer[511], 0xaa);
 
             info!("Raw disk I/O 2 succeeded");
         }