Some minor fixes and features for WASI and sockets (bytecodealliance#…

…6948)

* Use `command::add_to_linker` in tests to reduce the number of times
  all the `add_to_linker` are listed.
* Add all `wasi:sockets` interfaces currently implemented to both the
  sync and async `command` functions (this enables all the interfaces in
  the CLI for example).
* Use `tokio::net::TcpStream::try_io` whenever I/O is performed on a
  socket, ensuring that readable/writable flags are set/cleared
  appropriately (otherwise once readable a socket is infinitely readable).
* Add a `with_ambient_tokio_runtime` helper function to use when
  creating a `tokio::net::TcpStream` since otherwise it panics due to a
  lack of active runtime in a synchronous context.
* Add `WouldBlock` handling to return a 0-length read.
* Add an `--inherit-network` CLI flag to enable basic usage of sockets
  in the CLI.

This will conflict a small amount with bytecodealliance#6877 but should be easy to
resolve, and otherwise this targets different usability points/issues
than that PR.

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)?;