Some minor fixes and features for WASI and sockets

crates/test-programs/tests/reactor.rs

@@ -71,24 +71,9 @@ async fn instantiate(
     wasi_ctx: ReactorCtx,
 ) -> Result<(Store<ReactorCtx>, TestReactor)> {
     let mut linker = Linker::new(&ENGINE);
-
-    // All of the imports available to the world are provided by the wasi-common crate:
-    preview2::bindings::filesystem::types::add_to_linker(&mut linker, |x| x)?;
-    preview2::bindings::filesystem::preopens::add_to_linker(&mut linker, |x| x)?;
-    preview2::bindings::io::streams::add_to_linker(&mut linker, |x| x)?;
-    preview2::bindings::cli::environment::add_to_linker(&mut linker, |x| x)?;
-    preview2::bindings::cli::exit::add_to_linker(&mut linker, |x| x)?;
-    preview2::bindings::cli::stdin::add_to_linker(&mut linker, |x| x)?;
-    preview2::bindings::cli::stdout::add_to_linker(&mut linker, |x| x)?;
-    preview2::bindings::cli::stderr::add_to_linker(&mut linker, |x| x)?;
-    preview2::bindings::cli::terminal_input::add_to_linker(&mut linker, |x| x)?;
-    preview2::bindings::cli::terminal_output::add_to_linker(&mut linker, |x| x)?;
-    preview2::bindings::cli::terminal_stdin::add_to_linker(&mut linker, |x| x)?;
-    preview2::bindings::cli::terminal_stdout::add_to_linker(&mut linker, |x| x)?;
-    preview2::bindings::cli::terminal_stderr::add_to_linker(&mut linker, |x| x)?;
+    preview2::command::add_to_linker(&mut linker)?;
 
     let mut store = Store::new(&ENGINE, wasi_ctx);
-
     let (testreactor, _instance) =
         TestReactor::instantiate_async(&mut store, &component, &linker).await?;
     Ok((store, testreactor))

crates/test-programs/tests/wasi-sockets.rs

@@ -45,24 +45,7 @@ async fn run(name: &str) -> anyhow::Result<()> {
     let component = get_component(name);
     let mut linker = Linker::new(&ENGINE);
 
-    preview2::bindings::io::streams::add_to_linker(&mut linker, |x| x)?;
-    preview2::bindings::poll::poll::add_to_linker(&mut linker, |x| x)?;
-    preview2::bindings::cli::exit::add_to_linker(&mut linker, |x| x)?;
-    preview2::bindings::cli::stdin::add_to_linker(&mut linker, |x| x)?;
-    preview2::bindings::cli::stdout::add_to_linker(&mut linker, |x| x)?;
-    preview2::bindings::cli::stderr::add_to_linker(&mut linker, |x| x)?;
-    preview2::bindings::cli::terminal_input::add_to_linker(&mut linker, |x| x)?;
-    preview2::bindings::cli::terminal_output::add_to_linker(&mut linker, |x| x)?;
-    preview2::bindings::cli::terminal_stdin::add_to_linker(&mut linker, |x| x)?;
-    preview2::bindings::cli::terminal_stdout::add_to_linker(&mut linker, |x| x)?;
-    preview2::bindings::cli::terminal_stderr::add_to_linker(&mut linker, |x| x)?;
-    preview2::bindings::cli::environment::add_to_linker(&mut linker, |x| x)?;
-    preview2::bindings::filesystem::types::add_to_linker(&mut linker, |x| x)?;
-    preview2::bindings::filesystem::preopens::add_to_linker(&mut linker, |x| x)?;
-    preview2::bindings::sockets::tcp::add_to_linker(&mut linker, |x| x)?;
-    preview2::bindings::sockets::tcp_create_socket::add_to_linker(&mut linker, |x| x)?;
-    preview2::bindings::sockets::network::add_to_linker(&mut linker, |x| x)?;
-    preview2::bindings::sockets::instance_network::add_to_linker(&mut linker, |x| x)?;
+    preview2::command::add_to_linker(&mut linker)?;
 
     // Create our wasi context.
     let mut table = Table::new();

crates/wasi/src/preview2/command.rs

@@ -37,7 +37,6 @@ pub fn add_to_linker<T: WasiView>(l: &mut wasmtime::component::Linker<T>) -> any
     crate::preview2::bindings::clocks::timezone::add_to_linker(l, |t| t)?;
     crate::preview2::bindings::filesystem::types::add_to_linker(l, |t| t)?;
     crate::preview2::bindings::filesystem::preopens::add_to_linker(l, |t| t)?;
-    crate::preview2::bindings::sockets::tcp::add_to_linker(l, |t| t)?;
     crate::preview2::bindings::poll::poll::add_to_linker(l, |t| t)?;
     crate::preview2::bindings::io::streams::add_to_linker(l, |t| t)?;
     crate::preview2::bindings::random::random::add_to_linker(l, |t| t)?;
@@ -51,6 +50,10 @@ pub fn add_to_linker<T: WasiView>(l: &mut wasmtime::component::Linker<T>) -> any
     crate::preview2::bindings::cli::terminal_stdin::add_to_linker(l, |t| t)?;
     crate::preview2::bindings::cli::terminal_stdout::add_to_linker(l, |t| t)?;
     crate::preview2::bindings::cli::terminal_stderr::add_to_linker(l, |t| t)?;
+    crate::preview2::bindings::sockets::tcp::add_to_linker(l, |t| t)?;
+    crate::preview2::bindings::sockets::tcp_create_socket::add_to_linker(l, |t| t)?;
+    crate::preview2::bindings::sockets::instance_network::add_to_linker(l, |t| t)?;
+    crate::preview2::bindings::sockets::network::add_to_linker(l, |t| t)?;
     Ok(())
 }
 
@@ -110,6 +113,9 @@ pub mod sync {
         crate::preview2::bindings::cli::terminal_stdout::add_to_linker(l, |t| t)?;
         crate::preview2::bindings::cli::terminal_stderr::add_to_linker(l, |t| t)?;
         crate::preview2::bindings::sockets::tcp::add_to_linker(l, |t| t)?;
+        crate::preview2::bindings::sockets::tcp_create_socket::add_to_linker(l, |t| t)?;
+        crate::preview2::bindings::sockets::instance_network::add_to_linker(l, |t| t)?;
+        crate::preview2::bindings::sockets::network::add_to_linker(l, |t| t)?;
         Ok(())
     }
 }

crates/wasi/src/preview2/host/tcp.rs

@@ -201,10 +201,12 @@ impl<T: WasiView> tcp::Host for T {
         }
 
         // Do the OS accept call.
-        let (connection, _addr) = socket
-            .tcp_socket()
-            .as_socketlike_view::<TcpListener>()
-            .accept_with(Blocking::No)?;
+        let tcp_socket = socket.tcp_socket();
+        let (connection, _addr) = tcp_socket.try_io(Interest::READABLE, || {
+            tcp_socket
+                .as_socketlike_view::<TcpListener>()
+                .accept_with(Blocking::No)
+        })?;
         let tcp_socket = HostTcpSocket::from_tcp_stream(connection)?;
 
         let input_clone = tcp_socket.clone_inner();

crates/wasi/src/preview2/mod.rs

@@ -179,3 +179,13 @@ pub(crate) fn in_tokio<F: std::future::Future>(f: F) -> F::Output {
         }
     }
 }
+
+fn with_ambient_tokio_runtime<R>(f: impl FnOnce() -> R) -> R {
+    match tokio::runtime::Handle::try_current() {
+        Ok(_) => f(),
+        Err(_) => {
+            let _enter = RUNTIME.enter();
+            f()
+        }
+    }
+}

crates/wasi/src/preview2/tcp.rs

@@ -7,6 +7,7 @@ use io_lifetimes::AsSocketlike;
 use std::io;
 use std::sync::Arc;
 use system_interface::io::IoExt;
+use tokio::io::Interest;
 
 /// The state of a TCP socket.
 ///
@@ -64,15 +65,17 @@ impl HostTcpSocket {
         // by our async implementation.
         let tcp_socket = TcpListener::new(family, Blocking::No)?;
 
-        let tcp_socket = unsafe {
-            tokio::net::TcpStream::try_from(std::net::TcpStream::from_raw_socketlike(
-                tcp_socket.into_raw_socketlike(),
-            ))
-            .unwrap()
-        };
+        let std_socket =
+            unsafe { std::net::TcpStream::from_raw_socketlike(tcp_socket.into_raw_socketlike()) };
+
+        let tokio_tcp_socket = crate::preview2::with_ambient_tokio_runtime(|| {
+            tokio::net::TcpStream::try_from(std_socket).unwrap()
+        });
 
         Ok(Self {
-            inner: Arc::new(HostTcpSocketInner { tcp_socket }),
+            inner: Arc::new(HostTcpSocketInner {
+                tcp_socket: tokio_tcp_socket,
+            }),
             tcp_state: HostTcpState::Default,
         })
     }
@@ -84,15 +87,16 @@ impl HostTcpSocket {
         let fd = rustix::fd::OwnedFd::from(tcp_socket);
         let tcp_socket = TcpListener::from(fd);
 
-        let tcp_socket = unsafe {
-            tokio::net::TcpStream::try_from(std::net::TcpStream::from_raw_socketlike(
-                tcp_socket.into_raw_socketlike(),
-            ))
-            .unwrap()
-        };
+        let std_tcp_socket =
+            unsafe { std::net::TcpStream::from_raw_socketlike(tcp_socket.into_raw_socketlike()) };
+        let tokio_tcp_socket = crate::preview2::with_ambient_tokio_runtime(|| {
+            tokio::net::TcpStream::try_from(std_tcp_socket).unwrap()
+        });
 
         Ok(Self {
-            inner: Arc::new(HostTcpSocketInner { tcp_socket }),
+            inner: Arc::new(HostTcpSocketInner {
+                tcp_socket: tokio_tcp_socket,
+            }),
             tcp_state: HostTcpState::Default,
         })
     }
@@ -121,10 +125,10 @@ impl HostInputStream for Arc<HostTcpSocketInner> {
             return Ok((Bytes::new(), StreamState::Open));
         }
         let mut buf = BytesMut::zeroed(size);
-        let r = self
-            .tcp_socket()
-            .as_socketlike_view::<TcpStream>()
-            .read(&mut buf);
+        let socket = self.tcp_socket();
+        let r = socket.try_io(Interest::READABLE, || {
+            socket.as_socketlike_view::<TcpStream>().read(&mut buf)
+        });
         let (n, state) = read_result(r)?;
         buf.truncate(n);
         Ok((buf.freeze(), state))
@@ -142,10 +146,10 @@ impl HostOutputStream for Arc<HostTcpSocketInner> {
         if buf.is_empty() {
             return Ok((0, StreamState::Open));
         }
-        let r = self
-            .tcp_socket
-            .as_socketlike_view::<TcpStream>()
-            .write(buf.as_ref());
+        let socket = self.tcp_socket();
+        let r = socket.try_io(Interest::WRITABLE, || {
+            socket.as_socketlike_view::<TcpStream>().write(buf.as_ref())
+        });
         let (n, state) = write_result(r)?;
         Ok((n, state))
     }
@@ -186,7 +190,11 @@ pub(crate) fn read_result(r: io::Result<usize>) -> io::Result<(usize, StreamStat
     match r {
         Ok(0) => Ok((0, StreamState::Closed)),
         Ok(n) => Ok((n, StreamState::Open)),
-        Err(e) if e.kind() == io::ErrorKind::Interrupted => Ok((0, StreamState::Open)),
+        Err(e)
+            if e.kind() == io::ErrorKind::Interrupted || e.kind() == io::ErrorKind::WouldBlock =>
+        {
+            Ok((0, StreamState::Open))
+        }
         Err(e) => Err(e),
     }
 }

src/commands/run.rs

@@ -293,6 +293,11 @@ pub struct RunCommand {
     /// removed. For now this is primarily here for testing.
     #[clap(long)]
     preview2: bool,
+
+    /// Flag for WASI preview2 to inherit the host's network within the guest so
+    /// it has full access to all addresses/ports/etc.
+    #[clap(long)]
+    inherit_network: bool,
 }
 
 #[derive(Clone)]
@@ -1061,6 +1066,10 @@ impl RunCommand {
             );
         }
 
+        if self.inherit_network {
+            builder.inherit_network(ambient_authority());
+        }
+
         let data = store.data_mut();
         let table = Arc::get_mut(&mut data.preview2_table).unwrap();
         let ctx = builder.build(table)?;