Add DiskIo and DiskIo2 protocols

src/proto/media/disk.rs

@@ -0,0 +1,208 @@
+//! Disk I/O protocols.
+
+use crate::proto::Protocol;
+use crate::{unsafe_guid, Event, Result, Status};
+use core::ptr::NonNull;
+
+/// 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-8e4f-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.
+    pub event: Event,
+    /// Transaction status code.
+    pub 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-9e54-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: Option<NonNull<DiskIo2Token>>,
+        len: usize,
+        buffer: *mut u8,
+    ) -> Status,
+    write_disk_ex: extern "efiapi" fn(
+        this: &mut DiskIo2,
+        media_id: u32,
+        offset: u64,
+        token: Option<NonNull<DiskIo2Token>>,
+        len: usize,
+        buffer: *const u8,
+    ) -> Status,
+    flush_disk_ex:
+        extern "efiapi" fn(this: &mut DiskIo2, token: Option<NonNull<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.
+    /// * `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.
+    /// * `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 unsafe fn read_disk_raw(
+        &self,
+        media_id: u32,
+        offset: u64,
+        token: Option<NonNull<DiskIo2Token>>,
+        len: usize,
+        buffer: *mut u8,
+    ) -> Result {
+        (self.read_disk_ex)(self, media_id, offset, token, len, buffer).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.
+    /// * `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.
+    /// * `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 unsafe fn write_disk_raw(
+        &mut self,
+        media_id: u32,
+        offset: u64,
+        token: Option<NonNull<DiskIo2Token>>,
+        len: usize,
+        buffer: *const u8,
+    ) -> Result {
+        (self.write_disk_ex)(self, media_id, offset, token, len, buffer).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: Option<NonNull<DiskIo2Token>>) -> Result {
+        (self.flush_disk_ex)(self, token).into()
+    }
+}

src/proto/media/mod.rs

@@ -7,5 +7,6 @@
 pub mod file;
 
 pub mod block;
+pub mod disk;
 pub mod fs;
 pub mod partition;

uefi-test-runner/src/proto/media/known_disk.rs

@@ -1,10 +1,13 @@
 use alloc::string::ToString;
+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,
 };
 use uefi::proto::media::fs::SimpleFileSystem;
-use uefi::table::boot::{OpenProtocolAttributes, OpenProtocolParams};
+use uefi::table::boot::{EventType, OpenProtocolAttributes, OpenProtocolParams, Tpl};
 use uefi::table::runtime::{Daylight, Time, TimeParams};
 
 /// Test directory entry iteration.
@@ -138,6 +141,121 @@ 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 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>(
+            OpenProtocolParams {
+                handle,
+                agent: image,
+                controller: None,
+            },
+            OpenProtocolAttributes::GetProtocol,
+        )
+        .expect("Failed to get disk I/O protocol");
+
+    // Read from the first sector of the disk into the buffer
+    let mut buf = vec![0; 512];
+    disk_io
+        .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!(buf[510], 0x55);
+    assert_eq!(buf[511], 0xaa);
+
+    info!("Raw disk I/O succeeded");
+}
+
+/// Asynchronous disk I/O task context
+#[repr(C)]
+struct DiskIoTask {
+    /// Token for the transaction
+    token: DiskIo2Token,
+    /// Buffer holding the read data
+    buffer: [u8; 512],
+}
+
+/// 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
+    if let Ok(disk_io2) = bt.open_protocol::<DiskIo2>(
+        OpenProtocolParams {
+            handle,
+            agent: image,
+            controller: None,
+        },
+        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");
+
+        unsafe {
+            // Create the completion event
+            let mut event = bt
+                .create_event(EventType::empty(), Tpl::NOTIFY, None, None)
+                .expect("Failed to create disk I/O completion event");
+
+            // Initialise the task context
+            let mut task = DiskIoTask {
+                token: DiskIo2Token {
+                    event: event.unsafe_clone(),
+                    transaction_status: uefi::Status::NOT_READY,
+                },
+                buffer: [0; 512],
+            };
+
+            // Initiate the asynchronous read operation
+            disk_io2
+                .read_disk_raw(
+                    block_io.media().media_id(),
+                    0,
+                    NonNull::new(&mut task.token as _),
+                    task.buffer.len(),
+                    task.buffer.as_mut_ptr(),
+                )
+                .expect("Failed to initiate asynchronous disk I/O read");
+
+            // 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.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");
+        }
+    }
+}
+
 /// 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) {
@@ -178,6 +296,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);