diff --git a/ext/src/helpers/mod.rs b/ext/src/helpers/mod.rs index 8d03e6cf..5a531371 100644 --- 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 new file mode 100644 index 00000000..2b05745a --- /dev/null +++ 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, + /// 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, capacity: usize) -> Self { + Self { buffer, capacity } + } +} + +impl io::Write for OutputLimitedBuffer { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + // 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 index 116545d0..f18e0fe7 100644 --- 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 index 7764af9e..927d8e06 100644 --- 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), + Buffer(Opaque, 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] @@ -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 index b4044b5d..523828e9 100644 --- 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!")