c-api: wasi: Add custom callbacks for stdout/stderr

Fixes bytecodealliance#11963

Author: MangoPeachGrape

crates/c-api/include/wasi.h

@@ -145,6 +145,19 @@ WASI_API_EXTERN bool wasi_config_set_stdout_file(wasi_config_t *config,
  */
 WASI_API_EXTERN void wasi_config_inherit_stdout(wasi_config_t *config);
 
+/**
+ * \brief Configures standard output to be directed to \p callback
+ *
+ * \param callback The callback that will get called for each write with the
+ * buffer
+ * \param data An optional user provided data that will be passed to \p callback
+ * \param finalizer An optional callback to be called to destroy \p data
+ */
+WASI_API_EXTERN void wasi_config_set_stdout_custom(
+    wasi_config_t *config,
+    void (*callback)(void *, const unsigned char *, size_t), void *data,
+    void (*finalizer)(void *));
+
 /**
  * \brief Configures standard output to be written to the specified file.
  *
@@ -163,6 +176,19 @@ WASI_API_EXTERN bool wasi_config_set_stderr_file(wasi_config_t *config,
  */
 WASI_API_EXTERN void wasi_config_inherit_stderr(wasi_config_t *config);
 
+/**
+ * \brief Configures standard error output to be directed to \p callback
+ *
+ * \param callback The callback that will get called for each write with the
+ * buffer
+ * \param data An optional user provided data that will be passed to \p callback
+ * \param finalizer An optional callback to be called to destroy \p data
+ */
+WASI_API_EXTERN void wasi_config_set_stderr_custom(
+    wasi_config_t *config,
+    void (*callback)(void *, const unsigned char *, size_t), void *data,
+    void (*finalizer)(void *));
+
 /**
  * \brief The permissions granted for a directory when preopening it.
  */

crates/c-api/src/wasi.rs

@@ -2,7 +2,7 @@
 
 use crate::wasm_byte_vec_t;
 use anyhow::Result;
-use std::ffi::{CStr, c_char};
+use std::ffi::{CStr, c_char, c_void};
 use std::fs::File;
 use std::path::Path;
 use std::slice;
@@ -149,6 +149,34 @@ pub extern "C" fn wasi_config_inherit_stdout(config: &mut wasi_config_t) {
     config.builder.inherit_stdout();
 }
 
+struct CustomOutputStream {
+    foreign_data: crate::ForeignData,
+    callback: extern "C" fn(*mut c_void, *const u8, usize),
+}
+
+impl wasmtime_wasi::p2::pipe::SimpleCustomOutputWriter for CustomOutputStream {
+    fn write(&self, buf: &[u8]) {
+        (self.callback)(self.foreign_data.data, buf.as_ptr(), buf.len());
+    }
+}
+
+#[unsafe(no_mangle)]
+pub extern "C" fn wasi_config_set_stdout_custom(
+    config: &mut wasi_config_t,
+    callback: Option<extern "C" fn(*mut c_void, *const u8, usize)>,
+    data: *mut c_void,
+    finalizer: Option<extern "C" fn(*mut c_void)>,
+) {
+    config
+        .builder
+        .stdout(wasmtime_wasi::p2::pipe::SimpleCustomOutputStream::new(
+            CustomOutputStream {
+                foreign_data: crate::ForeignData { data, finalizer },
+                callback: callback.expect("Callback must be provided"),
+            },
+        ));
+}
+
 #[unsafe(no_mangle)]
 pub unsafe extern "C" fn wasi_config_set_stderr_file(
     config: &mut wasi_config_t,
@@ -171,6 +199,23 @@ pub extern "C" fn wasi_config_inherit_stderr(config: &mut wasi_config_t) {
     config.builder.inherit_stderr();
 }
 
+#[unsafe(no_mangle)]
+pub extern "C" fn wasi_config_set_stderr_custom(
+    config: &mut wasi_config_t,
+    callback: Option<extern "C" fn(*mut c_void, *const u8, usize)>,
+    data: *mut c_void,
+    finalizer: Option<extern "C" fn(*mut c_void)>,
+) {
+    config
+        .builder
+        .stderr(wasmtime_wasi::p2::pipe::SimpleCustomOutputStream::new(
+            CustomOutputStream {
+                foreign_data: crate::ForeignData { data, finalizer },
+                callback: callback.expect("Callback must be provided"),
+            },
+        ));
+}
+
 #[unsafe(no_mangle)]
 pub unsafe extern "C" fn wasi_config_preopen_dir(
     config: &mut wasi_config_t,

crates/wasi/src/cli/mem.rs

@@ -32,3 +32,16 @@ impl StdoutStream for p2::pipe::MemoryOutputPipe {
         Box::new(self.clone())
     }
 }
+
+// Implementation for p2::pipe::SimpleCustomOutputStream
+impl<T: p2::pipe::SimpleCustomOutputWriter> IsTerminal for p2::pipe::SimpleCustomOutputStream<T> {
+    fn is_terminal(&self) -> bool {
+        false
+    }
+}
+
+impl<T: p2::pipe::SimpleCustomOutputWriter> StdoutStream for p2::pipe::SimpleCustomOutputStream<T> {
+    fn async_stream(&self) -> Box<dyn AsyncWrite + Send + Sync> {
+        Box::new(Clone::clone(self))
+    }
+}

crates/wasi/src/p2/pipe.rs

@@ -146,6 +146,68 @@ impl AsyncWrite for MemoryOutputPipe {
     }
 }
 
+pub trait SimpleCustomOutputWriter: Send + Sync + 'static {
+    fn write(&self, buf: &[u8]);
+}
+
+pub struct SimpleCustomOutputStream<T: SimpleCustomOutputWriter> {
+    writer: Arc<Mutex<T>>,
+}
+
+impl<T: SimpleCustomOutputWriter> SimpleCustomOutputStream<T> {
+    pub fn new(x: T) -> Self {
+        Self {
+            writer: Arc::new(Mutex::new(x)),
+        }
+    }
+}
+
+impl<T: SimpleCustomOutputWriter> Clone for SimpleCustomOutputStream<T> {
+    fn clone(&self) -> Self {
+        Self {
+            writer: self.writer.clone(),
+        }
+    }
+}
+
+#[async_trait::async_trait]
+impl<T: SimpleCustomOutputWriter> Pollable for SimpleCustomOutputStream<T> {
+    async fn ready(&mut self) {}
+}
+
+#[async_trait::async_trait]
+impl<T: SimpleCustomOutputWriter> OutputStream for SimpleCustomOutputStream<T> {
+    fn write(&mut self, bytes: Bytes) -> Result<(), StreamError> {
+        self.writer.lock().unwrap().write(&bytes);
+        // Always ready for writing
+        Ok(())
+    }
+    fn flush(&mut self) -> Result<(), StreamError> {
+        // This stream is always flushed
+        Ok(())
+    }
+    fn check_write(&mut self) -> Result<usize, StreamError> {
+        Ok(8192)
+    }
+}
+
+impl<T: SimpleCustomOutputWriter> AsyncWrite for SimpleCustomOutputStream<T> {
+    fn poll_write(
+        self: Pin<&mut Self>,
+        _cx: &mut Context<'_>,
+        buf: &[u8],
+    ) -> Poll<io::Result<usize>> {
+        self.writer.lock().unwrap().write(buf);
+        Poll::Ready(Ok(buf.len()))
+    }
+    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        Poll::Ready(Ok(()))
+    }
+    fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        Poll::Ready(Ok(()))
+    }
+}
+
 /// Provides a [`InputStream`] impl from a [`tokio::io::AsyncRead`] impl
 pub struct AsyncReadStream {
     closed: bool,