fix(wasi): enable all WasiFiles to be pollable (#3913)

Currently, the use of the downcast method means that you have to use one
of the hard-coded types. But Enarx needs to define its own `WasiFile`
implementations. This works fine, except the resulting files cannot be
used in poll because they aren't part of the hard-coded list.

Replace this with an accessor method for the pollable type in
`WasiFile`. Because we provide a default implementation of the method
and manually implement it on all the hard-coded types, this is backwards
compatible.

Signed-off-by: Nathaniel McCallum <nathaniel@profian.com>

Cargo.lock

@@ -3204,6 +3204,7 @@ dependencies = [
  "bitflags",
  "cap-rand",
  "cap-std",
+ "io-extras",
  "rustix",
  "thiserror",
  "tracing",

crates/wasi-common/Cargo.toml

@@ -30,6 +30,7 @@ bitflags = "1.2"
 rustix = "0.33.0"
 
 [target.'cfg(windows)'.dependencies]
+io-extras = "0.13.2"
 winapi = "0.3"
 
 [badges]

crates/wasi-common/cap-std-sync/src/file.rs

@@ -26,6 +26,15 @@ impl WasiFile for File {
     fn as_any(&self) -> &dyn Any {
         self
     }
+    #[cfg(unix)]
+    fn pollable(&self) -> Option<rustix::fd::BorrowedFd> {
+        Some(self.0.as_fd())
+    }
+
+    #[cfg(windows)]
+    fn pollable(&self) -> Option<io_extras::os::windows::RawHandleOrSocket> {
+        Some(self.0.as_raw_handle_or_socket())
+    }
     async fn datasync(&mut self) -> Result<(), Error> {
         self.0.sync_data()?;
         Ok(())

crates/wasi-common/cap-std-sync/src/net.rs

@@ -85,6 +85,15 @@ macro_rules! wasi_listen_write_impl {
             fn as_any(&self) -> &dyn Any {
                 self
             }
+            #[cfg(unix)]
+            fn pollable(&self) -> Option<rustix::fd::BorrowedFd> {
+                Some(self.0.as_fd())
+            }
+
+            #[cfg(windows)]
+            fn pollable(&self) -> Option<io_extras::os::windows::RawHandleOrSocket> {
+                Some(self.0.as_raw_handle_or_socket())
+            }
             async fn sock_accept(&mut self, fdflags: FdFlags) -> Result<Box<dyn WasiFile>, Error> {
                 let (stream, _) = self.0.accept()?;
                 let mut stream = <$stream>::from_cap_std(stream);
@@ -170,6 +179,15 @@ macro_rules! wasi_stream_write_impl {
             fn as_any(&self) -> &dyn Any {
                 self
             }
+            #[cfg(unix)]
+            fn pollable(&self) -> Option<rustix::fd::BorrowedFd> {
+                Some(self.0.as_fd())
+            }
+
+            #[cfg(windows)]
+            fn pollable(&self) -> Option<io_extras::os::windows::RawHandleOrSocket> {
+                Some(self.0.as_raw_handle_or_socket())
+            }
             async fn get_filetype(&mut self) -> Result<FileType, Error> {
                 Ok(FileType::SocketStream)
             }

crates/wasi-common/cap-std-sync/src/sched/unix.rs

@@ -1,15 +1,8 @@
 use cap_std::time::Duration;
-use io_lifetimes::{AsFd, BorrowedFd};
 use rustix::io::{PollFd, PollFlags};
 use std::convert::TryInto;
-use wasi_common::{
+use wasi_common::sched::subscription::{RwEventFlags, Subscription};
-    file::WasiFile,
+use wasi_common::{sched::Poll, Error, ErrorExt};
-    sched::{
-        subscription::{RwEventFlags, Subscription},
-        Poll,
-    },
-    Error, ErrorExt,
-};
 
 pub async fn poll_oneoff<'a>(poll: &mut Poll<'a>) -> Result<(), Error> {
     if poll.is_empty() {
@@ -19,16 +12,18 @@ pub async fn poll_oneoff<'a>(poll: &mut Poll<'a>) -> Result<(), Error> {
     for s in poll.rw_subscriptions() {
         match s {
             Subscription::Read(f) => {
-                let fd = wasi_file_fd(f.file).ok_or(
+                let fd = f
-                    Error::invalid_argument().context("read subscription fd downcast failed"),
+                    .file
-                )?;
+                    .pollable()
+                    .ok_or(Error::invalid_argument().context("file is not pollable"))?;
                 pollfds.push(PollFd::from_borrowed_fd(fd, PollFlags::IN));
             }
 
             Subscription::Write(f) => {
-                let fd = wasi_file_fd(f.file).ok_or(
+                let fd = f
-                    Error::invalid_argument().context("write subscription fd downcast failed"),
+                    .file
-                )?;
+                    .pollable()
+                    .ok_or(Error::invalid_argument().context("file is not pollable"))?;
                 pollfds.push(PollFd::from_borrowed_fd(fd, PollFlags::OUT));
             }
             Subscription::MonotonicClock { .. } => unreachable!(),
@@ -85,30 +80,3 @@ pub async fn poll_oneoff<'a>(poll: &mut Poll<'a>) -> Result<(), Error> {
     }
     Ok(())
 }
-
-fn wasi_file_fd(f: &dyn WasiFile) -> Option<BorrowedFd<'_>> {
-    let a = f.as_any();
-    if a.is::<crate::file::File>() {
-        Some(a.downcast_ref::<crate::file::File>().unwrap().as_fd())
-    } else if a.is::<crate::net::TcpStream>() {
-        Some(a.downcast_ref::<crate::net::TcpStream>().unwrap().as_fd())
-    } else if a.is::<crate::net::TcpListener>() {
-        Some(a.downcast_ref::<crate::net::TcpListener>().unwrap().as_fd())
-    } else if a.is::<crate::net::UnixStream>() {
-        Some(a.downcast_ref::<crate::net::UnixStream>().unwrap().as_fd())
-    } else if a.is::<crate::net::UnixListener>() {
-        Some(
-            a.downcast_ref::<crate::net::UnixListener>()
-                .unwrap()
-                .as_fd(),
-        )
-    } else if a.is::<crate::stdio::Stdin>() {
-        Some(a.downcast_ref::<crate::stdio::Stdin>().unwrap().as_fd())
-    } else if a.is::<crate::stdio::Stdout>() {
-        Some(a.downcast_ref::<crate::stdio::Stdout>().unwrap().as_fd())
-    } else if a.is::<crate::stdio::Stderr>() {
-        Some(a.downcast_ref::<crate::stdio::Stderr>().unwrap().as_fd())
-    } else {
-        None
-    }
-}

crates/wasi-common/cap-std-sync/src/sched/windows.rs

@@ -9,30 +9,21 @@
 // taken the time to improve it. See bug #2880.
 
 use anyhow::Context;
-use io_extras::os::windows::{AsRawHandleOrSocket, RawHandleOrSocket};
 use std::ops::Deref;
 use std::sync::mpsc::{self, Receiver, RecvTimeoutError, Sender, TryRecvError};
 use std::sync::Mutex;
 use std::thread;
 use std::time::Duration;
-use wasi_common::{
+use wasi_common::sched::subscription::{RwEventFlags, Subscription};
-    file::WasiFile,
+use wasi_common::{file::WasiFile, sched::Poll, Error, ErrorExt};
-    sched::{
-        subscription::{RwEventFlags, Subscription},
-        Poll,
-    },
-    Error, ErrorExt,
-};
 
 pub async fn poll_oneoff<'a>(poll: &mut Poll<'a>) -> Result<(), Error> {
-    poll_oneoff_(poll, wasi_file_is_stdin, wasi_file_raw_handle).await
+    poll_oneoff_(poll, wasi_file_is_stdin).await
 }
 
-// For reuse by wasi-tokio, which has a different WasiFile -> RawHandle translator.
 pub async fn poll_oneoff_<'a>(
     poll: &mut Poll<'a>,
     file_is_stdin: impl Fn(&dyn WasiFile) -> bool,
-    file_to_handle: impl Fn(&dyn WasiFile) -> Option<RawHandleOrSocket>,
 ) -> Result<(), Error> {
     if poll.is_empty() {
         return Ok(());
@@ -61,21 +52,17 @@ pub async fn poll_oneoff_<'a>(
             Subscription::Read(r) => {
                 if file_is_stdin(r.file.deref()) {
                     stdin_read_subs.push(r);
-                } else if file_to_handle(r.file.deref()).is_some() {
+                } else if r.file.pollable().is_some() {
                     immediate_reads.push(r);
                 } else {
-                    return Err(
+                    return Err(Error::invalid_argument().context("file is not pollable"));
-                        Error::invalid_argument().context("read subscription fd downcast failed")
-                    );
                 }
             }
             Subscription::Write(w) => {
-                if file_to_handle(w.file.deref()).is_some() {
+                if w.file.pollable().is_some() {
                     immediate_writes.push(w);
                 } else {
-                    return Err(
+                    return Err(Error::invalid_argument().context("file is not pollable"));
-                        Error::invalid_argument().context("write subscription fd downcast failed")
-                    );
                 }
             }
             Subscription::MonotonicClock { .. } => unreachable!(),
@@ -139,49 +126,6 @@ pub fn wasi_file_is_stdin(f: &dyn WasiFile) -> bool {
     f.as_any().is::<crate::stdio::Stdin>()
 }
 
-pub fn wasi_file_raw_handle(f: &dyn WasiFile) -> Option<RawHandleOrSocket> {
-    let a = f.as_any();
-    if a.is::<crate::file::File>() {
-        Some(
-            a.downcast_ref::<crate::file::File>()
-                .unwrap()
-                .as_raw_handle_or_socket(),
-        )
-    } else if a.is::<crate::net::TcpStream>() {
-        Some(
-            a.downcast_ref::<crate::net::TcpStream>()
-                .unwrap()
-                .as_raw_handle_or_socket(),
-        )
-    } else if a.is::<crate::net::TcpListener>() {
-        Some(
-            a.downcast_ref::<crate::net::TcpListener>()
-                .unwrap()
-                .as_raw_handle_or_socket(),
-        )
-    } else if a.is::<crate::stdio::Stdin>() {
-        Some(
-            a.downcast_ref::<crate::stdio::Stdin>()
-                .unwrap()
-                .as_raw_handle_or_socket(),
-        )
-    } else if a.is::<crate::stdio::Stdout>() {
-        Some(
-            a.downcast_ref::<crate::stdio::Stdout>()
-                .unwrap()
-                .as_raw_handle_or_socket(),
-        )
-    } else if a.is::<crate::stdio::Stderr>() {
-        Some(
-            a.downcast_ref::<crate::stdio::Stderr>()
-                .unwrap()
-                .as_raw_handle_or_socket(),
-        )
-    } else {
-        None
-    }
-}
-
 enum PollState {
     Ready,
     NotReady, // Not ready, but did not wait

crates/wasi-common/cap-std-sync/src/stdio.rs

@@ -31,6 +31,15 @@ impl WasiFile for Stdin {
     fn as_any(&self) -> &dyn Any {
         self
     }
+    #[cfg(unix)]
+    fn pollable(&self) -> Option<rustix::fd::BorrowedFd> {
+        Some(self.0.as_fd())
+    }
+
+    #[cfg(windows)]
+    fn pollable(&self) -> Option<io_extras::os::windows::RawHandleOrSocket> {
+        Some(self.0.as_raw_handle_or_socket())
+    }
     async fn get_filetype(&mut self) -> Result<FileType, Error> {
         if self.isatty() {
             Ok(FileType::CharacterDevice)
@@ -98,6 +107,15 @@ macro_rules! wasi_file_write_impl {
             fn as_any(&self) -> &dyn Any {
                 self
             }
+            #[cfg(unix)]
+            fn pollable(&self) -> Option<rustix::fd::BorrowedFd> {
+                Some(self.0.as_fd())
+            }
+
+            #[cfg(windows)]
+            fn pollable(&self) -> Option<io_extras::os::windows::RawHandleOrSocket> {
+                Some(self.0.as_raw_handle_or_socket())
+            }
             async fn get_filetype(&mut self) -> Result<FileType, Error> {
                 if self.isatty() {
                     Ok(FileType::CharacterDevice)

crates/wasi-common/src/file.rs

@@ -7,6 +7,16 @@ pub trait WasiFile: Send + Sync {
     fn as_any(&self) -> &dyn Any;
     async fn get_filetype(&mut self) -> Result<FileType, Error>;
 
+    #[cfg(unix)]
+    fn pollable(&self) -> Option<rustix::fd::BorrowedFd> {
+        None
+    }
+
+    #[cfg(windows)]
+    fn pollable(&self) -> Option<io_extras::os::windows::RawHandleOrSocket> {
+        None
+    }
+
     fn isatty(&mut self) -> bool {
         false
     }

crates/wasi-common/tokio/src/file.rs

@@ -94,6 +94,15 @@ macro_rules! wasi_file_impl {
             fn as_any(&self) -> &dyn Any {
                 self
             }
+            #[cfg(unix)]
+            fn pollable(&self) -> Option<rustix::fd::BorrowedFd> {
+                Some(self.0.as_fd())
+            }
+
+            #[cfg(windows)]
+            fn pollable(&self) -> Option<io_extras::os::windows::RawHandleOrSocket> {
+                Some(self.0.as_raw_handle_or_socket())
+            }
             async fn datasync(&mut self) -> Result<(), Error> {
                 block_on_dummy_executor(|| self.0.datasync())
             }

crates/wasi-common/tokio/src/sched/windows.rs

@@ -1,5 +1,4 @@
 use crate::block_on_dummy_executor;
-use io_extras::os::windows::{AsRawHandleOrSocket, RawHandleOrSocket};
 use wasi_cap_std_sync::sched::windows::poll_oneoff_;
 use wasi_common::{file::WasiFile, sched::Poll, Error};
 
@@ -8,52 +7,9 @@ pub async fn poll_oneoff<'a>(poll: &mut Poll<'a>) -> Result<(), Error> {
     // we use the blocking poll_oneoff implementation from the wasi-cap-std-crate.
     // We provide a function specific to this crate's WasiFile types for downcasting
     // to a RawHandle.
-    block_on_dummy_executor(move || poll_oneoff_(poll, wasi_file_is_stdin, wasi_file_raw_handle))
+    block_on_dummy_executor(move || poll_oneoff_(poll, wasi_file_is_stdin))
 }
 
 pub fn wasi_file_is_stdin(f: &dyn WasiFile) -> bool {
     f.as_any().is::<crate::stdio::Stdin>()
 }
-
-fn wasi_file_raw_handle(f: &dyn WasiFile) -> Option<RawHandleOrSocket> {
-    let a = f.as_any();
-    if a.is::<crate::file::File>() {
-        Some(
-            a.downcast_ref::<crate::file::File>()
-                .unwrap()
-                .as_raw_handle_or_socket(),
-        )
-    } else if a.is::<crate::net::TcpListener>() {
-        Some(
-            a.downcast_ref::<crate::net::TcpListener>()
-                .unwrap()
-                .as_raw_handle_or_socket(),
-        )
-    } else if a.is::<crate::net::TcpStream>() {
-        Some(
-            a.downcast_ref::<crate::net::TcpStream>()
-                .unwrap()
-                .as_raw_handle_or_socket(),
-        )
-    } else if a.is::<crate::stdio::Stdin>() {
-        Some(
-            a.downcast_ref::<crate::stdio::Stdin>()
-                .unwrap()
-                .as_raw_handle_or_socket(),
-        )
-    } else if a.is::<crate::stdio::Stdout>() {
-        Some(
-            a.downcast_ref::<crate::stdio::Stdout>()
-                .unwrap()
-                .as_raw_handle_or_socket(),
-        )
-    } else if a.is::<crate::stdio::Stderr>() {
-        Some(
-            a.downcast_ref::<crate::stdio::Stderr>()
-                .unwrap()
-                .as_raw_handle_or_socket(),
-        )
-    } else {
-        None
-    }
-}