fix(driver,unix): set_result for OpenFile & CreateSocket (#701)

* fix(driver,unix): close fd correctly fr OpenFile & CreateSocket

* test: cancel OpenFile

* fix(driver,poll): imports

* fix(driver,poll): set accepted_fd on immediate success

Author: Berrysoft

compio-driver/src/sys/iour/op.rs

@@ -12,7 +12,7 @@ use io_uring::{
     types::{Fd, FsyncFlags},
 };
 use pin_project_lite::pin_project;
-use socket2::{SockAddr, SockAddrStorage, socklen_t};
+use socket2::{SockAddr, SockAddrStorage, Socket as Socket2, socklen_t};
 
 use super::OpCode;
 pub use crate::sys::unix_op::*;
@@ -73,6 +73,12 @@ unsafe impl<S: AsFd> OpCode for OpenFile<S> {
     fn call_blocking(self: Pin<&mut Self>) -> io::Result<usize> {
         self.call()
     }
+
+    unsafe fn set_result(self: Pin<&mut Self>, fd: usize) {
+        // SAFETY: fd is a valid fd returned from kernel
+        let fd = unsafe { OwnedFd::from_raw_fd(fd as _) };
+        *self.project().opened_fd = Some(fd);
+    }
 }
 
 unsafe impl OpCode for CloseFile {
@@ -470,6 +476,12 @@ unsafe impl OpCode for CreateSocket {
             self.protocol
         ))? as _)
     }
+
+    unsafe fn set_result(self: Pin<&mut Self>, fd: usize) {
+        // SAFETY: fd is a valid fd returned from kernel
+        let fd = unsafe { Socket2::from_raw_fd(fd as _) };
+        *self.project().opened_fd = Some(fd);
+    }
 }
 
 unsafe impl<S: AsFd> OpCode for ShutdownSocket<S> {
@@ -511,7 +523,7 @@ unsafe impl<S: AsFd> OpCode for Accept<S> {
 
     unsafe fn set_result(self: Pin<&mut Self>, fd: usize) {
         // SAFETY: fd is a valid fd returned from kernel
-        let fd = unsafe { OwnedFd::from_raw_fd(fd as _) };
+        let fd = unsafe { Socket2::from_raw_fd(fd as _) };
         *self.project().accepted_fd = Some(fd);
     }
 }

compio-driver/src/sys/poll/op.rs

@@ -4,7 +4,7 @@ use std::{
     ffi::CString,
     io,
     marker::PhantomPinned,
-    os::fd::{AsRawFd, FromRawFd, IntoRawFd, OwnedFd},
+    os::fd::{AsRawFd, FromRawFd, OwnedFd},
     pin::Pin,
     task::Poll,
 };
@@ -69,8 +69,10 @@ unsafe impl<S: AsFd> OpCode for OpenFile<S> {
         Ok(Decision::Blocking)
     }
 
-    fn operate(self: Pin<&mut Self>) -> Poll<io::Result<usize>> {
-        Poll::Ready(self.call())
+    fn operate(mut self: Pin<&mut Self>) -> Poll<io::Result<usize>> {
+        let fd = self.as_mut().call()?;
+        *self.project().opened_fd = Some(unsafe { OwnedFd::from_raw_fd(fd as _) });
+        Poll::Ready(Ok(fd))
     }
 }
 
@@ -638,7 +640,8 @@ impl CreateSocket {
             target_os = "cygwin",
         )))]
         socket.set_nonblocking(true)?;
-        Ok(socket.into_raw_fd())
+        *self.project().opened_fd = Some(socket);
+        Ok(fd)
     }
 }
 
@@ -673,40 +676,44 @@ unsafe impl OpCode for CloseSocket {
 }
 
 impl<S: AsFd> Accept<S> {
+    // If the first call succeeds, there won't be another call.
     unsafe fn call(self: Pin<&mut Self>) -> libc::c_int {
         let this = self.project();
-        #[cfg(any(
-            target_os = "android",
-            target_os = "dragonfly",
-            target_os = "freebsd",
-            target_os = "fuchsia",
-            target_os = "illumos",
-            target_os = "linux",
-            target_os = "netbsd",
-            target_os = "openbsd",
-            target_os = "cygwin",
-        ))]
-        unsafe {
-            libc::accept4(
-                this.fd.as_fd().as_raw_fd(),
-                this.buffer as *mut _ as *mut _,
-                this.addr_len,
-                libc::SOCK_NONBLOCK | libc::SOCK_CLOEXEC,
-            )
-        }
-        #[cfg(not(any(
-            target_os = "android",
-            target_os = "dragonfly",
-            target_os = "freebsd",
-            target_os = "fuchsia",
-            target_os = "illumos",
-            target_os = "linux",
-            target_os = "netbsd",
-            target_os = "openbsd",
-            target_os = "cygwin",
-        )))]
-        {
-            || -> io::Result<libc::c_int> {
+        || -> io::Result<libc::c_int> {
+            #[cfg(any(
+                target_os = "android",
+                target_os = "dragonfly",
+                target_os = "freebsd",
+                target_os = "fuchsia",
+                target_os = "illumos",
+                target_os = "linux",
+                target_os = "netbsd",
+                target_os = "openbsd",
+                target_os = "cygwin",
+            ))]
+            {
+                let fd = syscall!(libc::accept4(
+                    this.fd.as_fd().as_raw_fd(),
+                    this.buffer as *mut _ as *mut _,
+                    this.addr_len,
+                    libc::SOCK_NONBLOCK | libc::SOCK_CLOEXEC,
+                ))?;
+                let socket = unsafe { Socket2::from_raw_fd(fd) };
+                *this.accepted_fd = Some(socket);
+                Ok(fd)
+            }
+            #[cfg(not(any(
+                target_os = "android",
+                target_os = "dragonfly",
+                target_os = "freebsd",
+                target_os = "fuchsia",
+                target_os = "illumos",
+                target_os = "linux",
+                target_os = "netbsd",
+                target_os = "openbsd",
+                target_os = "cygwin",
+            )))]
+            {
                 let fd = syscall!(libc::accept(
                     this.fd.as_fd().as_raw_fd(),
                     this.buffer as *mut _ as *mut _,
@@ -715,10 +722,11 @@ impl<S: AsFd> Accept<S> {
                 let socket = unsafe { Socket2::from_raw_fd(fd) };
                 socket.set_cloexec(true)?;
                 socket.set_nonblocking(true)?;
-                Ok(socket.into_raw_fd())
-            }()
-            .unwrap_or(-1)
-        }
+                *this.accepted_fd = Some(socket);
+                Ok(fd)
+            }
+        }()
+        .unwrap_or(-1)
     }
 }
 
@@ -733,13 +741,7 @@ unsafe impl<S: AsFd> OpCode for Accept<S> {
     }
 
     fn operate(mut self: Pin<&mut Self>) -> Poll<io::Result<usize>> {
-        let res = syscall!(break self.as_mut().call());
-        if let Poll::Ready(Ok(fd)) = res {
-            // Safety: we own the fd returned by accept/accept4
-            let fd = unsafe { OwnedFd::from_raw_fd(fd as _) };
-            *self.project().accepted_fd = Some(fd);
-        }
-        res
+        syscall!(break self.as_mut().call())
     }
 }
 

compio-driver/src/sys/unix_op.rs

@@ -27,7 +27,7 @@ use libc::{ftruncate as ftruncate64, off_t as off64_t};
 ))]
 use libc::{ftruncate64, off64_t};
 use pin_project_lite::pin_project;
-use socket2::{SockAddr, SockAddrStorage, socklen_t};
+use socket2::{SockAddr, SockAddrStorage, Socket as Socket2, socklen_t};
 
 use crate::{op::*, sys::aio::*, sys_slice::*, syscall};
 
@@ -47,12 +47,15 @@ impl AsFd for CurrentDir {
     }
 }
 
-/// Open or create a file with flags and mode.
-pub struct OpenFile<S: AsFd> {
-    pub(crate) dirfd: S,
-    pub(crate) path: CString,
-    pub(crate) flags: i32,
-    pub(crate) mode: libc::mode_t,
+pin_project! {
+    /// Open or create a file with flags and mode.
+    pub struct OpenFile<S: AsFd> {
+        pub(crate) dirfd: S,
+        pub(crate) path: CString,
+        pub(crate) flags: i32,
+        pub(crate) mode: libc::mode_t,
+        pub(crate) opened_fd: Option<OwnedFd>,
+    }
 }
 
 impl<S: AsFd> OpenFile<S> {
@@ -63,6 +66,7 @@ impl<S: AsFd> OpenFile<S> {
             path,
             flags,
             mode,
+            opened_fd: None,
         }
     }
 
@@ -76,6 +80,14 @@ impl<S: AsFd> OpenFile<S> {
     }
 }
 
+impl<S: AsFd> IntoInner for OpenFile<S> {
+    type Inner = OwnedFd;
+
+    fn into_inner(self) -> Self::Inner {
+        self.opened_fd.expect("file not opened")
+    }
+}
+
 impl CloseFile {
     pub(crate) fn call(self: Pin<&mut Self>) -> io::Result<usize> {
         Ok(syscall!(libc::close(self.fd.as_fd().as_raw_fd()))? as _)
@@ -462,11 +474,14 @@ impl<S1: AsFd, S2: AsFd> HardLink<S1, S2> {
     }
 }
 
-/// Create a socket.
-pub struct CreateSocket {
-    pub(crate) domain: i32,
-    pub(crate) socket_type: i32,
-    pub(crate) protocol: i32,
+pin_project! {
+    /// Create a socket.
+    pub struct CreateSocket {
+        pub(crate) domain: i32,
+        pub(crate) socket_type: i32,
+        pub(crate) protocol: i32,
+        pub(crate) opened_fd: Option<Socket2>,
+    }
 }
 
 impl CreateSocket {
@@ -476,10 +491,19 @@ impl CreateSocket {
             domain,
             socket_type,
             protocol,
+            opened_fd: None,
         }
     }
 }
 
+impl IntoInner for CreateSocket {
+    type Inner = Socket2;
+
+    fn into_inner(self) -> Self::Inner {
+        self.opened_fd.expect("socket not created")
+    }
+}
+
 impl<S: AsFd> ShutdownSocket<S> {
     pub(crate) fn how(&self) -> i32 {
         match self.how {
@@ -506,7 +530,7 @@ pin_project! {
         pub(crate) fd: S,
         pub(crate) buffer: SockAddrStorage,
         pub(crate) addr_len: socklen_t,
-        pub(crate) accepted_fd: Option<OwnedFd>,
+        pub(crate) accepted_fd: Option<Socket2>,
         _p: PhantomPinned,
     }
 }
@@ -526,9 +550,9 @@ impl<S> Accept<S> {
     }
 
     /// Get the remote address from the inner buffer.
-    pub fn into_addr(mut self) -> SockAddr {
-        std::mem::forget(self.accepted_fd.take());
-        unsafe { SockAddr::new(self.buffer, self.addr_len) }
+    pub fn into_addr(mut self) -> (Socket2, SockAddr) {
+        let socket = self.accepted_fd.take().expect("socket not accepted");
+        (socket, unsafe { SockAddr::new(self.buffer, self.addr_len) })
     }
 }
 

compio-driver/tests/file.rs

@@ -52,8 +52,9 @@ fn open_file(driver: &mut Proactor) -> OwnedFd {
 
 #[cfg(unix)]
 fn open_file(driver: &mut Proactor) -> OwnedFd {
-    use std::{ffi::CString, os::fd::FromRawFd};
+    use std::ffi::CString;
 
+    use compio_buf::IntoInner;
     use compio_driver::op::{CurrentDir, OpenFile};
 
     let op = OpenFile::new(
@@ -62,8 +63,8 @@ fn open_file(driver: &mut Proactor) -> OwnedFd {
         libc::O_CLOEXEC | libc::O_RDONLY,
         0o666,
     );
-    let (fd, _) = push_and_wait(driver, op).unwrap();
-    unsafe { OwnedFd::from_raw_fd(fd as _) }
+    let (_, op) = push_and_wait(driver, op).unwrap();
+    op.into_inner()
 }
 
 fn push_and_wait_extra<O: OpCode + 'static>(

compio-driver/tests/personality.rs

@@ -1,6 +1,6 @@
 #![cfg(io_uring)]
 
-use compio_buf::BufResult;
+use compio_buf::{BufResult, IntoInner};
 use compio_driver::{
     op::{CurrentDir, ReadAt},
     *,
@@ -26,7 +26,7 @@ fn push_and_wait<O: OpCode + 'static>(
 }
 
 fn open_file(driver: &mut Proactor, personality: u16) -> OwnedFd {
-    use std::{ffi::CString, os::fd::FromRawFd};
+    use std::ffi::CString;
 
     use compio_driver::op::OpenFile;
 
@@ -36,8 +36,8 @@ fn open_file(driver: &mut Proactor, personality: u16) -> OwnedFd {
         libc::O_CLOEXEC | libc::O_RDONLY,
         0o666,
     );
-    let (fd, _) = push_and_wait(driver, op, personality).unwrap();
-    unsafe { OwnedFd::from_raw_fd(fd as _) }
+    let (_, op) = push_and_wait(driver, op, personality).unwrap();
+    op.into_inner()
 }
 
 #[test]

compio-fs/src/open_options/unix.rs

@@ -1,9 +1,7 @@
-use std::{io, os::fd::FromRawFd, path::Path};
+use std::{io, path::Path};
 
-use compio_driver::{
-    RawFd,
-    op::{CurrentDir, OpenFile},
-};
+use compio_buf::{IntoInner, buf_try};
+use compio_driver::op::{CurrentDir, OpenFile};
 
 use crate::{File, path_string};
 
@@ -89,7 +87,7 @@ impl OpenOptions {
             | (self.custom_flags as libc::c_int & !libc::O_ACCMODE);
         let p = path_string(p)?;
         let op = OpenFile::new(CurrentDir, p, flags, self.mode);
-        let fd = compio_runtime::submit(op).await.0? as RawFd;
-        File::from_std(unsafe { std::fs::File::from_raw_fd(fd) })
+        let (_, op) = buf_try!(@try compio_runtime::submit(op).await);
+        File::from_std(op.into_inner().into())
     }
 }