add string buffer for stdout and stderr (#316)

jenniferwills · web-flow · commit 3e9485b37c4f · 2024-04-24T08:26:41.000-04:00
* add string buffer for stdout and stderr

* handle case where buf can't be written

* add comment, rewrite condition to handle overflow

* return 0 and fmt

* check if strings are frozen before writing + tests

* raise and handle errors when a string buffer is frozen

* use eq not match

* document using frozen strings and encoding

* remove parentheses

* fix comments

* use wasm traps when string is frozen
diff --git a/ext/src/helpers/mod.rs b/ext/src/helpers/mod.rs
@@ -1,10 +1,12 @@
 mod macros;
 mod nogvl;
+mod output_limited_buffer;
 mod static_id;
 mod symbol_enum;
 mod tmplock;
 
 pub use nogvl::nogvl;
+pub use output_limited_buffer::OutputLimitedBuffer;
 pub use static_id::StaticId;
 pub use symbol_enum::SymbolEnum;
 pub use tmplock::Tmplock;
diff --git a/ext/src/helpers/output_limited_buffer.rs b/ext/src/helpers/output_limited_buffer.rs
@@ -0,0 +1,70 @@
+use magnus::{
+    value::{InnerValue, Opaque, ReprValue},
+    RString, Ruby,
+};
+use std::io;
+use std::io::ErrorKind;
+
+/// A buffer that limits the number of bytes that can be written to it.
+/// If the buffer is full, it will truncate the data.
+/// Is used in the buffer implementations of stdout and stderr in `WasiCtx` and `WasiCtxBuilder`.
+pub struct OutputLimitedBuffer {
+    buffer: Opaque<RString>,
+    /// The maximum number of bytes that can be written to the output stream buffer.
+    capacity: usize,
+}
+
+impl OutputLimitedBuffer {
+    #[must_use]
+    pub fn new(buffer: Opaque<RString>, capacity: usize) -> Self {
+        Self { buffer, capacity }
+    }
+}
+
+impl io::Write for OutputLimitedBuffer {
+    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+        // Append a buffer to the string and truncate when hitting the capacity.
+        // We return the input buffer size regardless of whether we truncated or not to avoid a panic.
+        let ruby = Ruby::get().unwrap();
+
+        let mut inner_buffer = self.buffer.get_inner_with(&ruby);
+
+        // Handling frozen case here is necessary because magnus does not check if a string is frozen before writing to it.
+        let is_frozen = inner_buffer.as_value().is_frozen();
+        if is_frozen {
+            return Err(io::Error::new(
+                ErrorKind::WriteZero,
+                "Cannot write to a frozen buffer.",
+            ));
+        }
+
+        if buf.is_empty() {
+            return Ok(0);
+        }
+
+        if inner_buffer
+            .len()
+            .checked_add(buf.len())
+            .is_some_and(|val| val < self.capacity)
+        {
+            let amount_written = inner_buffer.write(buf)?;
+            if amount_written < buf.len() {
+                return Ok(amount_written);
+            }
+        } else {
+            let portion = self.capacity - inner_buffer.len();
+            let amount_written = inner_buffer.write(&buf[0..portion])?;
+            if amount_written < portion {
+                return Ok(amount_written);
+            }
+        };
+
+        Ok(buf.len())
+    }
+
+    fn flush(&mut self) -> io::Result<()> {
+        let ruby = Ruby::get().unwrap();
+
+        self.buffer.get_inner_with(&ruby).flush()
+    }
+}
diff --git a/ext/src/ruby_api/wasi_ctx.rs b/ext/src/ruby_api/wasi_ctx.rs
@@ -4,13 +4,14 @@ use super::{
     WasiCtxBuilder,
 };
 use crate::error;
+use crate::helpers::OutputLimitedBuffer;
 use deterministic_wasi_ctx::build_wasi_ctx as wasi_deterministic_ctx;
 use magnus::{
     class, function, gc::Marker, method, prelude::*, typed_data::Obj, Error, Object, RString,
     RTypedData, Ruby, TypedData, Value,
 };
 use std::{borrow::Borrow, cell::RefCell, fs::File, path::PathBuf};
-use wasi_common::pipe::ReadPipe;
+use wasi_common::pipe::{ReadPipe, WritePipe};
 use wasi_common::WasiCtx as WasiCtxImpl;
 
 /// @yard
@@ -75,6 +76,21 @@ impl WasiCtx {
         rb_self
     }
 
+    /// @yard
+    /// Set stdout to write to a string buffer.
+    /// If the string buffer is frozen, Wasm execution will raise a Wasmtime::Error error.
+    /// No encoding checks are done on the resulting string, it is the caller's responsibility to ensure the string contains a valid encoding
+    /// @param buffer [String] The string buffer to write to.
+    /// @param capacity [Integer] The maximum number of bytes that can be written to the output buffer.
+    /// @def set_stout_buffer(buffer, capacity)
+    /// @return [WasiCtx] +self+
+    fn set_stdout_buffer(rb_self: RbSelf, buffer: RString, capacity: usize) -> RbSelf {
+        let inner = rb_self.inner.borrow_mut();
+        let pipe = WritePipe::new(OutputLimitedBuffer::new(buffer.into(), capacity));
+        inner.set_stdout(Box::new(pipe));
+        rb_self
+    }
+
     /// @yard
     /// Set stderr to write to a file. Will truncate the file if it exists,
     /// otherwise try to create it.
@@ -88,6 +104,21 @@ impl WasiCtx {
         rb_self
     }
 
+    /// @yard
+    /// Set stderr to write to a string buffer.
+    /// If the string buffer is frozen, Wasm execution will raise a Wasmtime::Error error.
+    /// No encoding checks are done on the resulting string, it is the caller's responsibility to ensure the string contains a valid encoding
+    /// @param buffer [String] The string buffer to write to.
+    /// @param capacity [Integer] The maximum number of bytes that can be written to the output buffer.
+    /// @def set_stout_buffer(buffer, capacity)
+    /// @return [WasiCtx] +self+
+    fn set_stderr_buffer(rb_self: RbSelf, buffer: RString, capacity: usize) -> RbSelf {
+        let inner = rb_self.inner.borrow_mut();
+        let pipe = WritePipe::new(OutputLimitedBuffer::new(buffer.into(), capacity));
+        inner.set_stderr(Box::new(pipe));
+        rb_self
+    }
+
     pub fn from_inner(inner: WasiCtxImpl) -> Self {
         Self {
             inner: RefCell::new(inner),
@@ -105,6 +136,8 @@ pub fn init() -> Result<(), Error> {
     class.define_method("set_stdin_file", method!(WasiCtx::set_stdin_file, 1))?;
     class.define_method("set_stdin_string", method!(WasiCtx::set_stdin_string, 1))?;
     class.define_method("set_stdout_file", method!(WasiCtx::set_stdout_file, 1))?;
+    class.define_method("set_stdout_buffer", method!(WasiCtx::set_stdout_buffer, 2))?;
     class.define_method("set_stderr_file", method!(WasiCtx::set_stderr_file, 1))?;
+    class.define_method("set_stderr_buffer", method!(WasiCtx::set_stderr_buffer, 2))?;
     Ok(())
 }
diff --git a/ext/src/ruby_api/wasi_ctx_builder.rs b/ext/src/ruby_api/wasi_ctx_builder.rs
@@ -1,12 +1,13 @@
 use super::{root, WasiCtx};
 use crate::error;
+use crate::helpers::OutputLimitedBuffer;
 use magnus::{
     class, function, gc::Marker, method, typed_data::Obj, value::Opaque, DataTypeFunctions, Error,
     Module, Object, RArray, RHash, RString, Ruby, TryConvert, TypedData,
 };
 use std::cell::RefCell;
 use std::{fs::File, path::PathBuf};
-use wasi_common::pipe::ReadPipe;
+use wasi_common::pipe::{ReadPipe, WritePipe};
 
 enum ReadStream {
     Inherit,
@@ -27,12 +28,14 @@ impl ReadStream {
 enum WriteStream {
     Inherit,
     Path(Opaque<RString>),
+    Buffer(Opaque<RString>, usize),
 }
 impl WriteStream {
     pub fn mark(&self, marker: &Marker) {
         match self {
             Self::Inherit => (),
             Self::Path(v) => marker.mark(*v),
+            Self::Buffer(v, _) => marker.mark(*v),
         }
     }
 }
@@ -149,6 +152,20 @@ impl WasiCtxBuilder {
         rb_self
     }
 
+    /// @yard
+    /// Set stdout to write to a string buffer.
+    /// If the string buffer is frozen, Wasm execution will raise a Wasmtime::Error error.
+    /// No encoding checks are done on the resulting string, it is the caller's responsibility to ensure the string contains a valid encoding
+    /// @param buffer [String] The string buffer to write to.
+    /// @param capacity [Integer] The maximum number of bytes that can be written to the output buffer.
+    /// @def set_stdout_buffer(buffer, capacity)
+    /// @return [WasiCtxBuilder] +self+
+    pub fn set_stdout_buffer(rb_self: RbSelf, buffer: RString, capacity: usize) -> RbSelf {
+        let mut inner = rb_self.inner.borrow_mut();
+        inner.stdout = Some(WriteStream::Buffer(buffer.into(), capacity));
+        rb_self
+    }
+
     /// @yard
     /// Inherit stderr from the current Ruby process.
     /// @return [WasiCtxBuilder] +self+
@@ -170,6 +187,19 @@ impl WasiCtxBuilder {
         rb_self
     }
 
+    /// @yard
+    /// Set stderr to write to a string buffer.
+    /// If the string buffer is frozen, Wasm execution will raise a Wasmtime::Error error.
+    /// No encoding checks are done on the resulting string, it is the caller's responsibility to ensure the string contains a valid encoding
+    /// @param buffer [String] The string buffer to write to.
+    /// @param capacity [Integer] The maximum number of bytes that can be written to the output buffer.
+    /// @def set_stderr_buffer(buffer, capacity)
+    /// @return [WasiCtxBuilder] +self+
+    pub fn set_stderr_buffer(rb_self: RbSelf, buffer: RString, capacity: usize) -> RbSelf {
+        let mut inner = rb_self.inner.borrow_mut();
+        inner.stderr = Some(WriteStream::Buffer(buffer.into(), capacity));
+        rb_self
+    }
     /// @yard
     /// Set env to the specified +Hash+.
     /// @param env [Hash<String, String>]
@@ -216,6 +246,10 @@ impl WasiCtxBuilder {
                 WriteStream::Path(path) => {
                     builder.stdout(file_w(ruby.get_inner(*path)).map(wasi_file)?)
                 }
+                WriteStream::Buffer(buffer, capacity) => {
+                    let buf = OutputLimitedBuffer::new(*buffer, *capacity);
+                    builder.stdout(Box::new(WritePipe::new(buf)))
+                }
             };
         }
 
@@ -225,6 +259,10 @@ impl WasiCtxBuilder {
                 WriteStream::Path(path) => {
                     builder.stderr(file_w(ruby.get_inner(*path)).map(wasi_file)?)
                 }
+                WriteStream::Buffer(buffer, capacity) => {
+                    let buf = OutputLimitedBuffer::new(*buffer, *capacity);
+                    builder.stderr(Box::new(WritePipe::new(buf)))
+                }
             };
         }
 
@@ -282,12 +320,20 @@ pub fn init() -> Result<(), Error> {
         "set_stdout_file",
         method!(WasiCtxBuilder::set_stdout_file, 1),
     )?;
+    class.define_method(
+        "set_stdout_buffer",
+        method!(WasiCtxBuilder::set_stdout_buffer, 2),
+    )?;
 
     class.define_method("inherit_stderr", method!(WasiCtxBuilder::inherit_stderr, 0))?;
     class.define_method(
         "set_stderr_file",
         method!(WasiCtxBuilder::set_stderr_file, 1),
     )?;
+    class.define_method(
+        "set_stderr_buffer",
+        method!(WasiCtxBuilder::set_stderr_buffer, 2),
+    )?;
 
     class.define_method("set_env", method!(WasiCtxBuilder::set_env, 1))?;
 
diff --git a/spec/unit/wasi_spec.rb b/spec/unit/wasi_spec.rb
@@ -55,6 +55,84 @@ module Wasmtime
         expect(stdout.dig("wasi", "stdin")).to eq("stdin content")
       end
 
+      it "writes std streams to buffers" do
+        File.write(tempfile_path("stdin"), "stdin content")
+
+        stdout_str = ""
+        stderr_str = ""
+        wasi_config = WasiCtxBuilder.new
+          .set_stdin_file(tempfile_path("stdin"))
+          .set_stdout_buffer(stdout_str, 40000)
+          .set_stderr_buffer(stderr_str, 40000)
+          .build
+
+        run_wasi_module(wasi_config)
+
+        parsed_stdout = JSON.parse(stdout_str)
+        parsed_stderr = JSON.parse(stderr_str)
+        expect(parsed_stdout.fetch("name")).to eq("stdout")
+        expect(parsed_stderr.fetch("name")).to eq("stderr")
+      end
+
+      it "writes std streams to buffers until capacity" do
+        File.write(tempfile_path("stdin"), "stdin content")
+
+        stdout_str = ""
+        stderr_str = ""
+        wasi_config = WasiCtxBuilder.new
+          .set_stdin_file(tempfile_path("stdin"))
+          .set_stdout_buffer(stdout_str, 5)
+          .set_stderr_buffer(stderr_str, 10)
+          .build
+
+        run_wasi_module(wasi_config)
+
+        expect(stdout_str).to eq("{\"nam")
+        expect(stderr_str).to eq("{\"name\":\"s")
+      end
+
+      it "frozen stdout string is not written to" do
+        File.write(tempfile_path("stdin"), "stdin content")
+
+        stdout_str = ""
+        stderr_str = ""
+        wasi_config = WasiCtxBuilder.new
+          .set_stdin_file(tempfile_path("stdin"))
+          .set_stdout_buffer(stdout_str, 40000)
+          .set_stderr_buffer(stderr_str, 40000)
+          .build
+
+        stdout_str.freeze
+        expect { run_wasi_module(wasi_config) }.to raise_error do |error|
+          expect(error).to be_a(Wasmtime::Error)
+          expect(error.message).to match(/error while executing at wasm backtrace:/)
+        end
+
+        parsed_stderr = JSON.parse(stderr_str)
+        expect(stdout_str).to eq("")
+        expect(parsed_stderr.fetch("name")).to eq("stderr")
+      end
+      it "frozen stderr string is not written to" do
+        File.write(tempfile_path("stdin"), "stdin content")
+
+        stderr_str = ""
+        stdout_str = ""
+        wasi_config = WasiCtxBuilder.new
+          .set_stdin_file(tempfile_path("stdin"))
+          .set_stderr_buffer(stderr_str, 40000)
+          .set_stdout_buffer(stdout_str, 40000)
+          .build
+
+        stderr_str.freeze
+        expect { run_wasi_module(wasi_config) }.to raise_error do |error|
+          expect(error).to be_a(Wasmtime::Error)
+          expect(error.message).to match(/error while executing at wasm backtrace:/)
+        end
+
+        expect(stderr_str).to eq("")
+        expect(stdout_str).to eq("")
+      end
+
       it "reads stdin from string" do
         env = wasi_module_env { |config| config.set_stdin_string("¡UTF-8 from Ruby!") }
         expect(env.fetch("stdin")).to eq("¡UTF-8 from Ruby!")
