uefi: Use RefCell for allocate_buffer and map

tmvkrpxl0 · tmvkrpxl0 · commit d2c59fbb0eac · 2025-06-22T04:26:45.000+09:00
diff --git a/uefi-test-runner/src/proto/pci/mod.rs b/uefi-test-runner/src/proto/pci/mod.rs
@@ -6,4 +6,5 @@ pub fn test() {
     root_bridge::test_io();
     root_bridge::test_buffer();
     root_bridge::test_mapping();
+    root_bridge::test_copy();
 }
diff --git a/uefi-test-runner/src/proto/pci/root_bridge.rs b/uefi-test-runner/src/proto/pci/root_bridge.rs
@@ -1,13 +1,16 @@
 // SPDX-License-Identifier: MIT OR Apache-2.0
 
+use core::cell::RefCell;
 use core::mem;
-use qcell::{QCell, QCellOwner};
-use uefi::Handle;
-use uefi::boot::{OpenProtocolAttributes, OpenProtocolParams, ScopedProtocol, image_handle};
-use uefi::proto::ProtocolPointer;
-use uefi::proto::pci::PciIoAddress;
+use uefi::boot::{image_handle, OpenProtocolAttributes, OpenProtocolParams, ScopedProtocol};
 use uefi::proto::pci::root_bridge::PciRootBridgeIo;
-use uefi_raw::protocol::pci::root_bridge::{PciRootBridgeIoProtocolAttribute, PciRootBridgeIoProtocolOperation, PciRootBridgeIoProtocolWidth};
+use uefi::proto::pci::PciIoAddress;
+use uefi::proto::ProtocolPointer;
+use uefi::Handle;
+use uefi_raw::protocol::pci::root_bridge::{
+    PciRootBridgeIoProtocolAttribute, PciRootBridgeIoProtocolOperation,
+    PciRootBridgeIoProtocolWidth,
+};
 use uefi_raw::table::boot::{MemoryType, PAGE_SIZE};
 
 const RED_HAT_PCI_VENDOR_ID: u16 = 0x1AF4;
@@ -75,14 +78,15 @@ pub fn test_buffer() {
 
     for pci_handle in pci_handles {
         let pci_proto = get_open_protocol::<PciRootBridgeIo>(pci_handle);
-
-        let mut buffer = pci_proto
-            .allocate_buffer::<[u8; 4096]>(
-                MemoryType::BOOT_SERVICES_DATA,
-                None,
-                PciRootBridgeIoProtocolAttribute::PCI_ATTRIBUTE_MEMORY_WRITE_COMBINE,
-            )
-            .unwrap();
+        let pci_proto: RefCell<_> = pci_proto.into();
+
+        let mut buffer = PciRootBridgeIo::allocate_buffer::<[u8; 4096]>(
+            &pci_proto,
+            MemoryType::BOOT_SERVICES_DATA,
+            None,
+            PciRootBridgeIoProtocolAttribute::PCI_ATTRIBUTE_MEMORY_WRITE_COMBINE,
+        )
+        .unwrap();
         let buffer = unsafe {
             buffer.assume_init_mut().fill(0);
             buffer.assume_init()
@@ -96,19 +100,25 @@ pub fn test_mapping() {
 
     for pci_handle in pci_handles {
         let pci_proto = get_open_protocol::<PciRootBridgeIo>(pci_handle);
-
-        let mut buffer = pci_proto
-            .allocate_buffer::<[u8; 4096]>(
-                MemoryType::BOOT_SERVICES_DATA,
-                None,
-                PciRootBridgeIoProtocolAttribute::PCI_ATTRIBUTE_MEMORY_WRITE_COMBINE,
-            )
-            .unwrap();
+        let pci_proto: RefCell<_> = pci_proto.into();
+
+        let mut buffer = PciRootBridgeIo::allocate_buffer::<[u8; 4096]>(
+            &pci_proto,
+            MemoryType::BOOT_SERVICES_DATA,
+            None,
+            PciRootBridgeIoProtocolAttribute::PCI_ATTRIBUTE_MEMORY_WRITE_COMBINE,
+        )
+        .unwrap();
         let buffer = unsafe {
             buffer.assume_init_mut().fill(0);
             buffer.assume_init()
         };
-        let mapped = pci_proto.map(PciRootBridgeIoProtocolOperation::BUS_MASTER_COMMON_BUFFER64, buffer.as_ref());
+
+        let mapped = PciRootBridgeIo::map(
+            &pci_proto,
+            PciRootBridgeIoProtocolOperation::BUS_MASTER_COMMON_BUFFER64,
+            buffer.as_ref(),
+        );
         if mapped.region().device_address == buffer.as_ptr().addr() as u64 {
             info!("This PCI device uses identity mapping");
         } else {
@@ -121,35 +131,48 @@ pub fn test_copy() {
     let pci_handles = uefi::boot::find_handles::<PciRootBridgeIo>().unwrap();
 
     for pci_handle in pci_handles {
-        let mut owner = QCellOwner::new();
-        let item = QCell::new(&owner, get_open_protocol::<PciRootBridgeIo>(pci_handle));
-        let pci_proto = owner.rw(&item);
-
-        let mut src = pci_proto
-            .allocate_buffer::<[u32; 4096 / 4]>(
-                MemoryType::BOOT_SERVICES_DATA,
-                None,
-                PciRootBridgeIoProtocolAttribute::PCI_ATTRIBUTE_MEMORY_WRITE_COMBINE,
-            )
-            .unwrap();
+        let pci_proto = get_open_protocol::<PciRootBridgeIo>(pci_handle);
+        let pci_proto: RefCell<_> = pci_proto.into();
+
+        let mut src = PciRootBridgeIo::allocate_buffer::<[u32; 4096 / 4]>(
+            &pci_proto,
+            MemoryType::BOOT_SERVICES_DATA,
+            None,
+            PciRootBridgeIoProtocolAttribute::PCI_ATTRIBUTE_MEMORY_WRITE_COMBINE,
+        )
+        .unwrap();
         assert_eq!(size_of_val(src.as_ref()), size_of::<[u8; PAGE_SIZE]>());
         let src = unsafe {
             src.assume_init_mut().fill(0xDEADBEEF);
             src.assume_init()
         };
-        let src_mapped = pci_proto.map(PciRootBridgeIoProtocolOperation::BUS_MASTER_READ, src.as_ref());
-
-        let dst = pci_proto
-            .allocate_buffer::<[u32; 4096 / 4]>(
-                MemoryType::BOOT_SERVICES_DATA,
-                None,
-                PciRootBridgeIoProtocolAttribute::PCI_ATTRIBUTE_MEMORY_WRITE_COMBINE,
-            )
-            .unwrap();
+        let src_mapped = PciRootBridgeIo::map(
+            &pci_proto,
+            PciRootBridgeIoProtocolOperation::BUS_MASTER_READ,
+            src.as_ref(),
+        );
+
+        let dst = PciRootBridgeIo::allocate_buffer::<[u32; 4096 / 4]>(
+            &pci_proto,
+            MemoryType::BOOT_SERVICES_DATA,
+            None,
+            PciRootBridgeIoProtocolAttribute::PCI_ATTRIBUTE_MEMORY_WRITE_COMBINE,
+        )
+        .unwrap();
         assert_eq!(size_of_val(dst.as_ref()), size_of::<[u8; PAGE_SIZE]>());
-        let dst_mapped = pci_proto.map(PciRootBridgeIoProtocolOperation::BUS_MASTER_WRITE, dst.as_ref());
-
-        pci_proto.copy(PciRootBridgeIoProtocolWidth::UINT32, dst_mapped.region(), src_mapped.region()).unwrap();
+        let dst_mapped = PciRootBridgeIo::map(
+            &pci_proto,
+            PciRootBridgeIoProtocolOperation::BUS_MASTER_WRITE,
+            dst.as_ref(),
+        );
+
+        PciRootBridgeIo::copy(
+            &pci_proto,
+            PciRootBridgeIoProtocolWidth::UINT32,
+            dst_mapped.region(),
+            src_mapped.region(),
+        )
+        .unwrap();
         drop(dst_mapped);
         let dst = unsafe { dst.assume_init() };
 
diff --git a/uefi/src/proto/pci/buffer.rs b/uefi/src/proto/pci/buffer.rs
@@ -2,22 +2,22 @@
 
 //! Defines wrapper allocated by PCI Root Bridge protocol.
 
+use core::cell::RefCell;
 use core::mem::{ManuallyDrop, MaybeUninit};
 use core::num::NonZeroUsize;
 use core::ops::{Deref, DerefMut};
-use core::ptr;
 use core::ptr::NonNull;
 use log::debug;
 use uefi_raw::Status;
-use uefi_raw::protocol::pci::root_bridge::PciRootBridgeIoProtocol;
+use crate::boot::ScopedProtocol;
+use crate::proto::pci::root_bridge::PciRootBridgeIo;
 
 /// Smart pointer for wrapping owned buffer allocated by PCI Root Bridge protocol.
 #[derive(Debug)]
 pub struct PciBuffer<'p, T> {
     base: NonNull<T>,
     pages: NonZeroUsize,
-    proto_lifetime: &'p (),
-    proto: *const PciRootBridgeIoProtocol,
+    proto: &'p RefCell<ScopedProtocol<PciRootBridgeIo>>,
 }
 
 impl<'p, T> PciBuffer<'p, MaybeUninit<T>> {
@@ -26,12 +26,11 @@ impl<'p, T> PciBuffer<'p, MaybeUninit<T>> {
     /// Passed protocol is stored as a pointer along with its lifetime so that it doesn't
     /// block others from using its mutable functions.
     #[must_use]
-    pub const fn new(base: NonNull<MaybeUninit<T>>, pages: NonZeroUsize, proto: &'p PciRootBridgeIoProtocol) -> Self {
+    pub const fn new(base: NonNull<MaybeUninit<T>>, pages: NonZeroUsize, proto: &'p RefCell<ScopedProtocol<PciRootBridgeIo>>) -> Self {
         Self {
             base,
             pages,
-            proto_lifetime: &(),
-            proto: ptr::from_ref(proto),
+            proto,
         }
     }
 
@@ -45,7 +44,6 @@ impl<'p, T> PciBuffer<'p, MaybeUninit<T>> {
         PciBuffer {
             base: old.base.cast(),
             pages: old.pages,
-            proto_lifetime: old.proto_lifetime,
             proto: old.proto,
         }
     }
@@ -79,10 +77,7 @@ impl<'p, T> DerefMut for PciBuffer<'p, T> {
 
 impl<'p, T> Drop for PciBuffer<'p, T> {
     fn drop(&mut self) {
-        let status = unsafe {
-            let proto = self.proto.as_ref().unwrap();
-            (proto.free_buffer)(proto, self.pages.get(), self.base.as_ptr().cast())
-        };
+        let status = PciRootBridgeIo::free_buffer(&self.proto, self.pages, self.base);
         match status {
             Status::SUCCESS => {
                 debug!(
diff --git a/uefi/src/proto/pci/region.rs b/uefi/src/proto/pci/region.rs
@@ -2,11 +2,14 @@
 
 //! Defines wrapper for region mapped by PCI Root Bridge I/O protocol.
 
+use core::cell::RefCell;
 use core::ffi::c_void;
+use core::marker::PhantomData;
 use core::ptr;
 use log::debug;
 use uefi_raw::Status;
-use uefi_raw::protocol::pci::root_bridge::PciRootBridgeIoProtocol;
+use crate::boot::ScopedProtocol;
+use crate::proto::pci::root_bridge::PciRootBridgeIo;
 
 /// Represents a region of memory mapped by PCI Root Bridge I/O protocol.
 /// The region will be unmapped automatically when it is dropped.
@@ -22,9 +25,9 @@ where
     'p: 'r,
 {
     region: PciRegion,
-    _lifetime_holder: &'r (),
+    _lifetime_holder: PhantomData<&'r ()>,
     key: *const c_void,
-    proto: &'p PciRootBridgeIoProtocol,
+    proto: &'p RefCell<ScopedProtocol<PciRootBridgeIo>>,
 }
 
 /// Represents a region of memory in PCI root bridge memory space.
@@ -45,22 +48,17 @@ impl<'p, 'r> PciMappedRegion<'p, 'r> where 'p: 'r {
         device_address: u64,
         length: usize,
         key: *const c_void,
-        to_map: &'r T,
-        proto: &'p PciRootBridgeIoProtocol,
+        _to_map: &'r T,
+        proto: &'p RefCell<ScopedProtocol<PciRootBridgeIo>>,
     ) -> Self {
-        let _lifetime_holder: &'r () = unsafe {
-            let ptr = ptr::from_ref(to_map);
-            ptr.cast::<()>().as_ref().unwrap()
-        };
-
         let end = device_address + length as u64;
         debug!("Mapped new region [0x{:X}..0x{:X}]", device_address, end);
         Self {
             region: PciRegion {
                 device_address,
                 length,
             },
-            _lifetime_holder,
+            _lifetime_holder: PhantomData,
             key,
             proto,
         }
@@ -79,7 +77,7 @@ impl<'p, 'r> PciMappedRegion<'p, 'r> where 'p: 'r {
 
 impl<'p, 'r> Drop for PciMappedRegion<'p, 'r> {
     fn drop(&mut self) {
-        let status = unsafe { (self.proto.unmap)(self.proto, self.key) };
+        let status = PciRootBridgeIo::unmap(self.proto, self.key);
         match status {
             Status::SUCCESS => {
                 let end = self.region.device_address + self.region.length as u64;
diff --git a/uefi/src/proto/pci/root_bridge.rs b/uefi/src/proto/pci/root_bridge.rs

Original file line number	Diff line number	Diff line change
`@@ -6,4 +6,5 @@ pub fn test() {`
`6`	`6`	`root_bridge::test_io();`
`7`	`7`	`root_bridge::test_buffer();`
`8`	`8`	`root_bridge::test_mapping();`
	`9`	`+ root_bridge::test_copy();`
`9`	`10`	`}`