c-api: wasi: Add custom callbacks for stdout/stderr (#11965)

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

Fixes #11963

* Move everything into c-api crate

* Return `ptrdiff_t`

* Smoke test

* Use `io::Error::from_raw_os_error()`

* Fixes

* Fix

* documentation prtest:full

Author: MangoPeachGrape

Cargo.lock

@@ -34,6 +34,9 @@ wat = { workspace = true, optional = true }
 cap-std = { workspace = true, optional = true }
 tokio = { workspace = true, optional = true, features = ["fs"] }
 wasmtime-wasi = { workspace = true, optional = true, features = ["p1"] }
+wasmtime-wasi-io = { workspace = true, optional = true, features = ["std"] }
+async-trait = { workspace = true, optional = true }
+bytes = { workspace = true, optional = true }
 
 # Optional dependencies for the `async` feature
 futures = { workspace = true, optional = true }
@@ -44,7 +47,7 @@ async = ['wasmtime/async', 'futures']
 profiling = ["wasmtime/profiling"]
 cache = ["wasmtime/cache"]
 parallel-compilation = ['wasmtime/parallel-compilation']
-wasi = ['cap-std', 'wasmtime-wasi', 'tokio']
+wasi = ['cap-std', 'wasmtime-wasi', 'wasmtime-wasi-io', 'tokio', 'async-trait', 'bytes']
 logging = ['dep:env_logger']
 disable-logging = ["log/max_level_off", "tracing/max_level_off"]
 coredump = ["wasmtime/coredump"]

crates/c-api/Cargo.toml

@@ -145,6 +145,21 @@ 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 config The config to operate on
+ * \param callback A non-null callback must be provided, that will get called
+ * for each write with the buffer. A positive return value indicates the amount
+ * of bytes written. Negative return values are treated as OS error codes.
+ * \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,
+    ptrdiff_t (*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 +178,21 @@ 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 config The config to operate on
+ * \param callback A non-null callback must be provided, that will get called
+ * for each write with the buffer. A positive return value indicates the amount
+ * of bytes written. Negative return values are treated as OS error codes.
+ * \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,
+    ptrdiff_t (*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/include/wasi.h

@@ -102,6 +102,9 @@ public:
     return wasi_config_preopen_dir(ptr.get(), path.c_str(), guest_path.c_str(),
                                    dir_perms, file_perms);
   }
+
+  /// \brief Returns the underlying C API pointer.
+  [[nodiscard]] wasi_config_t *capi() { return ptr.get(); }
 };
 
 } // namespace wasmtime

crates/c-api/include/wasmtime/wasi.hh

@@ -2,12 +2,17 @@
 
 use crate::wasm_byte_vec_t;
 use anyhow::Result;
-use std::ffi::{CStr, c_char};
+use bytes::Bytes;
+use std::ffi::{CStr, c_char, c_void};
 use std::fs::File;
 use std::path::Path;
+use std::pin::Pin;
 use std::slice;
+use std::task::{Context, Poll};
+use tokio::io::{self, AsyncWrite};
 use wasmtime_wasi::WasiCtxBuilder;
 use wasmtime_wasi::p1::WasiP1Ctx;
+use wasmtime_wasi_io::streams::StreamError;
 
 unsafe fn cstr_to_path<'a>(path: *const c_char) -> Option<&'a Path> {
     CStr::from_ptr(path).to_str().map(Path::new).ok()
@@ -149,6 +154,112 @@ pub extern "C" fn wasi_config_inherit_stdout(config: &mut wasi_config_t) {
     config.builder.inherit_stdout();
 }
 
+struct CustomOutputStreamInner {
+    foreign_data: crate::ForeignData,
+    callback: extern "C" fn(*mut c_void, *const u8, usize) -> isize,
+}
+
+impl CustomOutputStreamInner {
+    pub fn raw_write(&self, buf: &[u8]) -> io::Result<usize> {
+        let wrote = (self.callback)(self.foreign_data.data, buf.as_ptr(), buf.len());
+
+        if wrote >= 0 {
+            Ok(wrote as _)
+        } else {
+            Err(io::Error::from_raw_os_error(wrote.abs() as _))
+        }
+    }
+}
+
+#[derive(Clone)]
+pub struct CustomOutputStream {
+    inner: std::sync::Arc<CustomOutputStreamInner>,
+}
+
+impl CustomOutputStream {
+    pub fn new(
+        foreign_data: crate::ForeignData,
+        callback: extern "C" fn(*mut c_void, *const u8, usize) -> isize,
+    ) -> Self {
+        Self {
+            inner: std::sync::Arc::new(CustomOutputStreamInner {
+                foreign_data,
+                callback,
+            }),
+        }
+    }
+}
+
+#[async_trait::async_trait]
+impl wasmtime_wasi::p2::Pollable for CustomOutputStream {
+    async fn ready(&mut self) {}
+}
+
+#[async_trait::async_trait]
+impl wasmtime_wasi::p2::OutputStream for CustomOutputStream {
+    fn write(&mut self, bytes: Bytes) -> Result<(), StreamError> {
+        let wrote = self
+            .inner
+            .raw_write(&bytes)
+            .map_err(|e| StreamError::LastOperationFailed(e.into()))?;
+
+        if wrote != bytes.len() {
+            return Err(StreamError::LastOperationFailed(anyhow::anyhow!(
+                "Partial writes in wasip2 implementation are not allowed"
+            )));
+        }
+
+        Ok(())
+    }
+    fn flush(&mut self) -> Result<(), StreamError> {
+        Ok(())
+    }
+    fn check_write(&mut self) -> Result<usize, StreamError> {
+        Ok(usize::MAX)
+    }
+}
+
+impl AsyncWrite for CustomOutputStream {
+    fn poll_write(
+        self: Pin<&mut Self>,
+        _cx: &mut Context<'_>,
+        buf: &[u8],
+    ) -> Poll<io::Result<usize>> {
+        Poll::Ready(self.inner.raw_write(buf))
+    }
+    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(()))
+    }
+}
+
+impl wasmtime_wasi::cli::IsTerminal for CustomOutputStream {
+    fn is_terminal(&self) -> bool {
+        false
+    }
+}
+
+impl wasmtime_wasi::cli::StdoutStream for CustomOutputStream {
+    fn async_stream(&self) -> Box<dyn AsyncWrite + Send + Sync> {
+        Box::new(self.clone())
+    }
+}
+
+#[unsafe(no_mangle)]
+pub extern "C" fn wasi_config_set_stdout_custom(
+    config: &mut wasi_config_t,
+    callback: extern "C" fn(*mut c_void, *const u8, usize) -> isize,
+    data: *mut c_void,
+    finalizer: Option<extern "C" fn(*mut c_void)>,
+) {
+    config.builder.stdout(CustomOutputStream::new(
+        crate::ForeignData { data, finalizer },
+        callback,
+    ));
+}
+
 #[unsafe(no_mangle)]
 pub unsafe extern "C" fn wasi_config_set_stderr_file(
     config: &mut wasi_config_t,
@@ -171,6 +282,19 @@ 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: extern "C" fn(*mut c_void, *const u8, usize) -> isize,
+    data: *mut c_void,
+    finalizer: Option<extern "C" fn(*mut c_void)>,
+) {
+    config.builder.stderr(CustomOutputStream::new(
+        crate::ForeignData { data, finalizer },
+        callback,
+    ));
+}
+
 #[unsafe(no_mangle)]
 pub unsafe extern "C" fn wasi_config_preopen_dir(
     config: &mut wasi_config_t,

crates/c-api/src/wasi.rs

@@ -1,4 +1,5 @@
 #include <gtest/gtest.h>
+#include <string_view>
 #include <wasmtime/component.hh>
 #include <wasmtime/store.hh>
 
@@ -20,6 +21,22 @@ TEST(wasip2, smoke) {
   auto context = store.context();
 
   wasmtime::WasiConfig config;
+
+  wasi_config_set_stdout_custom(
+      config.capi(),
+      [](void *, const unsigned char *buf, size_t len) -> ptrdiff_t {
+        std::cout << std::string_view{(const char *)(buf), len};
+        return len;
+      },
+      nullptr, nullptr);
+  wasi_config_set_stderr_custom(
+      config.capi(),
+      [](void *, const unsigned char *buf, size_t len) -> ptrdiff_t {
+        std::cerr << std::string_view{(const char *)(buf), len};
+        return len;
+      },
+      nullptr, nullptr);
+
   context.set_wasi(std::move(config)).unwrap();
   Component component = Component::compile(engine, component_text).unwrap();