wasi:io: enforce the check-write permit in the remaining output streams

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<usize, StreamError> {
-        Ok(64 * 1024)
+        Ok(WRITE_BUDGET)
     }
 }
 

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");
+}

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:?}");
+    }
+}

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:?}");
+    }
+}

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"] {