uefi: Implement memory management and mapping functions

Author: tmvkrpxl0

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/src/proto/pci/root_bridge/page.rs uefi/src/proto/pci/root_bridge/buffer.rsuefi/src/proto/pci/root_bridge/page.rs renamed to uefi/src/proto/pci/root_bridge/buffer.rs

@@ -1,22 +1,33 @@
 // 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::ops::{Deref, DerefMut};
 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.
+///
+/// # Safety
+/// 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<T>,
+    pub(super) base: NonNull<UnsafeCell<T>>,
     pub(super) pages: NonZeroUsize,
     pub(super) proto: &'p PciRootBridgeIoProtocol,
 }
@@ -39,63 +50,56 @@ impl<'p, T> PciBuffer<'p, MaybeUninit<T>> {
 }
 
 impl<'p, T> PciBuffer<'p, T> {
-    /// Returns the base address of this buffer
+    /// Returns the base pointer of this buffer
     #[must_use]
-    pub fn base(&self) -> NonZeroUsize {
-        self.base.addr()
+    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
     }
-}
 
-impl<T> AsRef<T> for PciBuffer<'_, T> {
-    fn as_ref(&self) -> &T {
-        unsafe { self.base.as_ref() }
+    /// 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")
     }
-}
 
-impl<T> AsMut<T> for PciBuffer<'_, T> {
-    fn as_mut(&mut self) -> &mut T {
-        unsafe { self.base.as_mut() }
+    /// Frees underlying memory of this buffer.
+    /// It is recommended to use this over drop implementation.
+    pub fn free(self) -> crate::Result {
+        self.free_inner()
     }
-}
 
-impl<T> Deref for PciBuffer<'_, T> {
-    type Target = T;
-
-    fn deref(&self) -> &Self::Target {
-        self.as_ref()
-    }
-}
-
-impl<T> DerefMut for PciBuffer<'_, T> {
-    fn deref_mut(&mut self) -> &mut Self::Target {
-        self.as_mut()
+    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 status = unsafe {
-            (self.proto.free_buffer)(self.proto, self.pages.get(), self.base.as_ptr().cast())
+        let Err(status) = self.free_inner() else {
+            return;
         };
-        match status {
-            Status::SUCCESS => {
-                trace!(
-                    "Freed {} pages at 0x{:X}",
-                    self.pages.get(),
-                    self.base.as_ptr().addr()
-                );
-            }
+        match status.status() {
+            Status::SUCCESS => {}
             Status::INVALID_PARAMETER => {
                 error!("PciBuffer was not created through valid protocol usage!")
             }
             etc => {
-                error!("Failed to free PciBuffer: {:?}", etc);
+                error!("Unexpected error occurred when freeing memory: {:?}", etc);
             }
         }
     }

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

@@ -3,20 +3,29 @@
 //! PCI Root Bridge protocol.
 
 use super::{PciIoAddress, PciIoUnit, encode_io_mode_and_unit};
-#[cfg(doc)]
 use crate::Status;
 use crate::StatusExt;
 #[cfg(feature = "alloc")]
 use crate::proto::pci::configuration::QwordAddressSpaceDescriptor;
+use crate::proto::pci::root_bridge::buffer::PciBuffer;
+use crate::proto::pci::root_bridge::region::PciMappedRegion;
 #[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,
+    PciRootBridgeIoProtocolOperation,
+};
+use uefi_raw::table::boot::{AllocateType, MemoryType, PAGE_SIZE};
 
-pub mod page;
+pub mod buffer;
 pub mod region;
 
 /// Protocol that provides access to the PCI Root Bridge I/O protocol.
@@ -53,11 +62,142 @@ 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())
+        }
+    }
+
+    /// Map the given buffer into a PCI Controller-specific address
+    /// so that devices can read system memory through it.
+    ///
+    /// # Arguments
+    /// - `operation` - Indicates if bus master is going to read, write, or do both to the buffer.
+    /// - `to_map` - Buffer to map.
+    ///
+    /// # Returns
+    /// An mapped region and unmapped bytes. It can map up to one DMA operation. Meaning large values
+    /// will require multiple calls to this function.
+    pub fn map<'p: 'r, 'r, T>(
+        &'p self,
+        operation: PciRootBridgeIoProtocolOperation,
+        to_map: &'r PciBuffer<'p, T>,
+    ) -> crate::Result<(PciMappedRegion<'p, 'r, T>, usize)> {
+        let host_address = to_map.base_ptr();
+        let mut bytes = size_of::<T>();
+        let requested_bytes = bytes;
+        let mut device_address = 0usize;
+        let mut mapping: *mut c_void = null_mut();
+
+        let status = unsafe {
+            (self.0.map)(
+                &self.0,
+                operation,
+                host_address.cast(),
+                ptr::from_mut(&mut bytes),
+                ptr::from_mut(&mut device_address).cast(),
+                ptr::from_mut(&mut mapping),
+            )
+        };
+
+        status.to_result_with_val(|| {
+            let left_over = requested_bytes - bytes;
+            let region = PciMappedRegion {
+                host: to_map,
+                length: 0,
+                device_address,
+                key: mapping,
+                proto: &self.0,
+            };
+            (region, left_over)
+        })
+    }
+
     // 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.