Merge pull request #1793 from rust-osdev/bishop-nvme-ub-fix

uefi: Wrap NmvePassThruProtocol in UnsafeCell

Author: phip1611

uefi/src/proto/nvme/pass_thru.rs

@@ -7,6 +7,7 @@ use crate::StatusExt;
 use crate::mem::{AlignedBuffer, PoolAllocation};
 use crate::proto::device_path::PoolDevicePathNode;
 use core::alloc::LayoutError;
+use core::cell::UnsafeCell;
 use core::ptr::{self, NonNull};
 use uefi_macros::unsafe_protocol;
 use uefi_raw::Status;
@@ -40,7 +41,7 @@ pub type NvmeNamespaceId = u32;
 #[derive(Debug)]
 #[repr(transparent)]
 #[unsafe_protocol(NvmExpressPassThruProtocol::GUID)]
-pub struct NvmePassThru(NvmExpressPassThruProtocol);
+pub struct NvmePassThru(UnsafeCell<NvmExpressPassThruProtocol>);
 
 impl NvmePassThru {
     /// Retrieves the mode of the NVMe Pass Thru protocol.
@@ -49,7 +50,7 @@ impl NvmePassThru {
     /// An instance of [`NvmePassThruMode`] describing the NVMe controller's capabilities.
     #[must_use]
     pub fn mode(&self) -> NvmePassThruMode {
-        let mut mode = unsafe { (*self.0.mode).clone() };
+        let mut mode = unsafe { (*(*self.0.get()).mode).clone() };
         mode.io_align = mode.io_align.max(1); // 0 and 1 is the same, says UEFI spec
         mode
     }
@@ -116,15 +117,11 @@ impl NvmePassThru {
 /// Typically, consumer devices only have a single namespace where all the data resides (id 1).
 #[derive(Debug)]
 pub struct NvmeNamespace<'a> {
-    proto: &'a NvmExpressPassThruProtocol,
+    proto: &'a UnsafeCell<NvmExpressPassThruProtocol>,
     namespace_id: NvmeNamespaceId,
 }
 
 impl NvmeNamespace<'_> {
-    const fn proto_mut(&mut self) -> *mut NvmExpressPassThruProtocol {
-        ptr::from_ref(self.proto).cast_mut()
-    }
-
     /// Retrieves the namespace identifier (NSID) associated with this NVMe namespace.
     #[must_use]
     pub const fn namespace_id(&self) -> NvmeNamespaceId {
@@ -138,8 +135,12 @@ impl NvmeNamespace<'_> {
     pub fn path_node(&self) -> crate::Result<PoolDevicePathNode> {
         unsafe {
             let mut path_ptr: *const DevicePathProtocol = ptr::null();
-            (self.proto.build_device_path)(self.proto, self.namespace_id, &mut path_ptr)
-                .to_result()?;
+            ((*self.proto.get()).build_device_path)(
+                self.proto.get(),
+                self.namespace_id,
+                &mut path_ptr,
+            )
+            .to_result()?;
             NonNull::new(path_ptr.cast_mut())
                 .map(|p| PoolDevicePathNode(PoolAllocation::new(p.cast())))
                 .ok_or_else(|| Status::OUT_OF_RESOURCES.into())
@@ -178,8 +179,8 @@ impl NvmeNamespace<'_> {
         req.packet.nvme_cmd = &req.cmd;
         req.packet.nvme_completion = &mut completion;
         unsafe {
-            (self.proto.pass_thru)(
-                self.proto_mut(),
+            ((*self.proto.get()).pass_thru)(
+                self.proto.get(),
                 self.namespace_id,
                 &mut req.packet,
                 ptr::null_mut(),
@@ -195,15 +196,16 @@ impl NvmeNamespace<'_> {
 /// on the NVMe controller.
 #[derive(Debug)]
 pub struct NvmeNamespaceIterator<'a> {
-    proto: &'a NvmExpressPassThruProtocol,
+    proto: &'a UnsafeCell<NvmExpressPassThruProtocol>,
     prev: NvmeNamespaceId,
 }
 
 impl<'a> Iterator for NvmeNamespaceIterator<'a> {
     type Item = NvmeNamespace<'a>;
 
     fn next(&mut self) -> Option<Self::Item> {
-        let result = unsafe { (self.proto.get_next_namespace)(self.proto, &mut self.prev) };
+        let result =
+            unsafe { ((*self.proto.get()).get_next_namespace)(self.proto.get(), &mut self.prev) };
         match result {
             Status::SUCCESS => Some(NvmeNamespace {
                 proto: self.proto,