numcodecs_wasm_host_reproducible/
stdio.rs

1use std::{fmt, io::Write, sync::OnceLock};
2
3use wasm_component_layer::{
4    AsContextMut, Func, FuncType, InterfaceIdentifier, Linker, ListType, PackageIdentifier,
5    PackageName, Value, ValueType,
6};
7
8pub fn add_to_linker(linker: &mut Linker, mut ctx: impl AsContextMut) -> Result<(), anyhow::Error> {
9    let WasiSandboxedStdioInterface {
10        stdio: simple_stdio_interface,
11    } = WasiSandboxedStdioInterface::get();
12
13    let simple_stdio_instance = linker.define_instance(simple_stdio_interface.clone())?;
14
15    simple_stdio_instance.define_func(
16        "write-stdout",
17        OutputStream::Stdout.create_write_func(ctx.as_context_mut()),
18    )?;
19    simple_stdio_instance.define_func(
20        "flush-stdout",
21        OutputStream::Stdout.create_flush_func(ctx.as_context_mut()),
22    )?;
23    simple_stdio_instance.define_func(
24        "write-stderr",
25        OutputStream::Stderr.create_write_func(ctx.as_context_mut()),
26    )?;
27    simple_stdio_instance.define_func(
28        "flush-stderr",
29        OutputStream::Stderr.create_flush_func(ctx.as_context_mut()),
30    )?;
31
32    Ok(())
33}
34
35#[derive(Copy, Clone)]
36enum OutputStream {
37    Stdout,
38    Stderr,
39}
40
41impl OutputStream {
42    fn create_write_func(self, ctx: impl AsContextMut) -> Func {
43        Func::new(
44            ctx,
45            FuncType::new([ValueType::List(ListType::new(ValueType::U8))], []),
46            move |_ctx, args, results| {
47                let [Value::List(contents)] = args else {
48                    anyhow::bail!("invalid wasi-sandboxed:io/stdio#write-{self} arguments");
49                };
50                let Ok(contents) = contents.typed::<u8>() else {
51                    anyhow::bail!("invalid wasi-sandboxed:io/stdio#write-{self} argument type");
52                };
53
54                anyhow::ensure!(
55                    results.is_empty(),
56                    "invalid wasi-sandboxed:io/stdio#write-{self} results"
57                );
58
59                if let Err(err) = match self {
60                    Self::Stdout => std::io::stdout().write_all(contents),
61                    Self::Stderr => std::io::stderr().write_all(contents),
62                } {
63                    error!(
64                        "Failed to write {} byte{} to {self}: {err}",
65                        contents.len(),
66                        if contents.len() == 1 { "" } else { "s" }
67                    );
68                }
69
70                Ok(())
71            },
72        )
73    }
74
75    fn create_flush_func(self, ctx: impl AsContextMut) -> Func {
76        Func::new(ctx, FuncType::new([], []), move |_ctx, args, results| {
77            anyhow::ensure!(
78                args.is_empty(),
79                "invalid wasi-sandboxed:io/stdio#flush-{self} arguments"
80            );
81
82            anyhow::ensure!(
83                results.is_empty(),
84                "invalid wasi-sandboxed:io/stdio#flush-{self} results"
85            );
86
87            if let Err(err) = match self {
88                Self::Stdout => std::io::stdout().flush(),
89                Self::Stderr => std::io::stderr().flush(),
90            } {
91                error!("Failed to flush {self}: {err}");
92            }
93
94            Ok(())
95        })
96    }
97}
98
99impl fmt::Display for OutputStream {
100    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
101        fmt.write_str(match self {
102            Self::Stdout => "stdout",
103            Self::Stderr => "stderr",
104        })
105    }
106}
107
108#[non_exhaustive]
109pub struct WasiSandboxedStdioInterface {
110    pub stdio: InterfaceIdentifier,
111}
112
113impl WasiSandboxedStdioInterface {
114    #[must_use]
115    pub fn get() -> &'static Self {
116        static WASI_SANDBOXED_STDIO_INTERFACE: OnceLock<WasiSandboxedStdioInterface> =
117            OnceLock::new();
118
119        WASI_SANDBOXED_STDIO_INTERFACE.get_or_init(|| Self {
120            stdio: InterfaceIdentifier::new(
121                PackageIdentifier::new(
122                    PackageName::new("wasi-sandboxed", "io"),
123                    Some(semver::Version::new(0, 2, 3)),
124                ),
125                "stdio",
126            ),
127        })
128    }
129}
numcodecs_wasm_host_reproducible/stdio.rs

numcodecs_wasm_host_reproducible/
stdio.rs
