numcodecs_wasm_host_reproducible/
stdio.rs1use 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}