diff --git a/crates/test-programs/tests/reactor.rs b/crates/test-programs/tests/reactor.rs index 7dff956b6179..521c67af8c9e 100644 --- a/crates/test-programs/tests/reactor.rs +++ b/crates/test-programs/tests/reactor.rs @@ -71,24 +71,9 @@ async fn instantiate( wasi_ctx: ReactorCtx, ) -> Result<(Store, 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)) diff --git a/crates/test-programs/tests/wasi-sockets.rs b/crates/test-programs/tests/wasi-sockets.rs index ff119108fc0c..c16e9d4fadf7 100644 --- a/crates/test-programs/tests/wasi-sockets.rs +++ b/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(); diff --git a/crates/wasi/src/preview2/command.rs b/crates/wasi/src/preview2/command.rs index ae01eeb0219d..1ba80b93923c 100644 --- a/crates/wasi/src/preview2/command.rs +++ b/crates/wasi/src/preview2/command.rs @@ -37,7 +37,6 @@ pub fn add_to_linker(l: &mut wasmtime::component::Linker) -> 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(l: &mut wasmtime::component::Linker) -> 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(()) } } diff --git a/crates/wasi/src/preview2/host/tcp.rs b/crates/wasi/src/preview2/host/tcp.rs index 2ec61f26b95d..acdca7ce018c 100644 --- a/crates/wasi/src/preview2/host/tcp.rs +++ b/crates/wasi/src/preview2/host/tcp.rs @@ -201,10 +201,12 @@ impl tcp::Host for T { } // Do the OS accept call. - let (connection, _addr) = socket - .tcp_socket() - .as_socketlike_view::() - .accept_with(Blocking::No)?; + let tcp_socket = socket.tcp_socket(); + let (connection, _addr) = tcp_socket.try_io(Interest::READABLE, || { + tcp_socket + .as_socketlike_view::() + .accept_with(Blocking::No) + })?; let tcp_socket = HostTcpSocket::from_tcp_stream(connection)?; let input_clone = tcp_socket.clone_inner(); diff --git a/crates/wasi/src/preview2/mod.rs b/crates/wasi/src/preview2/mod.rs index 2fcf78a7efe6..170465ba38ef 100644 --- a/crates/wasi/src/preview2/mod.rs +++ b/crates/wasi/src/preview2/mod.rs @@ -179,3 +179,13 @@ pub(crate) fn in_tokio(f: F) -> F::Output { } } } + +fn with_ambient_tokio_runtime(f: impl FnOnce() -> R) -> R { + match tokio::runtime::Handle::try_current() { + Ok(_) => f(), + Err(_) => { + let _enter = RUNTIME.enter(); + f() + } + } +} diff --git a/crates/wasi/src/preview2/tcp.rs b/crates/wasi/src/preview2/tcp.rs index 3821f4e55a1d..a563a6e09b3b 100644 --- a/crates/wasi/src/preview2/tcp.rs +++ b/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 { return Ok((Bytes::new(), StreamState::Open)); } let mut buf = BytesMut::zeroed(size); - let r = self - .tcp_socket() - .as_socketlike_view::() - .read(&mut buf); + let socket = self.tcp_socket(); + let r = socket.try_io(Interest::READABLE, || { + socket.as_socketlike_view::().read(&mut buf) + }); let (n, state) = read_result(r)?; buf.truncate(n); Ok((buf.freeze(), state)) @@ -142,10 +146,10 @@ impl HostOutputStream for Arc { if buf.is_empty() { return Ok((0, StreamState::Open)); } - let r = self - .tcp_socket - .as_socketlike_view::() - .write(buf.as_ref()); + let socket = self.tcp_socket(); + let r = socket.try_io(Interest::WRITABLE, || { + socket.as_socketlike_view::().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) -> 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), } } diff --git a/src/commands/run.rs b/src/commands/run.rs index 398b9bf686dc..a93eaecd0227 100644 --- a/src/commands/run.rs +++ b/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)?;