diff --git a/crates/c-api/src/wasi.rs b/crates/c-api/src/wasi.rs index 0f9722b87f39..a70bd3e92bfc 100644 --- a/crates/c-api/src/wasi.rs +++ b/crates/c-api/src/wasi.rs @@ -203,6 +203,9 @@ impl CustomOutputStream { } } +/// The number of bytes `check_write` permits and `write` accepts in one call. +const WRITE_BUDGET: usize = 64 * 1024; + #[async_trait::async_trait] impl wasmtime_wasi::p2::Pollable for CustomOutputStream { async fn ready(&mut self) {} @@ -211,6 +214,11 @@ impl wasmtime_wasi::p2::Pollable for CustomOutputStream { #[async_trait::async_trait] impl wasmtime_wasi::p2::OutputStream for CustomOutputStream { fn write(&mut self, bytes: Bytes) -> Result<(), StreamError> { + if bytes.len() > WRITE_BUDGET { + return Err(StreamError::Trap(wasmtime::format_err!( + "write exceeded budget" + ))); + } let wrote = self .inner .raw_write(&bytes) @@ -228,7 +236,7 @@ impl wasmtime_wasi::p2::OutputStream for CustomOutputStream { Ok(()) } fn check_write(&mut self) -> Result { - Ok(64 * 1024) + Ok(WRITE_BUDGET) } } diff --git a/crates/test-programs/src/bin/p2_cli_stdout_write_too_much.rs b/crates/test-programs/src/bin/p2_cli_stdout_write_too_much.rs new file mode 100644 index 000000000000..421e1e337f89 --- /dev/null +++ b/crates/test-programs/src/bin/p2_cli_stdout_write_too_much.rs @@ -0,0 +1,14 @@ +/// An output stream should trap when it is asked to write more bytes than +/// `check-write` permitted. +fn main() { + let stdout = test_programs::wasi::cli::stdout::get_stdout(); + + let permit = stdout.check_write().unwrap(); + assert!(permit > 0); + + // This should trap according to the `wasi:io/streams` spec since we're + // trying to write more bytes than `check-write` gave permission for: + let contents = vec![0; usize::try_from(permit + 1).unwrap()]; + stdout.write(&contents).unwrap(); + unreachable!("attempt to write more bytes than permitted should have trapped"); +} diff --git a/crates/wasi/src/cli/file.rs b/crates/wasi/src/cli/file.rs index ddf0705f9eec..d1dd60dbee0b 100644 --- a/crates/wasi/src/cli/file.rs +++ b/crates/wasi/src/cli/file.rs @@ -48,6 +48,11 @@ impl Pollable for OutputFile { impl OutputStream for OutputFile { fn write(&mut self, bytes: Bytes) -> StreamResult<()> { + if bytes.len() > crate::MAX_READ_SIZE_ALLOC { + return Err(StreamError::Trap(wasmtime::format_err!( + "write exceeded budget" + ))); + } (&*self.file) .write_all(&bytes) .map_err(|e| StreamError::LastOperationFailed(wasmtime::format_err!(e))) @@ -152,3 +157,22 @@ impl AsyncRead for InputFile { } } } + +#[cfg(test)] +mod tests { + use super::*; + + /// A write larger than the budget `check_write` reports has to trap instead + /// of being written out. + #[test] + fn write_larger_than_budget_traps() { + let path = std::env::temp_dir().join("wasmtime-output-file-write-budget"); + let file = std::fs::File::create(&path).unwrap(); + let mut stream = OutputFile::new(file); + assert_eq!(stream.check_write().unwrap(), crate::MAX_READ_SIZE_ALLOC); + let err = stream + .write(Bytes::from(vec![0; crate::MAX_READ_SIZE_ALLOC + 1])) + .unwrap_err(); + assert!(matches!(err, StreamError::Trap(_)), "{err:?}"); + } +} diff --git a/crates/wasi/src/p2/pipe.rs b/crates/wasi/src/p2/pipe.rs index 7b5bce9ac216..8c04fe7c108b 100644 --- a/crates/wasi/src/p2/pipe.rs +++ b/crates/wasi/src/p2/pipe.rs @@ -275,7 +275,10 @@ pub struct SinkOutputStream; #[async_trait::async_trait] impl OutputStream for SinkOutputStream { - fn write(&mut self, _buf: Bytes) -> Result<(), StreamError> { + fn write(&mut self, buf: Bytes) -> Result<(), StreamError> { + if buf.len() > crate::MAX_READ_SIZE_ALLOC { + return Err(StreamError::Trap(format_err!("write exceeded budget"))); + } Ok(()) } fn flush(&mut self) -> Result<(), StreamError> { @@ -905,3 +908,23 @@ mod test { assert_eq!(permit, 1024); } } + +#[cfg(test)] +mod write_budget_tests { + use super::*; + + /// A write larger than the budget `check_write` reports has to trap instead + /// of being silently discarded. + #[test] + fn sink_write_larger_than_budget_traps() { + let mut stream = SinkOutputStream; + assert_eq!(stream.check_write().unwrap(), crate::MAX_READ_SIZE_ALLOC); + stream + .write(Bytes::from(vec![0; crate::MAX_READ_SIZE_ALLOC])) + .unwrap(); + let err = stream + .write(Bytes::from(vec![0; crate::MAX_READ_SIZE_ALLOC + 1])) + .unwrap_err(); + assert!(matches!(err, StreamError::Trap(_)), "{err:?}"); + } +} diff --git a/tests/all/cli_tests.rs b/tests/all/cli_tests.rs index 27fd6fbda02d..5165e19cbff9 100644 --- a/tests/all/cli_tests.rs +++ b/tests/all/cli_tests.rs @@ -3549,6 +3549,26 @@ start a print 1234 Ok(()) } + #[test] + fn p2_cli_stdout_write_too_much() -> Result<()> { + // With `-Sinherit-stdout=n` stdout is a sink, so writing past the + // permit reported by `check-write` has to trap. + let output = super::wasmtime(&[ + "run", + "-Sinherit-stdout=n", + P2_CLI_STDOUT_WRITE_TOO_MUCH_COMPONENT, + ])? + .output()?; + // The stdout stream is a sink, so nothing is expected on stdout. + assert_eq!(output.stdout, b""); + let stderr = String::from_utf8_lossy(&output.stderr); + assert!( + stderr.contains("write exceeded budget"), + "bad stderr: {stderr}" + ); + Ok(()) + } + #[test] fn p2_cli_stdout_write_zeros_to_sink() -> Result<()> { for n in ["0", "100", "10000"] {