Poll support for Winsock2.

sunfishcode · sunfishcode · commit 15185a6377fe · 2022-02-03T10:40:37.000-08:00
Winsock2 has a POSIX-like `poll` function named `WSAPoll`, along with
associated `POLL*` flags, so add rustix support for using it.

To support this, fold the `AsSocketAsFd` trait into `AsFd` so that users
can simply use `AsFd` on Windows when working with the Winsock2 APIs.
diff --git a/src/imp/libc/c.rs b/src/imp/libc/c.rs
@@ -26,8 +26,8 @@ pub(crate) use winapi::ctypes::*;
 pub(crate) use winapi::shared::basetsd::SSIZE_T as ssize_t;
 
 pub(crate) use winapi::um::winsock2::{
-    closesocket as close, ioctlsocket as ioctl, linger, SOL_SOCKET, SO_BROADCAST, SO_LINGER,
-    SO_RCVTIMEO, SO_REUSEADDR, SO_SNDTIMEO, SO_TYPE, WSAEACCES as EACCES,
+    closesocket as close, ioctlsocket as ioctl, linger, WSAPoll as poll, SOL_SOCKET, SO_BROADCAST,
+    SO_LINGER, SO_RCVTIMEO, SO_REUSEADDR, SO_SNDTIMEO, SO_TYPE, WSAEACCES as EACCES,
     WSAEADDRINUSE as EADDRINUSE, WSAEADDRNOTAVAIL as EADDRNOTAVAIL,
     WSAEAFNOSUPPORT as EAFNOSUPPORT, WSAEALREADY as EALREADY, WSAEBADF as EBADF,
     WSAECANCELLED as ECANCELED, WSAECONNABORTED as ECONNABORTED, WSAECONNREFUSED as ECONNREFUSED,
@@ -45,7 +45,7 @@ pub(crate) use winapi::um::winsock2::{
     WSAEPROVIDERFAILEDINIT as EPROVIDERFAILEDINIT, WSAEREFUSED as EREFUSED, WSAEREMOTE as EREMOTE,
     WSAESHUTDOWN as ESHUTDOWN, WSAESOCKTNOSUPPORT as ESOCKTNOSUPPORT, WSAESTALE as ESTALE,
     WSAETIMEDOUT as ETIMEDOUT, WSAETOOMANYREFS as ETOOMANYREFS, WSAEUSERS as EUSERS,
-    WSAEWOULDBLOCK as EWOULDBLOCK, WSAEWOULDBLOCK as EAGAIN, *,
+    WSAEWOULDBLOCK as EWOULDBLOCK, WSAEWOULDBLOCK as EAGAIN, WSAPOLLFD as pollfd, *,
 };
 
 // [Documented] values for the `how` argument to `shutdown`.
diff --git a/src/imp/libc/io/mod.rs b/src/imp/libc/io/mod.rs
@@ -2,7 +2,6 @@ mod error;
 #[cfg(not(windows))]
 #[cfg(not(feature = "std"))]
 mod io_slice;
-#[cfg(not(windows))]
 mod poll_fd;
 #[cfg(not(windows))]
 mod types;
@@ -12,11 +11,13 @@ pub(crate) mod syscalls;
 
 #[cfg(windows)]
 pub(crate) mod syscalls {
-    use super::super::conv::{borrowed_fd, ret};
+    use super::super::conv::{borrowed_fd, ret, ret_c_int};
     use super::super::fd::LibcFd;
     use super::c;
     use crate::fd::{BorrowedFd, RawFd};
     use crate::io;
+    use crate::io::PollFd;
+    use core::convert::TryInto;
 
     pub(crate) unsafe fn close(raw_fd: RawFd) {
         let _ = c::close(raw_fd as LibcFd);
@@ -28,6 +29,16 @@ pub(crate) mod syscalls {
             ret(c::ioctl(borrowed_fd(fd), c::FIONBIO, &mut data))
         }
     }
+
+    pub(crate) fn poll(fds: &mut [PollFd<'_>], timeout: c::c_int) -> io::Result<usize> {
+        let nfds = fds
+            .len()
+            .try_into()
+            .map_err(|_convert_err| io::Error::INVAL)?;
+
+        ret_c_int(unsafe { c::poll(fds.as_mut_ptr().cast(), nfds, timeout) })
+            .map(|nready| nready as usize)
+    }
 }
 
 #[cfg(not(windows))]
@@ -36,7 +47,6 @@ pub use io_slice::{IoSlice, IoSliceMut};
 #[cfg(any(target_os = "android", target_os = "linux"))]
 pub mod epoll;
 pub use error::Error;
-#[cfg(not(windows))]
 pub use poll_fd::{PollFd, PollFlags};
 #[cfg(not(any(windows, target_os = "redox", target_os = "wasi")))]
 pub use types::Advice;
diff --git a/src/imp/libc/io/poll_fd.rs b/src/imp/libc/io/poll_fd.rs
@@ -1,8 +1,12 @@
 use super::super::c;
 use super::super::conv::borrowed_fd;
-use crate::fd::{AsFd, AsRawFd, BorrowedFd};
+#[cfg(windows)]
+use super::super::fd::RawFd;
+use super::super::fd::{AsFd, AsRawFd, BorrowedFd, LibcFd};
 use bitflags::bitflags;
 use core::marker::PhantomData;
+#[cfg(windows)]
+use std::fmt;
 
 bitflags! {
     /// `POLL*` flags for use with [`poll`].
@@ -47,13 +51,25 @@ bitflags! {
 ///
 /// [`poll`]: rustix::io::poll
 #[doc(alias = "pollfd")]
-#[derive(Clone, Debug)]
+#[derive(Clone)]
+#[cfg_attr(not(windows), derive(Debug))]
 #[repr(transparent)]
 pub struct PollFd<'fd> {
     pollfd: c::pollfd,
     _phantom: PhantomData<BorrowedFd<'fd>>,
 }
 
+#[cfg(windows)]
+impl<'fd> fmt::Debug for PollFd<'fd> {
+    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
+        fmt.debug_struct("pollfd")
+            .field("fd", &self.pollfd.fd)
+            .field("events", &self.pollfd.events)
+            .field("revents", &self.pollfd.revents)
+            .finish()
+    }
+}
+
 impl<'fd> PollFd<'fd> {
     /// Constructs a new `PollFd` holding `fd` and `events`.
     #[inline]
@@ -64,7 +80,7 @@ impl<'fd> PollFd<'fd> {
     /// Sets the contained file descriptor to `fd`.
     #[inline]
     pub fn set_fd<Fd: AsFd>(&mut self, fd: &'fd Fd) {
-        self.pollfd.fd = fd.as_fd().as_raw_fd();
+        self.pollfd.fd = fd.as_fd().as_raw_fd() as LibcFd;
     }
 
     /// Constructs a new `PollFd` holding `fd` and `events`.
@@ -93,6 +109,7 @@ impl<'fd> PollFd<'fd> {
     }
 }
 
+#[cfg(not(windows))]
 impl<'fd> AsFd for PollFd<'fd> {
     #[inline]
     fn as_fd(&self) -> BorrowedFd<'_> {
@@ -103,3 +120,15 @@ impl<'fd> AsFd for PollFd<'fd> {
         unsafe { BorrowedFd::borrow_raw_fd(self.pollfd.fd) }
     }
 }
+
+#[cfg(windows)]
+impl<'fd> io_lifetimes::AsSocket for PollFd<'fd> {
+    #[inline]
+    fn as_socket(&self) -> BorrowedFd<'_> {
+        // Safety:
+        //
+        // Our constructors and `set_fd` require `pollfd.fd` to be valid
+        // for the `fd lifetime.
+        unsafe { BorrowedFd::borrow_raw_socket(self.pollfd.fd as RawFd) }
+    }
+}
diff --git a/src/imp/libc/io_lifetimes.rs b/src/imp/libc/io_lifetimes.rs
@@ -62,15 +62,11 @@ impl<T: std::os::windows::io::FromRawSocket> FromRawFd for T {
 /// A version of [`AsFd`] for use with Winsock2 API.
 ///
 /// [`AsFd`]: https://doc.rust-lang.org/stable/std/os/unix/io/trait.AsFd.html
-pub use io_lifetimes::AsSocket as AsFd;
-
-/// We define `AsFd` as an alias for `AsSocket`, but that doesn't provide
-/// an `as_fd` function. This trait adapts an `AsSocket` implementation to
-/// provide `as_fd` using a blanket implementation.
-pub(crate) trait AsSocketAsFd {
+pub trait AsFd {
+    /// An `as_fd` function for Winsock2, where a `Fd` is a `Socket`.
     fn as_fd(&self) -> BorrowedFd;
 }
-impl<T: io_lifetimes::AsSocket> AsSocketAsFd for T {
+impl<T: io_lifetimes::AsSocket> AsFd for T {
     #[inline]
     fn as_fd(&self) -> BorrowedFd {
         self.as_socket()
diff --git a/src/io/ioctl.rs b/src/io/ioctl.rs
@@ -7,8 +7,6 @@
 use crate::io::{Termios, Winsize};
 use crate::{imp, io};
 use imp::fd::AsFd;
-#[cfg(windows)]
-use imp::fd::AsSocketAsFd;
 
 /// `ioctl(fd, TCGETS)`—Get terminal attributes.
 ///
diff --git a/src/io/mod.rs b/src/io/mod.rs
@@ -26,7 +26,6 @@ mod msync;
 mod owned_fd;
 #[cfg(not(any(windows, target_os = "wasi")))]
 mod pipe;
-#[cfg(not(windows))]
 mod poll;
 #[cfg(all(feature = "procfs", any(target_os = "android", target_os = "linux")))]
 mod procfs;
@@ -82,7 +81,6 @@ pub use owned_fd::OwnedFd;
 pub use pipe::pipe;
 #[cfg(not(any(windows, target_os = "ios", target_os = "macos", target_os = "wasi")))]
 pub use pipe::{pipe_with, PipeFlags};
-#[cfg(not(windows))]
 pub use poll::{poll, PollFd, PollFlags};
 #[cfg(all(feature = "procfs", any(target_os = "android", target_os = "linux")))]
 pub use procfs::{proc_self_fd, proc_self_fdinfo_fd, proc_self_maps, proc_self_pagemap};
diff --git a/src/io/owned_fd.rs b/src/io/owned_fd.rs
@@ -11,8 +11,6 @@
 //! file descriptor and close it ourselves.
 #![allow(unsafe_code)]
 
-#[cfg(windows)]
-use crate::imp::fd::AsSocketAsFd;
 use crate::imp::fd::{AsFd, AsRawFd, BorrowedFd, FromRawFd, IntoRawFd, RawFd};
 #[cfg(all(not(io_lifetimes_use_std), feature = "std"))]
 use crate::imp::fd::{FromFd, IntoFd};
@@ -29,14 +27,16 @@ pub struct OwnedFd {
     inner: ManuallyDrop<crate::imp::fd::OwnedFd>,
 }
 
+#[cfg(not(windows))]
 impl AsFd for OwnedFd {
-    #[cfg(not(windows))]
     #[inline]
     fn as_fd(&self) -> BorrowedFd<'_> {
         self.inner.as_fd()
     }
+}
 
-    #[cfg(windows)]
+#[cfg(windows)]
+impl io_lifetimes::AsSocket for OwnedFd {
     #[inline]
     fn as_socket(&self) -> BorrowedFd<'_> {
         self.inner.as_socket()
diff --git a/src/net/send_recv.rs b/src/net/send_recv.rs
@@ -5,8 +5,6 @@ use crate::net::SocketAddrUnix;
 use crate::net::{SocketAddrAny, SocketAddrV4, SocketAddrV6};
 use crate::{imp, io};
 use imp::fd::AsFd;
-#[cfg(windows)]
-use imp::fd::AsSocketAsFd;
 
 pub use imp::net::{RecvFlags, SendFlags};
 
diff --git a/src/net/socket.rs b/src/net/socket.rs
@@ -4,8 +4,6 @@ use crate::io::{self, OwnedFd};
 use crate::net::SocketAddrUnix;
 use crate::net::{SocketAddrAny, SocketAddrV4, SocketAddrV6};
 use imp::fd::AsFd;
-#[cfg(windows)]
-use imp::fd::AsSocketAsFd;
 
 pub use imp::net::{AcceptFlags, AddressFamily, Protocol, Shutdown, SocketFlags, SocketType};
 
diff --git a/src/net/sockopt.rs b/src/net/sockopt.rs
@@ -10,8 +10,6 @@ use crate::net::{Ipv4Addr, Ipv6Addr, SocketType};
 use crate::{imp, io};
 use core::time::Duration;
 use imp::fd::AsFd;
-#[cfg(windows)]
-use imp::fd::AsSocketAsFd;
 
 pub use imp::net::Timeout;
 
diff --git a/tests/io/main.rs b/tests/io/main.rs
@@ -18,7 +18,6 @@ mod from_into;
 mod isatty;
 #[cfg(not(windows))]
 mod mmap;
-#[cfg(not(windows))]
 mod poll;
 #[cfg(all(feature = "procfs", any(target_os = "android", target_os = "linux")))]
 mod procfs;
diff --git a/tests/io/poll.rs b/tests/io/poll.rs
@@ -1,8 +1,13 @@
 use rustix::fd::{AsFd, AsRawFd, FromRawFd, IntoRawFd, OwnedFd};
-use rustix::io::{pipe, poll, read, with_retrying, write, PollFd, PollFlags};
+#[cfg(not(windows))]
+use rustix::io::{poll, with_retrying};
+use rustix::io::{PollFd, PollFlags};
 
+#[cfg(not(windows))]
 #[test]
 fn test_poll() {
+    use rustix::io::{pipe, read, write};
+
     // Create a pipe.
     let (reader, writer) = pipe().unwrap();
     let mut poll_fds = [PollFd::new(&reader, PollFlags::IN)];
diff --git a/tests/net/main.rs b/tests/net/main.rs
@@ -5,6 +5,7 @@
 #![cfg_attr(io_lifetimes_use_std, feature(io_safety))]
 
 mod addr;
+mod poll;
 #[cfg(not(windows))]
 mod unix;
 mod v4;
diff --git a/tests/net/poll.rs b/tests/net/poll.rs
@@ -0,0 +1,121 @@
+//! The same as v6.rs, but with `poll` calls.
+
+#![cfg(not(any(target_os = "redox", target_os = "wasi")))]
+
+use rustix::io::{poll, PollFd, PollFlags};
+use rustix::net::{
+    accept, bind_v6, connect_v6, getsockname, listen, recv, send, socket, AddressFamily, Ipv6Addr,
+    Protocol, RecvFlags, SendFlags, SocketAddrAny, SocketAddrV6, SocketType,
+};
+use std::sync::{Arc, Condvar, Mutex};
+use std::thread;
+
+const BUFFER_SIZE: usize = 20;
+
+fn server(ready: Arc<(Mutex<u16>, Condvar)>) {
+    let connection_socket = socket(
+        AddressFamily::INET6,
+        SocketType::STREAM,
+        Protocol::default(),
+    )
+    .unwrap();
+
+    let name = SocketAddrV6::new(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1), 0, 0, 0);
+    bind_v6(&connection_socket, &name).unwrap();
+
+    let who = match getsockname(&connection_socket).unwrap() {
+        SocketAddrAny::V6(addr) => addr,
+        _ => panic!(),
+    };
+
+    listen(&connection_socket, 1).unwrap();
+
+    {
+        let (lock, cvar) = &*ready;
+        let mut port = lock.lock().unwrap();
+        *port = who.port();
+        cvar.notify_all();
+    }
+
+    let mut buffer = vec![0; BUFFER_SIZE];
+    let data_socket = accept(&connection_socket).unwrap();
+
+    let mut fds = [PollFd::new(&data_socket, PollFlags::IN)];
+    assert_eq!(poll(&mut fds, -1).unwrap(), 1);
+    assert!(fds[0].revents().intersects(PollFlags::IN));
+    assert!(!fds[0].revents().intersects(PollFlags::OUT));
+
+    let nread = recv(&data_socket, &mut buffer, RecvFlags::empty()).unwrap();
+    assert_eq!(String::from_utf8_lossy(&buffer[..nread]), "hello, world");
+
+    let mut fds = [PollFd::new(&data_socket, PollFlags::OUT)];
+    assert_eq!(poll(&mut fds, -1).unwrap(), 1);
+    assert!(!fds[0].revents().intersects(PollFlags::IN));
+    assert!(fds[0].revents().intersects(PollFlags::OUT));
+
+    send(&data_socket, b"goodnight, moon", SendFlags::empty()).unwrap();
+}
+
+fn client(ready: Arc<(Mutex<u16>, Condvar)>) {
+    let port = {
+        let (lock, cvar) = &*ready;
+        let mut port = lock.lock().unwrap();
+        while *port == 0 {
+            port = cvar.wait(port).unwrap();
+        }
+        *port
+    };
+
+    let addr = SocketAddrV6::new(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1), port, 0, 0);
+    let mut buffer = vec![0; BUFFER_SIZE];
+
+    let data_socket = socket(
+        AddressFamily::INET6,
+        SocketType::STREAM,
+        Protocol::default(),
+    )
+    .unwrap();
+    connect_v6(&data_socket, &addr).unwrap();
+
+    let mut fds = [PollFd::new(&data_socket, PollFlags::OUT)];
+    assert_eq!(poll(&mut fds, -1).unwrap(), 1);
+    assert!(!fds[0].revents().intersects(PollFlags::IN));
+    assert!(fds[0].revents().intersects(PollFlags::OUT));
+
+    send(&data_socket, b"hello, world", SendFlags::empty()).unwrap();
+
+    let mut fds = [PollFd::new(&data_socket, PollFlags::IN)];
+    assert_eq!(poll(&mut fds, -1).unwrap(), 1);
+    assert!(fds[0].revents().intersects(PollFlags::IN));
+    assert!(!fds[0].revents().intersects(PollFlags::OUT));
+
+    let nread = recv(&data_socket, &mut buffer, RecvFlags::empty()).unwrap();
+    assert_eq!(String::from_utf8_lossy(&buffer[..nread]), "goodnight, moon");
+}
+
+#[test]
+fn test_v6() {
+    #[cfg(windows)]
+    rustix::net::wsa_startup().unwrap();
+
+    let ready = Arc::new((Mutex::new(0_u16), Condvar::new()));
+    let ready_clone = Arc::clone(&ready);
+
+    let server = thread::Builder::new()
+        .name("server".to_string())
+        .spawn(move || {
+            server(ready);
+        })
+        .unwrap();
+    let client = thread::Builder::new()
+        .name("client".to_string())
+        .spawn(move || {
+            client(ready_clone);
+        })
+        .unwrap();
+    client.join().unwrap();
+    server.join().unwrap();
+
+    #[cfg(windows)]
+    rustix::net::wsa_cleanup().unwrap();
+}
