uefi: Add wrapper for pages allocated by PCI protocol

Author: tmvkrpxl0

uefi-test-runner/src/proto/pci/mod.rs

@@ -3,5 +3,6 @@
 pub mod root_bridge;
 
 pub fn test() {
-    root_bridge::test();
+    root_bridge::test_io();
+    root_bridge::test_buffer();
 }

uefi-test-runner/src/proto/pci/root_bridge.rs

@@ -9,14 +9,16 @@ use uefi::proto::device_path::DevicePath;
 use uefi::proto::device_path::text::{AllowShortcuts, DisplayOnly};
 use uefi::proto::pci::root_bridge::PciRootBridgeIo;
 use uefi::proto::scsi::pass_thru::ExtScsiPassThru;
+use uefi_raw::protocol::pci::root_bridge::PciRootBridgeIoProtocolAttribute;
+use uefi_raw::table::boot::MemoryType;
 
 const RED_HAT_PCI_VENDOR_ID: u16 = 0x1AF4;
 const MASS_STORAGE_CTRL_CLASS_CODE: u8 = 0x1;
 const SATA_CTRL_SUBCLASS_CODE: u8 = 0x6;
 
 const REG_SIZE: u8 = size_of::<u32>() as u8;
 
-pub fn test() {
+pub fn test_io() {
     let pci_handles = uefi::boot::find_handles::<PciRootBridgeIo>().unwrap();
 
     let mut sata_ctrl_cnt = 0;
@@ -88,6 +90,31 @@ pub fn test() {
     }
 }
 
+pub fn test_buffer() {
+    let pci_handles = uefi::boot::find_handles::<PciRootBridgeIo>().unwrap();
+
+    for pci_handle in pci_handles {
+        let pci_proto = get_open_protocol::<PciRootBridgeIo>(pci_handle);
+
+        let buffer = pci_proto
+            .allocate_buffer::<[u8; 4096]>(
+                MemoryType::BOOT_SERVICES_DATA,
+                None,
+                PciRootBridgeIoProtocolAttribute::PCI_ATTRIBUTE_MEMORY_WRITE_COMBINE,
+            )
+            .unwrap();
+        let buffer = unsafe {
+            let buffer = buffer.assume_init();
+            buffer.base_ptr().as_mut().unwrap().fill(0);
+            buffer
+        };
+        assert_eq!(buffer.base_ptr().addr() % 4096, 0);
+        unsafe {
+            assert!(buffer.base_ptr().as_mut().unwrap().iter().all(|v| *v == 0));
+        }
+    }
+}
+
 fn get_open_protocol<P: ProtocolPointer + ?Sized>(handle: Handle) -> ScopedProtocol<P> {
     let open_opts = OpenProtocolParams {
         handle,

uefi/src/proto/pci/root_bridge/buffer.rs

@@ -0,0 +1,104 @@
+// SPDX-License-Identifier: MIT OR Apache-2.0
+
+//! Defines wrapper for pages allocated by PCI Root Bridge protocol.
+
+use crate::StatusExt;
+use core::cell::UnsafeCell;
+use core::fmt::Debug;
+use core::mem::{ManuallyDrop, MaybeUninit};
+use core::num::NonZeroUsize;
+use core::ptr::NonNull;
+use log::{error, trace};
+use uefi_raw::Status;
+use uefi_raw::protocol::pci::root_bridge::PciRootBridgeIoProtocol;
+use uefi_raw::table::boot::PAGE_SIZE;
+
+/// Smart pointer for wrapping owned pages allocated by PCI Root Bridge protocol.
+/// Value stored in this buffer maybe modified by a PCI device.
+///
+/// # Lifetime
+/// `'p` is the lifetime for Protocol.
+///
+/// # Invariant
+/// * Value stored in this memory cannot have a larger alignment requirement
+///   than page size, which is 4096.
+/// * Value stored in this memory cannot be larger than the buffer's size, which is 4096 * `pages`
+#[derive(Debug)]
+pub struct PciBuffer<'p, T> {
+    pub(super) base: NonNull<UnsafeCell<T>>,
+    pub(super) pages: NonZeroUsize,
+    pub(super) proto: &'p PciRootBridgeIoProtocol,
+}
+
+impl<'p, T> PciBuffer<'p, MaybeUninit<T>> {
+    /// Assumes the contents of this buffer have been initialized.
+    ///
+    /// # Safety
+    /// Callers of this function must guarantee that the value stored is valid.
+    #[must_use]
+    pub const unsafe fn assume_init(self) -> PciBuffer<'p, T> {
+        let initialized = PciBuffer {
+            base: self.base.cast(),
+            pages: self.pages,
+            proto: self.proto,
+        };
+        let _ = ManuallyDrop::new(self);
+        initialized
+    }
+}
+
+impl<'p, T> PciBuffer<'p, T> {
+    /// Returns the base pointer of this buffer
+    #[must_use]
+    pub const fn base_ptr(&self) -> *mut T {
+        self.base.as_ptr().cast()
+    }
+
+    /// Returns the number of pages this buffer uses
+    #[must_use]
+    pub const fn pages(&self) -> NonZeroUsize {
+        self.pages
+    }
+
+    /// Returns the size of this buffer in bytes
+    #[must_use]
+    pub const fn bytes_size(&self) -> NonZeroUsize {
+        self.pages
+            .checked_mul(NonZeroUsize::new(PAGE_SIZE).unwrap())
+            .expect("Memory size Overflow")
+    }
+
+    /// Frees underlying memory of this buffer.
+    /// It is recommended to use this over drop implementation.
+    pub fn free(self) -> crate::Result {
+        self.free_inner()
+    }
+
+    fn free_inner(&self) -> crate::Result {
+        unsafe { (self.proto.free_buffer)(self.proto, self.pages.get(), self.base.as_ptr().cast()) }
+            .to_result_with_val(|| {
+                trace!(
+                    "Freed {} pages at 0x{:X}",
+                    self.pages,
+                    self.base.as_ptr().addr()
+                )
+            })
+    }
+}
+
+impl<T> Drop for PciBuffer<'_, T> {
+    fn drop(&mut self) {
+        let Err(status) = self.free_inner() else {
+            return;
+        };
+        match status.status() {
+            Status::SUCCESS => {}
+            Status::INVALID_PARAMETER => {
+                error!("PciBuffer was not created through valid protocol usage!")
+            }
+            etc => {
+                error!("Unexpected error occurred when freeing memory: {:?}", etc);
+            }
+        }
+    }
+}

uefi/src/proto/pci/root_bridge.rs uefi/src/proto/pci/root_bridge/mod.rsuefi/src/proto/pci/root_bridge.rs renamed to uefi/src/proto/pci/root_bridge/mod.rs

@@ -3,19 +3,26 @@
 //! PCI Root Bridge protocol.
 
 use super::{PciIoAddress, PciIoUnit, encode_io_mode_and_unit};
-use crate::StatusExt;
 #[cfg(feature = "alloc")]
 use crate::proto::pci::configuration::QwordAddressSpaceDescriptor;
+use crate::proto::pci::root_bridge::buffer::PciBuffer;
+use crate::{Status, StatusExt};
 #[cfg(feature = "alloc")]
 use alloc::vec::Vec;
 #[cfg(feature = "alloc")]
 use core::ffi::c_void;
+use core::mem::MaybeUninit;
+use core::num::NonZeroUsize;
 use core::ptr;
+use core::ptr::{NonNull, null_mut};
+use log::debug;
 use uefi_macros::unsafe_protocol;
-use uefi_raw::protocol::pci::root_bridge::{PciRootBridgeIoAccess, PciRootBridgeIoProtocol};
+use uefi_raw::protocol::pci::root_bridge::{
+    PciRootBridgeIoAccess, PciRootBridgeIoProtocol, PciRootBridgeIoProtocolAttribute,
+};
+use uefi_raw::table::boot::{AllocateType, MemoryType, PAGE_SIZE};
 
-#[cfg(doc)]
-use crate::Status;
+pub mod buffer;
 
 /// Protocol that provides access to the PCI Root Bridge I/O protocol.
 ///
@@ -51,11 +58,98 @@ impl PciRootBridgeIo {
         unsafe { (self.0.flush)(&mut self.0).to_result() }
     }
 
+    /// Allocates pages suitable for communicating with PCI devices.
+    ///
+    /// # Errors
+    /// - [`Status::INVALID_PARAMETER`] MemoryType is invalid.
+    /// - [`Status::UNSUPPORTED`] Attributes is unsupported. The only legal attribute bits are:
+    ///   - [`PciRootBridgeIoProtocolAttribute::PCI_ATTRIBUTE_MEMORY_WRITE_COMBINE`]
+    ///   - [`PciRootBridgeIoProtocolAttribute::PCI_ATTRIBUTE_MEMORY_CACHED`]
+    ///   - [`PciRootBridgeIoProtocolAttribute::PCI_ATTRIBUTE_DUAL_ADDRESS_CYCLE`]
+    /// - [`Status::OUT_OF_RESOURCES`] The memory pages could not be allocated.
+    pub fn allocate_buffer<T>(
+        &self,
+        memory_type: MemoryType,
+        pages: Option<NonZeroUsize>,
+        attributes: PciRootBridgeIoProtocolAttribute,
+    ) -> crate::Result<PciBuffer<'_, MaybeUninit<T>>> {
+        let original_alignment = align_of::<T>();
+        // TODO switch to const block once it lands on stable. These checks should be done in compile time.
+        assert_ne!(original_alignment, 0);
+        assert!(PAGE_SIZE >= original_alignment);
+        assert_eq!(PAGE_SIZE % original_alignment, 0);
+
+        let alignment = PAGE_SIZE;
+
+        let pages = if let Some(pages) = pages {
+            pages
+        } else {
+            let size = size_of::<T>();
+            assert_ne!(size, 0);
+
+            NonZeroUsize::new(size.div_ceil(alignment)).unwrap()
+        };
+        let size = size_of::<T>();
+        // TODO switch to const block once it lands on stable.
+        assert!(pages.get() * PAGE_SIZE >= size);
+
+        let mut address: *mut T = null_mut();
+        let status = unsafe {
+            (self.0.allocate_buffer)(
+                &self.0,
+                AllocateType(0),
+                memory_type,
+                pages.get(),
+                ptr::from_mut(&mut address).cast(),
+                attributes.bits(),
+            )
+        };
+
+        match status {
+            Status::SUCCESS => {
+                let base = NonNull::new(address.cast()).unwrap();
+                debug!("Allocated {} pages at 0x{:X}", pages.get(), address.addr());
+                Ok(PciBuffer {
+                    base,
+                    pages,
+                    proto: &self.0,
+                })
+            }
+            error
+            @ (Status::INVALID_PARAMETER | Status::UNSUPPORTED | Status::OUT_OF_RESOURCES) => {
+                Err(error.into())
+            }
+            _ => unreachable!(),
+        }
+    }
+
+    /// Allocates pages suitable for communicating with PCI devices and initialize it right away.
+    ///
+    /// # Errors
+    /// Same as [`Self::allocate_buffer`]
+    /// - [`Status::INVALID_PARAMETER`] MemoryType is invalid.
+    /// - [`Status::UNSUPPORTED`] Attributes is unsupported. The only legal attribute bits are:
+    ///   - [`PciRootBridgeIoProtocolAttribute::PCI_ATTRIBUTE_MEMORY_WRITE_COMBINE`]
+    ///   - [`PciRootBridgeIoProtocolAttribute::PCI_ATTRIBUTE_MEMORY_CACHED`]
+    ///   - [`PciRootBridgeIoProtocolAttribute::PCI_ATTRIBUTE_DUAL_ADDRESS_CYCLE`]
+    /// - [`Status::OUT_OF_RESOURCES`] The memory pages could not be allocated.
+    pub fn allocate_buffer_init<T>(
+        &self,
+        memory_type: MemoryType,
+        value: T,
+        attributes: PciRootBridgeIoProtocolAttribute,
+    ) -> crate::Result<PciBuffer<'_, T>> {
+        let buffer = self.allocate_buffer(memory_type, None, attributes)?;
+        unsafe {
+            buffer.base_ptr().write(MaybeUninit::new(value));
+            Ok(buffer.assume_init())
+        }
+    }
+
     // TODO: poll I/O
     // TODO: mem I/O access
     // TODO: io I/O access
     // TODO: map & unmap & copy memory
-    // TODO: buffer management
     // TODO: get/set attributes
 
     /// Retrieves the current resource settings of this PCI root bridge in the form of a set of ACPI resource descriptors.