Implement memory management and mapping functions for PCI Root bridge io protocols

uefi-raw/src/protocol/pci/root_bridge.rs

@@ -2,6 +2,7 @@
 
 use crate::table::boot::{AllocateType, MemoryType};
 use crate::{Handle, PhysicalAddress, Status, newtype_enum};
+use bitflags::bitflags;
 use core::ffi::c_void;
 use uguid::{Guid, guid};
 
@@ -37,6 +38,29 @@ newtype_enum! {
     }
 }
 
+bitflags! {
+    /// Describes PCI I/O Protocol Attribute bitflags specified in UEFI specification.
+    /// https://uefi.org/specs/UEFI/2.10_A/14_Protocols_PCI_Bus_Support.html
+    #[repr(transparent)]
+    #[derive(Copy, Clone, Debug, Eq, PartialEq)]
+    pub struct PciRootBridgeIoProtocolAttribute: u64 {
+        const PCI_ATTRIBUTE_ISA_MOTHERBOARD_IO     = 0x0001;
+        const PCI_ATTRIBUTE_ISA_IO                 = 0x0002;
+        const PCI_ATTRIBUTE_VGA_PALETTE_IO         = 0x0004;
+        const PCI_ATTRIBUTE_VGA_MEMORY             = 0x0008;
+        const PCI_ATTRIBUTE_VGA_IO                 = 0x0010;
+        const PCI_ATTRIBUTE_IDE_PRIMARY_IO         = 0x0020;
+        const PCI_ATTRIBUTE_IDE_SECONDARY_IO       = 0x0040;
+        const PCI_ATTRIBUTE_MEMORY_WRITE_COMBINE   = 0x0080;
+        const PCI_ATTRIBUTE_MEMORY_CACHED          = 0x0800;
+        const PCI_ATTRIBUTE_MEMORY_DISABLE         = 0x1000;
+        const PCI_ATTRIBUTE_DUAL_ADDRESS_CYCLE     = 0x8000;
+        const PCI_ATTRIBUTE_ISA_IO_16              = 0x10000;
+        const PCI_ATTRIBUTE_VGA_PALETTE_IO_16      = 0x20000;
+        const PCI_ATTRIBUTE_VGA_IO_16              = 0x40000;
+    }
+}
+
 #[derive(Debug)]
 #[repr(C)]
 pub struct PciRootBridgeIoAccess {

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

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

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

@@ -9,14 +9,18 @@ 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, PciRootBridgeIoProtocolOperation,
+};
+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 +92,70 @@ 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));
+        }
+    }
+}
+
+pub fn test_mapping() {
+    let pci_handles = uefi::boot::find_handles::<PciRootBridgeIo>().unwrap();
+    const BUFFER_SIZE: usize = 12342;
+
+    for pci_handle in pci_handles {
+        let pci_proto = get_open_protocol::<PciRootBridgeIo>(pci_handle);
+
+        let buffer = pci_proto
+            .allocate_buffer::<[u8; BUFFER_SIZE]>(
+                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
+        };
+
+        let mut mapped_regions = vec![];
+        let mut offset = 0;
+        loop {
+            let (mapped, mapped_size) = pci_proto
+                .map(
+                    PciRootBridgeIoProtocolOperation::BUS_MASTER_COMMON_BUFFER64,
+                    &buffer,
+                    offset,
+                )
+                .unwrap();
+            mapped_regions.push(mapped);
+            offset += mapped_size;
+            if offset == size_of::<[u8; BUFFER_SIZE]>() {
+                break;
+            }
+        }
+    }
+}
+
 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);
+            }
+        }
+    }
+}