Tests where two commands have their stdin/stdout piped together

Author: elliottt

crates/test-programs/artifacts/build.rs

@@ -74,6 +74,7 @@ fn build_and_generate_tests() {
             s if s.starts_with("cli_") => "cli",
             s if s.starts_with("api_") => "api",
             s if s.starts_with("nn_") => "nn",
+            s if s.starts_with("piped_") => "piped",
             // If you're reading this because you hit this panic, either add it
             // to a test suite above or add a new "suite". The purpose of the
             // categorization above is to have a static assertion that tests

crates/test-programs/src/bin/piped_polling.rs

@@ -0,0 +1,51 @@
+use test_programs::wasi::cli::{stdin, stdout};
+
+fn main() {
+    match std::env::var("PIPED_SIDE")
+        .expect("piped tests require the PIPED_SIDE env var")
+        .as_str()
+    {
+        "PRODUCER" => producer(),
+        "CONSUMER" => consumer(),
+        side => panic!("unknown piped test side: {side}"),
+    }
+}
+
+fn producer() {
+    let out = stdout::get_stdout();
+    let out_pollable = out.subscribe();
+
+    for i in 1..100 {
+        let message = format!("{i}");
+        loop {
+            let available = out.check_write().unwrap() as usize;
+            if available >= message.len() {
+                break;
+            }
+
+            out_pollable.block();
+            assert!(out_pollable.ready());
+        }
+
+        out.write(message.as_bytes()).unwrap()
+    }
+
+    drop(out_pollable);
+}
+
+fn consumer() {
+    let stdin = stdin::get_stdin();
+    let stdin_pollable = stdin.subscribe();
+
+    for i in 1..100 {
+        let expected = format!("{i}");
+
+        stdin_pollable.block();
+        assert!(stdin_pollable.ready());
+
+        let bytes = stdin.read(expected.len() as u64).unwrap();
+        assert_eq!(&bytes, expected.as_bytes());
+    }
+
+    drop(stdin_pollable);
+}

crates/test-programs/src/bin/piped_simple.rs

@@ -0,0 +1,30 @@
+use test_programs::wasi::cli::{stdin, stdout};
+
+fn main() {
+    match std::env::var("PIPED_SIDE")
+        .expect("piped tests require the PIPED_SIDE env var")
+        .as_str()
+    {
+        "PRODUCER" => producer(),
+        "CONSUMER" => consumer(),
+        side => panic!("unknown piped test side: {side}"),
+    }
+}
+
+const CHUNK: &[u8] = &[b'a'; 50];
+
+fn producer() {
+    let out = stdout::get_stdout();
+    let n = out.check_write().unwrap() as usize;
+    assert!(n > CHUNK.len());
+    out.write(CHUNK).unwrap();
+}
+
+fn consumer() {
+    let stdin = stdin::get_stdin();
+    let stdin_pollable = stdin.subscribe();
+    stdin_pollable.block();
+    assert!(stdin_pollable.ready());
+    let bytes = stdin.read(CHUNK.len() as u64).unwrap();
+    assert_eq!(&bytes, CHUNK);
+}

crates/wasi/src/preview2/stdio.rs

@@ -36,6 +36,16 @@ pub trait StdinStream: Send + Sync {
     fn isatty(&self) -> bool;
 }
 
+impl<T: StdinStream + ?Sized> StdinStream for Box<T> {
+    fn stream(&self) -> Box<dyn HostInputStream> {
+        self.as_ref().stream()
+    }
+
+    fn isatty(&self) -> bool {
+        self.as_ref().isatty()
+    }
+}
+
 impl StdinStream for pipe::MemoryInputPipe {
     fn stream(&self) -> Box<dyn HostInputStream> {
         Box::new(self.clone())
@@ -79,6 +89,16 @@ pub trait StdoutStream: Send + Sync {
     fn isatty(&self) -> bool;
 }
 
+impl<T: StdoutStream + ?Sized> StdoutStream for Box<T> {
+    fn stream(&self) -> Box<dyn HostOutputStream> {
+        self.as_ref().stream()
+    }
+
+    fn isatty(&self) -> bool {
+        self.as_ref().isatty()
+    }
+}
+
 impl StdoutStream for pipe::MemoryOutputPipe {
     fn stream(&self) -> Box<dyn HostOutputStream> {
         Box::new(self.clone())

crates/wasi/tests/all/main.rs

@@ -7,7 +7,7 @@ use wasmtime::{
 use wasmtime_wasi::preview2::{
     pipe::MemoryOutputPipe,
     preview1::{WasiPreview1Adapter, WasiPreview1View},
-    DirPerms, FilePerms, WasiCtx, WasiCtxBuilder, WasiView,
+    DirPerms, FilePerms, StdinStream, StdoutStream, WasiCtx, WasiCtxBuilder, WasiView,
 };
 
 struct Ctx {
@@ -48,40 +48,98 @@ fn prepare_workspace(exe_name: &str) -> Result<TempDir> {
     Ok(tempdir)
 }
 
-fn store(engine: &Engine, name: &str, inherit_stdio: bool) -> Result<(Store<Ctx>, TempDir)> {
-    let stdout = MemoryOutputPipe::new(4096);
-    let stderr = MemoryOutputPipe::new(4096);
-    let workspace = prepare_workspace(name)?;
-
-    // Create our wasi context.
-    let mut builder = WasiCtxBuilder::new();
-    if inherit_stdio {
-        builder.inherit_stdio();
-    } else {
-        builder.stdout(stdout.clone()).stderr(stderr.clone());
+struct StoreBuilder {
+    builder: WasiCtxBuilder,
+    inherit_stdio: bool,
+    workspace: TempDir,
+    stdin: Option<Box<dyn StdinStream>>,
+    stdout: Option<Box<dyn StdoutStream>>,
+    stderr: Option<Box<dyn StdoutStream>>,
+}
+
+impl StoreBuilder {
+    fn new(name: &str) -> Result<Self> {
+        let mut builder = WasiCtxBuilder::new();
+
+        builder
+            .args(&[name, "."])
+            .inherit_network()
+            .allow_ip_name_lookup(true);
+
+        let workspace = prepare_workspace(name)?;
+
+        Ok(Self {
+            builder,
+            inherit_stdio: false,
+            workspace,
+            stdin: None,
+            stdout: None,
+            stderr: None,
+        })
     }
 
-    builder
-        .args(&[name, "."])
-        .inherit_network()
-        .allow_ip_name_lookup(true);
-    println!("preopen: {:?}", workspace);
-    let preopen_dir =
-        cap_std::fs::Dir::open_ambient_dir(workspace.path(), cap_std::ambient_authority())?;
-    builder.preopened_dir(preopen_dir, DirPerms::all(), FilePerms::all(), ".");
-    for (var, val) in test_programs_artifacts::wasi_tests_environment() {
-        builder.env(var, val);
+    fn stdout(&mut self, stdout: Box<dyn StdoutStream>) -> &mut Self {
+        self.stdout.replace(stdout);
+        self
     }
 
-    let ctx = Ctx {
-        table: ResourceTable::new(),
-        wasi: builder.build(),
-        stderr,
-        stdout,
-        adapter: WasiPreview1Adapter::new(),
-    };
+    fn stdin(&mut self, stdin: Box<dyn StdinStream>) -> &mut Self {
+        self.stdin.replace(stdin);
+        self
+    }
+
+    fn build(mut self, engine: &Engine) -> Result<(Store<Ctx>, TempDir)> {
+        let stdout = MemoryOutputPipe::new(4096);
+        let stderr = MemoryOutputPipe::new(4096);
+
+        // Create our wasi context.
+        if self.inherit_stdio {
+            self.builder.inherit_stdio();
+        } else {
+            if let Some(stdin) = self.stdin {
+                self.builder.stdin(stdin);
+            }
 
-    Ok((Store::new(&engine, ctx), workspace))
+            if let Some(stdout) = self.stdout {
+                self.builder.stdout(stdout);
+            } else {
+                self.builder.stdout(stdout.clone());
+            }
+
+            if let Some(stderr) = self.stderr {
+                self.builder.stderr(stderr);
+            } else {
+                self.builder.stderr(stderr.clone());
+            }
+        }
+
+        println!("preopen: {:?}", self.workspace);
+        let preopen_dir = cap_std::fs::Dir::open_ambient_dir(
+            self.workspace.path(),
+            cap_std::ambient_authority(),
+        )?;
+        self.builder
+            .preopened_dir(preopen_dir, DirPerms::all(), FilePerms::all(), ".");
+        for (var, val) in test_programs_artifacts::wasi_tests_environment() {
+            self.builder.env(var, val);
+        }
+
+        let ctx = Ctx {
+            table: ResourceTable::new(),
+            wasi: self.builder.build(),
+            stderr,
+            stdout,
+            adapter: WasiPreview1Adapter::new(),
+        };
+
+        Ok((Store::new(&engine, ctx), self.workspace))
+    }
+}
+
+fn store(engine: &Engine, name: &str, inherit_stdio: bool) -> Result<(Store<Ctx>, TempDir)> {
+    let mut builder = StoreBuilder::new(name)?;
+    builder.inherit_stdio = inherit_stdio;
+    builder.build(engine)
 }
 
 impl Drop for Ctx {
@@ -108,5 +166,6 @@ macro_rules! assert_test_exists {
 
 mod api;
 mod async_;
+mod piped;
 mod preview1;
 mod sync;