diff --git a/examples/stdio.rs b/examples/stdio.rs new file mode 100644 index 0000000..c167b0f --- /dev/null +++ b/examples/stdio.rs @@ -0,0 +1,84 @@ +#![cfg_attr(not(target_os = "wasi"), no_main)] +#![cfg(target_os = "wasi")] + +//! Demonstrates async line-oriented stdin, stdout, stderr, and flushing. + +use anyhow::Result; +use wstd::io::{AsyncRead, AsyncWrite}; + +struct LineReader { + reader: R, + buffer: Vec, +} + +impl LineReader { + fn new(reader: R) -> Self { + Self { + reader, + buffer: Vec::new(), + } + } + + async fn read_line(&mut self) -> std::io::Result>> { + loop { + if let Some(newline) = self.buffer.iter().position(|byte| *byte == b'\n') { + return Ok(Some(self.buffer.drain(..=newline).collect())); + } + + let mut chunk = [0; 1024]; + match self.reader.read(&mut chunk).await? { + 0 if self.buffer.is_empty() => return Ok(None), + 0 => return Ok(Some(std::mem::take(&mut self.buffer))), + len => { + self.buffer.extend_from_slice(&chunk[..len]); + } + } + } + } +} + +#[wstd::main] +async fn main() -> Result<()> { + let mut stdin = LineReader::new(wstd::io::stdin()); + let mut stdout = wstd::io::stdout(); + let mut stderr = wstd::io::stderr(); + + let mut stdin_error = None; + while let Some(line) = match stdin.read_line().await { + Ok(Some(line)) => Some(line), + Ok(None) => { + stdin_error = Some(std::io::ErrorKind::UnexpectedEof.into()); + None + } + Err(error) => { + stdin_error = Some(error); + None + } + } { + stdout.write_all(b"stdout: ").await.unwrap(); + stdout.write_all(&line).await.unwrap(); + stdout.flush().await.unwrap(); + + stderr.write_all(b"stderr: ").await.unwrap(); + stderr.write_all(&line).await.unwrap(); + stderr.flush().await.unwrap(); + } + let stdin_error = stdin_error.expect("the loop only exits after a stdin error"); + + let stdout_result = match stdout.write_all(b"stdin closed\n").await { + Ok(()) => stdout.flush().await, + Err(error) => Err(error), + }; + + stderr + .write_all(format!("stdin error: {:?}\n", stdin_error.kind()).as_bytes()) + .await?; + if let Err(error) = stdout_result { + stderr + .write_all(format!("stdout error: {:?}\n", error.kind()).as_bytes()) + .await?; + } + stderr.flush().await?; + + Ok(()) +} diff --git a/src/io/mod.rs b/src/io/mod.rs index 0f34b1b..14a0ed0 100644 --- a/src/io/mod.rs +++ b/src/io/mod.rs @@ -5,10 +5,36 @@ mod cursor; mod empty; mod read; mod seek; -mod stdio; -mod streams; + +mod streams { + mod sys { + #[cfg(target_env = "p2")] + pub(super) mod p2; + #[cfg(target_env = "p3")] + pub(super) mod p3; + } + #[cfg(target_env = "p2")] + pub use sys::p2::*; + #[cfg(target_env = "p3")] + pub use sys::p3::*; +} + +mod stdio { + mod sys { + #[cfg(target_env = "p2")] + pub(super) mod p2; + #[cfg(target_env = "p3")] + pub(super) mod p3; + } + #[cfg(target_env = "p2")] + pub use sys::p2::*; + #[cfg(target_env = "p3")] + pub use sys::p3::*; +} + mod write; +#[cfg(target_env = "p2")] pub use crate::runtime::AsyncPollable; pub use copy::*; pub use cursor::*; @@ -26,3 +52,37 @@ pub use std::io::Error; /// A specialized Result type for I/O operations. /// pub use std::io::Result; + +#[cfg(test)] +mod test { + use super::{AsyncWrite, stderr, stdout}; + use crate::runtime::block_on; + + #[test] + // No internal predicate. Run test with --nocapture and inspect output manually. + fn stdout_println_hello_world() { + block_on(async { + let mut stdout = stdout(); + let term = if stdout.is_terminal() { "is" } else { "is not" }; + stdout + .write_all(format!("hello, world! stdout {term} a terminal\n",).as_bytes()) + .await + .unwrap(); + stdout.flush().await.unwrap(); + }) + } + + #[test] + // No internal predicate. Run test with --nocapture and inspect output manually. + fn stderr_println_hello_world() { + block_on(async { + let mut stderr = stderr(); + let term = if stderr.is_terminal() { "is" } else { "is not" }; + stderr + .write_all(format!("hello, world! stderr {term} a terminal\n",).as_bytes()) + .await + .unwrap(); + stderr.flush().await.unwrap(); + }) + } +} diff --git a/src/io/stdio.rs b/src/io/stdio/sys/p2.rs similarity index 74% rename from src/io/stdio.rs rename to src/io/stdio/sys/p2.rs index e403986..696a946 100644 --- a/src/io/stdio.rs +++ b/src/io/stdio/sys/p2.rs @@ -1,7 +1,7 @@ -use super::{AsyncInputStream, AsyncOutputStream, AsyncRead, AsyncWrite, Result}; +use crate::io::{AsyncInputStream, AsyncOutputStream, AsyncRead, AsyncWrite, Result}; use std::cell::LazyCell; -use wasip2::cli::terminal_input::TerminalInput; -use wasip2::cli::terminal_output::TerminalOutput; + +use wasip2::cli::{terminal_input::TerminalInput, terminal_output::TerminalOutput}; /// Use the program's stdin as an `AsyncInputStream`. #[derive(Debug)] @@ -34,10 +34,10 @@ impl Stdin { impl AsyncRead for Stdin { #[inline] async fn read(&mut self, buf: &mut [u8]) -> Result { - self.stream.read(buf).await + let read = self.stream.read(buf).await?; + Ok(read) } - #[inline] async fn read_to_end(&mut self, buf: &mut Vec) -> Result { self.stream.read_to_end(buf).await } @@ -147,33 +147,3 @@ impl AsyncWrite for Stderr { self.stream.as_async_output_stream() } } - -#[cfg(test)] -mod test { - use crate::io::AsyncWrite; - use crate::runtime::block_on; - #[test] - // No internal predicate. Run test with --nocapture and inspect output manually. - fn stdout_println_hello_world() { - block_on(async { - let mut stdout = super::stdout(); - let term = if stdout.is_terminal() { "is" } else { "is not" }; - stdout - .write_all(format!("hello, world! stdout {term} a terminal\n",).as_bytes()) - .await - .unwrap(); - }) - } - #[test] - // No internal predicate. Run test with --nocapture and inspect output manually. - fn stderr_println_hello_world() { - block_on(async { - let mut stdout = super::stdout(); - let term = if stdout.is_terminal() { "is" } else { "is not" }; - stdout - .write_all(format!("hello, world! stderr {term} a terminal\n",).as_bytes()) - .await - .unwrap(); - }) - } -} diff --git a/src/io/stdio/sys/p3.rs b/src/io/stdio/sys/p3.rs new file mode 100644 index 0000000..71502f4 --- /dev/null +++ b/src/io/stdio/sys/p3.rs @@ -0,0 +1,224 @@ +use crate::io::{AsyncInputStream, AsyncOutputStream, AsyncRead, AsyncWrite, Result}; +use std::cell::LazyCell; +use std::pin::Pin; + +use wasip3::{ + cli::{terminal_input::TerminalInput, terminal_output::TerminalOutput, types::ErrorCode}, + wit_bindgen::FutureRead, +}; + +type Completion = FutureRead>; + +fn into_io_error(error: ErrorCode) -> std::io::Error { + let kind = match error { + ErrorCode::Io => std::io::ErrorKind::Other, + ErrorCode::IllegalByteSequence => std::io::ErrorKind::InvalidData, + ErrorCode::Pipe => std::io::ErrorKind::BrokenPipe, + }; + std::io::Error::new(kind, format!("WASI CLI error: {error:?}")) +} + +/// Use the program's stdin as an `AsyncInputStream`. +pub struct Stdin { + stream: AsyncInputStream, + completion: Pin>, + terminput: LazyCell>, +} + +pub fn stdin() -> Stdin { + let (stream, completion) = wasip3::cli::stdin::read_via_stream(); + Stdin { + stream: AsyncInputStream::new(stream), + completion: Box::pin(completion.into_future()), + terminput: LazyCell::new(wasip3::cli::terminal_stdin::get_terminal_stdin), + } +} + +impl Stdin { + /// Check if stdin is a terminal. + pub fn is_terminal(&self) -> bool { + LazyCell::force(&self.terminput).is_some() + } + + /// Get the `AsyncInputStream` used to implement `Stdin` + pub fn into_inner(self) -> AsyncInputStream { + self.stream + } + + async fn check_error(&mut self) -> Result<()> { + let (stream, completion) = wasip3::cli::stdin::read_via_stream(); + let old_stream = std::mem::replace(&mut self.stream, AsyncInputStream::new(stream)); + drop(old_stream); + let result = self.completion.as_mut().await.map_err(into_io_error); + self.completion = Box::pin(completion.into_future()); + result + } +} + +impl std::fmt::Debug for Stdin { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Stdin") + .field("stream", &self.stream) + .field("terminput", &self.terminput) + .finish_non_exhaustive() + } +} + +impl AsyncRead for Stdin { + #[inline] + async fn read(&mut self, buf: &mut [u8]) -> Result { + let read = self.stream.read(buf).await?; + #[cfg(target_env = "p3")] + if read == 0 && buf.is_empty() { + self.check_error().await?; + } + Ok(read) + } + + async fn read_to_end(&mut self, buf: &mut Vec) -> Result { + self.stream.read_to_end(buf).await + } + + #[inline] + fn as_async_input_stream(&mut self) -> Option<&mut AsyncInputStream> { + Some(&mut self.stream) + } +} + +/// Use the program's stdout as an `AsyncOutputStream`. +pub struct Stdout { + stream: AsyncOutputStream, + completion: Pin>, + termoutput: LazyCell>, +} + +pub fn stdout() -> Stdout { + let (tx, rx) = wasip3::wit_stream::new(); + let completion = wasip3::cli::stdout::write_via_stream(rx); + Stdout { + stream: AsyncOutputStream::new(tx), + completion: Box::pin(completion.into_future()), + termoutput: LazyCell::new(wasip3::cli::terminal_stdout::get_terminal_stdout), + } +} + +impl Stdout { + /// Check if stdout is a terminal. + pub fn is_terminal(&self) -> bool { + LazyCell::force(&self.termoutput).is_some() + } + + /// Get the `AsyncOutputStream` used to implement `Stdout` + pub fn into_inner(self) -> AsyncOutputStream { + self.stream + } + + async fn flush(&mut self) -> Result<()> { + let Self { + stream, completion, .. + } = std::mem::replace(self, stdout()); + drop(stream); + + completion.await.map_err(into_io_error) + } +} + +impl std::fmt::Debug for Stdout { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Stdout") + .field("stream", &self.stream) + .field("termoutput", &self.termoutput) + .finish_non_exhaustive() + } +} + +impl AsyncWrite for Stdout { + #[inline] + async fn write(&mut self, buf: &[u8]) -> Result { + self.stream.write(buf).await + } + + #[inline] + async fn flush(&mut self) -> Result<()> { + Self::flush(self).await + } + + #[inline] + async fn write_all(&mut self, buf: &[u8]) -> Result<()> { + self.stream.write_all(buf).await + } + + #[inline] + fn as_async_output_stream(&mut self) -> Option<&mut AsyncOutputStream> { + self.stream.as_async_output_stream() + } +} + +/// Use the program's stdout as an `AsyncOutputStream`. +pub struct Stderr { + stream: AsyncOutputStream, + completion: Pin>, + termoutput: LazyCell>, +} + +pub fn stderr() -> Stderr { + let (tx, rx) = wasip3::wit_stream::new(); + let completion = wasip3::cli::stderr::write_via_stream(rx); + Stderr { + stream: AsyncOutputStream::new(tx), + completion: Box::pin(completion.into_future()), + termoutput: LazyCell::new(wasip3::cli::terminal_stderr::get_terminal_stderr), + } +} + +impl Stderr { + /// Check if stderr is a terminal. + pub fn is_terminal(&self) -> bool { + LazyCell::force(&self.termoutput).is_some() + } + + /// Get the `AsyncOutputStream` used to implement `Stderr` + pub fn into_inner(self) -> AsyncOutputStream { + self.stream + } + + async fn flush(&mut self) -> Result<()> { + let Self { + stream, completion, .. + } = std::mem::replace(self, stderr()); + drop(stream); + + completion.await.map_err(into_io_error) + } +} + +impl std::fmt::Debug for Stderr { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Stderr") + .field("stream", &self.stream) + .field("termoutput", &self.termoutput) + .finish_non_exhaustive() + } +} + +impl AsyncWrite for Stderr { + #[inline] + async fn write(&mut self, buf: &[u8]) -> Result { + self.stream.write(buf).await + } + + #[inline] + async fn flush(&mut self) -> Result<()> { + Self::flush(self).await + } + + #[inline] + async fn write_all(&mut self, buf: &[u8]) -> Result<()> { + self.stream.write_all(buf).await + } + + #[inline] + fn as_async_output_stream(&mut self) -> Option<&mut AsyncOutputStream> { + self.stream.as_async_output_stream() + } +} diff --git a/src/io/streams.rs b/src/io/streams/sys/p2.rs similarity index 98% rename from src/io/streams.rs rename to src/io/streams/sys/p2.rs index 34fd3b0..2955825 100644 --- a/src/io/streams.rs +++ b/src/io/streams/sys/p2.rs @@ -1,4 +1,4 @@ -use super::{AsyncPollable, AsyncRead, AsyncWrite}; +use crate::io::{AsyncPollable, AsyncRead, AsyncWrite}; use crate::runtime::WaitFor; use std::future::{Future, poll_fn}; use std::pin::Pin; @@ -78,7 +78,12 @@ impl AsyncInputStream { /// Use this `AsyncInputStream` as a `futures_lite::stream::Stream` with /// items of `Result, std::io::Error>`. The returned byte vectors /// will be at most the `chunk_size` argument specified. + /// + /// # Panics + /// + /// Panics if `chunk_size` is zero. pub fn into_stream_of(self, chunk_size: usize) -> AsyncInputChunkStream { + assert!(chunk_size > 0, "chunk size must be non-zero"); AsyncInputChunkStream { stream: self, chunk_size, diff --git a/src/io/streams/sys/p3.rs b/src/io/streams/sys/p3.rs new file mode 100644 index 0000000..f86e6df --- /dev/null +++ b/src/io/streams/sys/p3.rs @@ -0,0 +1,379 @@ +use crate::io::{AsyncRead, AsyncWrite}; + +use wasip3::wit_bindgen::{StreamReader, StreamResult, StreamWriter}; + +use std::future::Future; +use std::pin::Pin; +use std::task::{Context, Poll}; + +/// A wrapper for the readable end of a `stream` that provides an +/// implementation of `AsyncRead`. +#[derive(Debug)] +pub struct AsyncInputStream { + stream: StreamReader, +} + +impl AsyncInputStream { + pub fn new(stream: StreamReader) -> Self { + Self { stream } + } + + /// Move the entire contents of an input stream directly into an output + /// stream, until the input stream has closed. This operation is optimized + /// to avoid copying stream contents into and out of memory. + pub async fn copy_to(&mut self, writer: &mut AsyncOutputStream) -> std::io::Result { + // TODO: The current implementation avoids the extra copy within Wasm + // that occurs with `AsyncRead::read`, but it still requires the host to + // copy bytes into Wasm in the first place and then run guest code to + // move those bytes to the output stream. This should all be further + // optimized away by switching to `stream.forward`. + const CHUNK_SIZE: usize = 1024; + let mut vec = Vec::with_capacity(CHUNK_SIZE); + let mut written = 0; + loop { + let (result, new_vec) = self.stream.read(vec).await; + vec = new_vec; + match result { + StreamResult::Complete(r) => { + writer.write_all(&vec).await?; + written += r as u64; + } + StreamResult::Dropped => break, + StreamResult::Cancelled => return Err(std::io::ErrorKind::Interrupted.into()), + } + vec.clear(); + } + Ok(written) + } + + /// Use this `AsyncInputStream` as a `futures_lite::stream::Stream` with + /// items of `Result, std::io::Error>`. The returned byte vectors + /// will be at most 8k. If you want to control chunk size, use + /// `Self::into_stream_of`. + pub fn into_stream(self) -> AsyncInputChunkStream { + AsyncInputChunkStream { + state: AsyncInputChunkStreamState::Ready(self), + chunk_size: 8 * 1024, + } + } + + /// Use this `AsyncInputStream` as a `futures_lite::stream::Stream` with + /// items of `Result, std::io::Error>`. The returned byte vectors + /// will be at most the `chunk_size` argument specified. + /// + /// # Panics + /// + /// Panics if `chunk_size` is zero. + pub fn into_stream_of(self, chunk_size: usize) -> AsyncInputChunkStream { + assert!(chunk_size > 0, "chunk size must be non-zero"); + AsyncInputChunkStream { + state: AsyncInputChunkStreamState::Ready(self), + chunk_size, + } + } + + /// Use this `AsyncInputStream` as a `futures_lite::stream::Stream` with + /// items of `Result`. + pub fn into_bytestream(self) -> AsyncInputByteStream { + AsyncInputByteStream { + stream: self.into_stream(), + buffer: std::io::Read::bytes(std::io::Cursor::new(Vec::new())), + } + } +} + +impl AsyncRead for AsyncInputStream { + async fn read(&mut self, buf: &mut [u8]) -> std::io::Result { + let vec = Vec::with_capacity(buf.len()); + let (result, vec) = self.stream.read(vec).await; + match result { + StreamResult::Complete(_) => { + buf[..vec.len()].copy_from_slice(&vec); + Ok(vec.len()) + } + StreamResult::Dropped => Ok(0), + StreamResult::Cancelled => Err(std::io::ErrorKind::Interrupted.into()), + } + } + + #[inline] + fn as_async_input_stream(&mut self) -> Option<&mut AsyncInputStream> { + Some(self) + } +} + +/// Wrapper of `AsyncInputStream` that impls `futures_lite::stream::Stream` +/// with an item of `Result, std::io::Error>` +pub struct AsyncInputChunkStream { + state: AsyncInputChunkStreamState, + chunk_size: usize, +} + +enum AsyncInputChunkStreamState { + Ready(AsyncInputStream), + Reading(Pin>>), + Done, +} + +enum AsyncInputChunkReadResult { + Chunk { + chunk: Vec, + stream: AsyncInputStream, + }, + Done(std::io::Result<()>), +} + +impl AsyncInputChunkStream { + /// Extract the `AsyncInputStream` which backs this stream. The operation + /// will fail if a read is currently in progress or the stream is done in + /// which case `self` is returned. + pub fn into_inner(self) -> Result { + match self.state { + AsyncInputChunkStreamState::Ready(stream) => Ok(stream), + AsyncInputChunkStreamState::Reading(_) | AsyncInputChunkStreamState::Done => Err(self), + } + } +} + +impl futures_lite::stream::Stream for AsyncInputChunkStream { + type Item = Result, std::io::Error>; + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + loop { + match &mut self.state { + AsyncInputChunkStreamState::Ready(_) => { + // We need to take ownership of the stream before we setting + // the new state, which requires this temporary `replace`. + let AsyncInputChunkStreamState::Ready(mut stream) = + std::mem::replace(&mut self.state, AsyncInputChunkStreamState::Done) + else { + unreachable!(); + }; + let chunk_size = self.chunk_size; + self.state = AsyncInputChunkStreamState::Reading(Box::pin(async move { + let (result, chunk) = + stream.stream.read(Vec::with_capacity(chunk_size)).await; + match result { + StreamResult::Complete(_) => { + AsyncInputChunkReadResult::Chunk { chunk, stream } + } + StreamResult::Dropped => AsyncInputChunkReadResult::Done(Ok(())), + StreamResult::Cancelled => AsyncInputChunkReadResult::Done(Err( + std::io::ErrorKind::Interrupted.into(), + )), + } + })); + } + AsyncInputChunkStreamState::Reading(read) => { + match std::task::ready!(read.as_mut().poll(cx)) { + AsyncInputChunkReadResult::Chunk { chunk, stream } => { + self.state = AsyncInputChunkStreamState::Ready(stream); + return Poll::Ready(Some(Ok(chunk))); + } + AsyncInputChunkReadResult::Done(result) => { + self.state = AsyncInputChunkStreamState::Done; + return match result { + Ok(()) => Poll::Ready(None), + Err(error) => Poll::Ready(Some(Err(error))), + }; + } + } + } + AsyncInputChunkStreamState::Done => return Poll::Ready(None), + } + } + } +} + +pin_project_lite::pin_project! { + /// Wrapper of `AsyncInputStream` that impls + /// `futures_lite::stream::Stream` with item `Result`. + pub struct AsyncInputByteStream { + #[pin] + stream: AsyncInputChunkStream, + buffer: std::io::Bytes>>, + } +} + +impl AsyncInputByteStream { + /// Extract the `AsyncInputStream` which backs this stream, and any bytes + /// read from the `AsyncInputStream` which have not yet been yielded by + /// the byte stream. If a `read` is in progress or the stream is completed + /// this will error and return back `self`. + pub fn into_inner(self) -> Result<(AsyncInputStream, Vec), Self> { + match self.stream.into_inner() { + Ok(inner) => Ok(( + inner, + self.buffer + .collect::, std::io::Error>>() + .expect("read of Cursor> is infallible"), + )), + Err(stream) => Err(Self { + stream, + buffer: self.buffer, + }), + } + } +} + +impl futures_lite::stream::Stream for AsyncInputByteStream { + type Item = std::result::Result; + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.project(); + match this.buffer.next() { + Some(byte) => Poll::Ready(Some(Ok(byte.expect("cursor on Vec is infallible")))), + None => match futures_lite::stream::Stream::poll_next(this.stream, cx) { + Poll::Ready(Some(Ok(bytes))) => { + let mut bytes = std::io::Read::bytes(std::io::Cursor::new(bytes)); + match bytes.next() { + Some(Ok(byte)) => { + *this.buffer = bytes; + Poll::Ready(Some(Ok(byte))) + } + Some(Err(err)) => Poll::Ready(Some(Err(err))), + None => Poll::Ready(None), + } + } + Poll::Ready(Some(Err(err))) => Poll::Ready(Some(Err(err))), + Poll::Ready(None) => Poll::Ready(None), + Poll::Pending => Poll::Pending, + }, + } + } +} + +/// A wrapper for the writable end of a `stream` resource that provides +/// implementations of `AsyncWrite`. +#[derive(Debug)] +pub struct AsyncOutputStream { + stream: StreamWriter, +} + +impl AsyncOutputStream { + pub fn new(stream: StreamWriter) -> Self { + Self { stream } + } +} + +impl AsyncWrite for AsyncOutputStream { + /// Asynchronously write to the output stream. + /// + /// Performs at most one write to the output stream. Returns how much of the + /// argument `buf` was written, or a `std::io::Error`. + async fn write(&mut self, buf: &[u8]) -> std::io::Result { + let mut vec = Vec::with_capacity(buf.len()); + vec.extend_from_slice(buf); + let (result, _abi_buf) = self.stream.write(vec).await; + match result { + StreamResult::Complete(sent) => Ok(sent), + StreamResult::Dropped => Err(std::io::ErrorKind::ConnectionReset.into()), + StreamResult::Cancelled => unreachable!("Write operation cannot be cancelled"), + } + } + + async fn write_all(&mut self, buf: &[u8]) -> std::io::Result<()> { + let remaining = self.stream.write_all(buf.to_vec()).await; + if remaining.is_empty() { + Ok(()) + } else { + Err(std::io::ErrorKind::ConnectionReset.into()) + } + } + + /// # Warning + /// + /// This is a no-op on generic p3 streams. Use interface-specific flush + /// methods when available (e.g. [`crate::io::Stdout::flush`] or + /// [`crate::io::Stderr::flush`]). + async fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } + + #[inline] + fn as_async_output_stream(&mut self) -> Option<&mut AsyncOutputStream> { + Some(self) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use futures_lite::StreamExt; + + #[test] + fn chunk_stream_partial_read_works() { + crate::runtime::block_on(async { + let (mut writer, reader) = wasip3::wit_stream::new(); + let read = crate::runtime::spawn(async move { + let mut chunks = AsyncInputStream::new(reader).into_stream_of(4); + let chunk = chunks.next().await.unwrap().unwrap(); + let end = chunks.next().await; + (chunk, end) + }); + + assert!(writer.write_all(vec![1, 2]).await.is_empty()); + drop(writer); + + let (chunk, end) = read.await; + assert_eq!(chunk, [1, 2]); + assert!(end.is_none()); + }); + } + + #[test] + fn output_stream_reports_closed_reader() { + crate::runtime::block_on(async { + let (writer, reader) = wasip3::wit_stream::new(); + drop(reader); + + let error = AsyncOutputStream::new(writer) + .write(&[1]) + .await + .unwrap_err(); + + assert_eq!(error.kind(), std::io::ErrorKind::ConnectionReset); + }); + } + + #[test] + fn output_stream_write_all_sends_entire_buffer() { + crate::runtime::block_on(async { + let (writer, reader) = wasip3::wit_stream::new(); + let collect = crate::runtime::spawn(reader.collect()); + let expected: Vec<_> = (0..=255).collect(); + let mut output = AsyncOutputStream::new(writer); + + output.write_all(&expected).await.unwrap(); + drop(output); + + assert_eq!(collect.await, expected); + }); + } + + #[test] + fn copy_to_works() { + crate::runtime::block_on(async { + let (mut source_writer, source_reader) = wasip3::wit_stream::new(); + let (destination_writer, destination_reader) = wasip3::wit_stream::new(); + + let copy = crate::runtime::spawn(async move { + let mut source = AsyncInputStream::new(source_reader); + let mut destination = AsyncOutputStream::new(destination_writer); + let copied = source.copy_to(&mut destination).await.unwrap(); + drop(destination); + copied + }); + let collect = crate::runtime::spawn(destination_reader.collect()); + + assert!(source_writer.write_all(vec![1]).await.is_empty()); + assert!(source_writer.write_all(vec![2]).await.is_empty()); + assert!(source_writer.write_all(vec![3, 4, 5]).await.is_empty()); + drop(source_writer); + + let copied = copy.await; + let actual = collect.await; + + assert_eq!(copied, 5); + assert_eq!(actual, [1, 2, 3, 4, 5]); + }); + } +} diff --git a/src/lib.rs b/src/lib.rs index e3fb694..7d20a2f 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -60,7 +60,6 @@ pub mod future; #[cfg(all(target_os = "wasi", target_env = "p2"))] #[macro_use] pub mod http; -#[cfg(all(target_os = "wasi", target_env = "p2"))] pub mod io; pub mod iter; #[cfg(all(target_os = "wasi", target_env = "p2"))] @@ -90,8 +89,8 @@ pub mod __internal { pub use wasip3; } -#[cfg(all(target_os = "wasi", target_env = "p2"))] pub mod prelude { + #[cfg(all(target_os = "wasi", target_env = "p2"))] pub use crate::future::FutureExt as _; pub use crate::io::AsyncRead as _; pub use crate::io::AsyncWrite as _; diff --git a/test-programs/build.rs b/test-programs/build.rs index 4beee02..cb81811 100644 --- a/test-programs/build.rs +++ b/test-programs/build.rs @@ -6,6 +6,7 @@ use std::process::Command; fn main() { let out_dir = PathBuf::from(var_os("OUT_DIR").expect("OUT_DIR env var exists")); + let nightly_toolchain = is_nightly_toolchain(); let meta = MetadataCommand::new() .exec() @@ -16,13 +17,20 @@ fn main() { meta.workspace_root.as_os_str().to_str().unwrap() ); - fn build_targets(pkg: &str, manifest: &str, kind: &str, out_dir: &PathBuf) { + // TODO: Remove this nightly only cfg once wasm32-wasip3 is available on + // stable. + println!("cargo::rustc-check-cfg=cfg(wstd_nightly)"); + if nightly_toolchain { + println!("cargo:rustc-cfg=wstd_nightly"); + } + + fn build_target(pkg: &str, manifest: &str, kind: &str, target: &str, out_dir: &Path) { // release build is required for aws sdk to not overflow wasm locals let status = Command::new("cargo") .arg("build") .arg(kind) .arg("--release") - .arg("--target=wasm32-wasip2") + .arg(format!("--target={target}")) .arg(format!("--package={pkg}")) .arg(format!("--manifest-path={manifest}")) .env("CARGO_TARGET_DIR", out_dir) @@ -33,12 +41,41 @@ fn main() { .expect("cargo build wstd examples"); assert!(status.success()); } - build_targets("wstd", "../Cargo.toml", "--examples", &out_dir); - build_targets("wstd-axum", "../Cargo.toml", "--examples", &out_dir); + + fn build_targets( + pkg: &str, + manifest: &str, + kind: &str, + nightly_toolchain: bool, + out_dir: &Path, + ) { + build_target(pkg, manifest, kind, "wasm32-wasip2", out_dir); + + if nightly_toolchain { + build_target(pkg, manifest, kind, "wasm32-wasip3", out_dir); + } + } + + build_targets( + "wstd", + "../Cargo.toml", + "--examples", + nightly_toolchain, + &out_dir, + ); + build_targets( + "wstd-axum", + "../Cargo.toml", + "--examples", + nightly_toolchain, + &out_dir, + ); build_targets( "wstd-aws-example", "../aws-example/Cargo.toml", "--bins", + // TODO: enable when aws example is running for WASIp3 + false, &out_dir, ); @@ -52,20 +89,27 @@ fn main() { generated_code += &format!("pub mod {name} {{"); for binary in meta.targets.iter().filter(|t| t.kind == [kind.clone()]) { let mut component_path = out_dir.join("wasm32-wasip2").join("release"); + let mut p3_component_path = out_dir.join("wasm32-wasip3").join("release"); match kind { TargetKind::Bin => {} TargetKind::Example => { component_path = component_path.join("examples"); + p3_component_path = p3_component_path.join("examples"); } _ => unimplemented!("path interpolation for TargetKind {kind:?}"), } component_path = component_path.join(format!("{}.wasm", binary.name)); + p3_component_path = p3_component_path.join(format!("{}.wasm", binary.name)); let const_name = binary.name.to_shouty_snake_case(); generated_code += &format!( "pub const {const_name}: &str = {:?};\n", component_path.as_os_str().to_str().expect("path is str") ); + generated_code += &format!( + "pub const {const_name}_P3: &str = {:?};\n", + p3_component_path.as_os_str().to_str().expect("path is str") + ); } generated_code += "}\n\n"; // end `pub mod {name}` generated_code @@ -108,6 +152,21 @@ fn main() { std::fs::write(out_dir.join("gen.rs"), generated_code).unwrap(); } +fn is_nightly_toolchain() -> bool { + let rustc = var_os("RUSTC").unwrap_or_else(|| "rustc".into()); + let output = Command::new(rustc) + .arg("--version") + .output() + .expect("query active rustc version"); + assert!( + output.status.success(), + "failed to query active rustc version" + ); + String::from_utf8(output.stdout) + .expect("rustc version is UTF-8") + .contains("-nightly") +} + fn rustflags() -> &'static str { match option_env!("RUSTFLAGS") { Some(s) if s.contains("-D warnings") => "-D warnings", diff --git a/test-programs/tests/stdio.rs b/test-programs/tests/stdio.rs new file mode 100644 index 0000000..d5d20b3 --- /dev/null +++ b/test-programs/tests/stdio.rs @@ -0,0 +1,61 @@ +use anyhow::{Context, Result}; +use std::io::{BufRead, BufReader, Read, Write}; +use std::process::{Command, Stdio}; + +fn run(component: &str, p3: bool) -> Result<()> { + let mut command = Command::new("wasmtime"); + command.arg("run"); + if p3 { + command.arg("-Sp3"); + } + + let mut child = command + .arg(component) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn()?; + + let mut stdin = child.stdin.take().context("child stdin")?; + let mut stdout = BufReader::new(child.stdout.take().context("child stdout")?); + let mut stderr = BufReader::new(child.stderr.take().context("child stderr")?); + + stdin.write_all(b"hello from stdin\n")?; + stdin.flush()?; + + let mut line = String::new(); + stdout.read_line(&mut line)?; + assert_eq!(line, "stdout: hello from stdin\n"); + + line.clear(); + stderr.read_line(&mut line)?; + assert_eq!(line, "stderr: hello from stdin\n"); + + // Closing stdin ends the guest's read loop. Closing stdout makes its next + // write fail, which the guest reports through the still-open stderr pipe. + drop(stdin); + drop(stdout); + + let mut errors = String::new(); + stderr.read_to_string(&mut errors)?; + + let status = child.wait()?; + assert!(status.success(), "stdio example failed: {status}"); + assert_eq!( + errors, + format!("stdin error: UnexpectedEof\nstdout error: ConnectionReset\n") + ); + + Ok(()) +} + +#[test_log::test] +fn stdio_p2() -> Result<()> { + run(test_programs::STDIO, false) +} + +#[cfg(wstd_nightly)] +#[test_log::test] +fn stdio_p3() -> Result<()> { + run(test_programs::STDIO_P3, true) +}