fix Buffer ownership issue

Author: mox692

tokio/src/fs/mocks.rs

@@ -2,6 +2,8 @@
 use mockall::mock;
 
 use crate::sync::oneshot;
+#[cfg(all(test, unix))]
+use std::os::fd::{AsRawFd, FromRawFd, OwnedFd};
 use std::{
     cell::RefCell,
     collections::VecDeque,
@@ -89,6 +91,14 @@ impl Write for &'_ MockFile {
     }
 }
 
+#[cfg(all(test, unix))]
+impl From<MockFile> for OwnedFd {
+    #[inline]
+    fn from(file: MockFile) -> OwnedFd {
+        unsafe { OwnedFd::from_raw_fd(file.as_raw_fd()) }
+    }
+}
+
 tokio_thread_local! {
     static QUEUE: RefCell<VecDeque<Box<dyn FnOnce() + Send>>> = RefCell::new(VecDeque::new())
 }

tokio/src/fs/write.rs

@@ -40,9 +40,9 @@ pub async fn write(path: impl AsRef<Path>, contents: impl AsRef<[u8]>) -> io::Re
 }
 
 #[cfg(all(tokio_uring, feature = "rt", feature = "fs", target_os = "linux"))]
-async fn write_uring(path: &Path, contents: OwnedBuf) -> io::Result<()> {
+async fn write_uring(path: &Path, mut buf: OwnedBuf) -> io::Result<()> {
     use crate::{fs::OpenOptions, runtime::driver::op::Op};
-    use std::os::fd::AsFd;
+    use std::os::fd::OwnedFd;
 
     let file = OpenOptions::new()
         .write(true)
@@ -51,25 +51,27 @@ async fn write_uring(path: &Path, contents: OwnedBuf) -> io::Result<()> {
         .open(path)
         .await?;
 
-    let fd = file.as_fd();
+    let mut fd: OwnedFd = file
+        .try_into_std()
+        .expect("unexpected in-flight operation detected")
+        .into();
 
-    let mut pos = 0;
-    let mut buf = contents.as_ref();
-    while !buf.is_empty() {
-        // SAFETY:
-        // If the operation completes successfully, `fd` and `buf` are still
-        // alive within the scope of this function, so remain valid.
-        //
-        // In the case of cancellation, local variables within the scope of
-        // this `async fn` are dropped in the reverse order of their declaration.
-        // Therefore, `Op` is dropped before `fd` and `buf`, ensuring that the
-        // operation finishes gracefully before these resources are dropped.
-        let n = unsafe { Op::write_at(fd, buf, pos) }?.await? as usize;
+    let total: usize = buf.len();
+    let mut offset: usize = 0;
+
+    while offset < total {
+        // There is a cap on how many bytes we can write in a single uring write operation.
+        // ref: https://github.com/axboe/liburing/discussions/497
+        let len = std::cmp::min(total - offset, u32::MAX as usize) as u32;
+
+        let (n, _buf, _fd) = Op::write_at(fd, buf, offset, len, offset as u64)?.await?;
         if n == 0 {
             return Err(io::ErrorKind::WriteZero.into());
         }
-        buf = &buf[n..];
-        pos += n as u64;
+
+        buf = _buf;
+        fd = _fd;
+        offset += n as usize;
     }
 
     Ok(())

tokio/src/io/uring/write.rs

@@ -1,18 +1,23 @@
-use crate::runtime::driver::op::{CancelData, Cancellable, Completable, CqeResult, Op};
+use crate::{
+    runtime::driver::op::{CancelData, Cancellable, Completable, CqeResult, Op},
+    util::as_ref::OwnedBuf,
+};
 use io_uring::{opcode, types};
 use std::{
-    cmp, io,
-    os::fd::{AsRawFd, BorrowedFd},
+    io,
+    os::fd::{AsRawFd, OwnedFd},
 };
 
 #[derive(Debug)]
-pub(crate) struct Write;
+pub(crate) struct Write {
+    buf: OwnedBuf,
+    fd: OwnedFd,
+}
 
 impl Completable for Write {
-    // The number of bytes written.
-    type Output = u32;
+    type Output = (u32, OwnedBuf, OwnedFd);
     fn complete(self, cqe: CqeResult) -> io::Result<Self::Output> {
-        cqe.result
+        Ok((cqe.result?, self.buf, self.fd))
     }
 }
 
@@ -23,23 +28,29 @@ impl Cancellable for Write {
 }
 
 impl Op<Write> {
-    /// # SAFETY
-    ///
-    /// The caller must ensure that `fd` and `buf` remain valid until the
-    /// operation finishes (or gets cancelled) and the `Op::drop` completes.
-    /// Otherwise, the kernel could access freed memory, breaking soundness.
-    pub(crate) unsafe fn write_at(fd: BorrowedFd<'_>, buf: &[u8], offset: u64) -> io::Result<Self> {
-        // There is a cap on how many bytes we can write in a single uring write operation.
-        // ref: https://github.com/axboe/liburing/discussions/497
-        let len: u32 = cmp::min(buf.len(), u32::MAX as usize) as u32;
+    /// Issue a write that starts at `buf_offset` within `buf` and writes `len` bytes
+    /// into `file` at `file_offset`.
+    pub(crate) fn write_at(
+        fd: OwnedFd,
+        buf: OwnedBuf,
+        buf_offset: usize,
+        len: u32,
+        file_offset: u64,
+    ) -> io::Result<Self> {
+        // Check if `buf_offset` stays in bounds of the allocation
+        debug_assert!(buf_offset + len as usize <= buf.len());
 
-        let sqe = opcode::Write::new(types::Fd(fd.as_raw_fd()), buf.as_ptr(), len)
-            .offset(offset)
+        // SAFETY:
+        // - `buf_offset` stays in bounds of the allocation.
+        // - `buf` is derived from an actual allocation, and the entire memory
+        //    range is in bounds of that allocation.
+        let ptr = unsafe { buf.as_ptr().add(buf_offset) };
+        let sqe = opcode::Write::new(types::Fd(fd.as_raw_fd()), ptr, len)
+            .offset(file_offset)
             .build();
 
-        // SAFETY: `fd` and `buf` valid until the operation completes or gets cancelled
-        // and the `Op::drop` completes.
-        let op = unsafe { Op::new(sqe, Write) };
+        // SAFETY: parameters of the entry is valid until this operation completes.
+        let op = unsafe { Op::new(sqe, Write { buf, fd }) };
         Ok(op)
     }
 }

tokio/src/runtime/driver/op.rs

@@ -17,7 +17,11 @@ pub(crate) enum CancelData {
         // so `#[allow(dead_code)]` is needed.
         #[allow(dead_code)] Open,
     ),
-    Write(Write),
+    Write(
+        // This field isn't accessed directly, but it holds cancellation data,
+        // so `#[allow(dead_code)]` is needed.
+        #[allow(dead_code)] Write,
+    ),
 }
 
 #[derive(Debug)]

tokio/src/util/as_ref.rs

@@ -7,6 +7,22 @@ pub(crate) enum OwnedBuf {
     Bytes(bytes::Bytes),
 }
 
+impl OwnedBuf {
+    cfg_tokio_uring! {
+        pub(crate) fn len(&self) -> usize {
+            match self {
+                Self::Vec(vec) => vec.len(),
+                #[cfg(feature = "io-util")]
+                Self::Bytes(bytes) => bytes.len(),
+            }
+        }
+
+        pub(crate) fn as_ptr(&self) -> *const u8 {
+            self.as_ref().as_ptr()
+        }
+    }
+}
+
 impl AsRef<[u8]> for OwnedBuf {
     fn as_ref(&self) -> &[u8] {
         match self {