wasmtime_wasi/p3/cli/
host.rs

1use crate::I32Exit;
2use crate::cli::{IsTerminal, WasiCli, WasiCliCtxView};
3use crate::p3::DEFAULT_BUFFER_CAPACITY;
4use crate::p3::bindings::cli::types::ErrorCode;
5use crate::p3::bindings::cli::{
6    environment, exit, stderr, stdin, stdout, terminal_input, terminal_output, terminal_stderr,
7    terminal_stdin, terminal_stdout,
8};
9use crate::p3::cli::{TerminalInput, TerminalOutput};
10use bytes::BytesMut;
11use core::pin::Pin;
12use core::task::{Context, Poll};
13use std::io;
14use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
15use tokio::sync::oneshot;
16use wasmtime::AsContextMut;
17use wasmtime::component::{
18    Access, Destination, FutureReader, Resource, Source, StreamConsumer, StreamProducer,
19    StreamReader, StreamResult,
20};
21use wasmtime::{StoreContextMut, error::Context as _, format_err};
22
23struct InputStreamProducer {
24    rx: Pin<Box<dyn AsyncRead + Send + Sync>>,
25    result_tx: Option<oneshot::Sender<ErrorCode>>,
26}
27
28fn io_error_to_error_code(err: io::Error) -> ErrorCode {
29    match err.kind() {
30        io::ErrorKind::BrokenPipe => ErrorCode::Pipe,
31        other => {
32            tracing::warn!("stdio error: {other}");
33            ErrorCode::Io
34        }
35    }
36}
37
38impl<D> StreamProducer<D> for InputStreamProducer {
39    type Item = u8;
40    type Buffer = BytesMut;
41
42    fn poll_produce<'a>(
43        mut self: Pin<&mut Self>,
44        cx: &mut Context<'_>,
45        mut store: StoreContextMut<'a, D>,
46        dst: Destination<'a, Self::Item, Self::Buffer>,
47        finish: bool,
48    ) -> Poll<wasmtime::Result<StreamResult>> {
49        // If the destination buffer is empty then this is a request on
50        // behalf of the guest to wait for this input stream to be readable.
51        // The `AsyncRead` trait abstraction does not provide the ability to
52        // await this event so we're forced to basically just lie here and
53        // say we're ready read data later.
54        //
55        // See WebAssembly/component-model#561 for some more information.
56        if dst.remaining(store.as_context_mut()) == Some(0) {
57            return Poll::Ready(Ok(StreamResult::Completed));
58        }
59
60        let mut dst = dst.as_direct(store, DEFAULT_BUFFER_CAPACITY);
61        let mut buf = ReadBuf::new(dst.remaining());
62        match self.rx.as_mut().poll_read(cx, &mut buf) {
63            Poll::Ready(Ok(())) if buf.filled().is_empty() => {
64                Poll::Ready(Ok(StreamResult::Dropped))
65            }
66            Poll::Ready(Ok(())) => {
67                let n = buf.filled().len();
68                dst.mark_written(n);
69                Poll::Ready(Ok(StreamResult::Completed))
70            }
71            Poll::Ready(Err(e)) => {
72                let _ = self
73                    .result_tx
74                    .take()
75                    .unwrap()
76                    .send(io_error_to_error_code(e));
77                Poll::Ready(Ok(StreamResult::Dropped))
78            }
79            Poll::Pending if finish => Poll::Ready(Ok(StreamResult::Cancelled)),
80            Poll::Pending => Poll::Pending,
81        }
82    }
83}
84
85struct OutputStreamConsumer {
86    tx: Pin<Box<dyn AsyncWrite + Send + Sync>>,
87    result_tx: Option<oneshot::Sender<ErrorCode>>,
88    flush_pending: bool,
89}
90
91impl OutputStreamConsumer {
92    fn poll_flush(
93        &mut self,
94        cx: &mut Context<'_>,
95        finish: bool,
96    ) -> Poll<wasmtime::Result<StreamResult>> {
97        match self.tx.as_mut().poll_flush(cx) {
98            Poll::Ready(Ok(())) => {
99                self.flush_pending = false;
100                Poll::Ready(Ok(StreamResult::Completed))
101            }
102            Poll::Ready(Err(e)) => self.dropped(e),
103            Poll::Pending => {
104                if finish {
105                    self.flush_pending = false;
106                    Poll::Ready(Ok(StreamResult::Cancelled))
107                } else {
108                    self.flush_pending = true;
109                    Poll::Pending
110                }
111            }
112        }
113    }
114
115    fn dropped(&mut self, err: io::Error) -> Poll<wasmtime::Result<StreamResult>> {
116        if let Some(tx) = self.result_tx.take() {
117            let _ = tx.send(io_error_to_error_code(err));
118        }
119        Poll::Ready(Ok(StreamResult::Dropped))
120    }
121}
122
123impl<D> StreamConsumer<D> for OutputStreamConsumer {
124    type Item = u8;
125
126    fn poll_consume(
127        mut self: Pin<&mut Self>,
128        cx: &mut Context<'_>,
129        store: StoreContextMut<D>,
130        src: Source<Self::Item>,
131        finish: bool,
132    ) -> Poll<wasmtime::Result<StreamResult>> {
133        if self.flush_pending {
134            return self.poll_flush(cx, finish);
135        }
136
137        let mut src = src.as_direct(store);
138        let buf = src.remaining();
139
140        // If the source buffer is empty then this is a request on behalf of
141        // the guest to wait for this output stream to be writable. The
142        // `AsyncWrite` trait abstraction does not provide the ability to await
143        // this event so we're forced to basically just lie here and say we're
144        // ready write data later.
145        //
146        // See WebAssembly/component-model#561 for some more information.
147        if buf.len() == 0 {
148            return Poll::Ready(Ok(StreamResult::Completed));
149        }
150        match self.tx.as_mut().poll_write(cx, buf) {
151            Poll::Ready(Ok(0)) => self.dropped(io::ErrorKind::WriteZero.into()),
152            Poll::Ready(Ok(n)) => {
153                src.mark_read(n);
154                self.poll_flush(cx, finish)
155            }
156            Poll::Ready(Err(e)) => self.dropped(e),
157            Poll::Pending if finish => Poll::Ready(Ok(StreamResult::Cancelled)),
158            Poll::Pending => Poll::Pending,
159        }
160    }
161}
162
163impl terminal_input::Host for WasiCliCtxView<'_> {}
164impl terminal_output::Host for WasiCliCtxView<'_> {}
165
166impl terminal_input::HostTerminalInput for WasiCliCtxView<'_> {
167    fn drop(&mut self, rep: Resource<TerminalInput>) -> wasmtime::Result<()> {
168        self.table
169            .delete(rep)
170            .context("failed to delete terminal input resource from table")?;
171        Ok(())
172    }
173}
174
175impl terminal_output::HostTerminalOutput for WasiCliCtxView<'_> {
176    fn drop(&mut self, rep: Resource<TerminalOutput>) -> wasmtime::Result<()> {
177        self.table
178            .delete(rep)
179            .context("failed to delete terminal output resource from table")?;
180        Ok(())
181    }
182}
183
184impl terminal_stdin::Host for WasiCliCtxView<'_> {
185    fn get_terminal_stdin(&mut self) -> wasmtime::Result<Option<Resource<TerminalInput>>> {
186        if self.ctx.stdin.is_terminal() {
187            let fd = self
188                .table
189                .push(TerminalInput)
190                .context("failed to push terminal stdin resource to table")?;
191            Ok(Some(fd))
192        } else {
193            Ok(None)
194        }
195    }
196}
197
198impl terminal_stdout::Host for WasiCliCtxView<'_> {
199    fn get_terminal_stdout(&mut self) -> wasmtime::Result<Option<Resource<TerminalOutput>>> {
200        if self.ctx.stdout.is_terminal() {
201            let fd = self
202                .table
203                .push(TerminalOutput)
204                .context("failed to push terminal stdout resource to table")?;
205            Ok(Some(fd))
206        } else {
207            Ok(None)
208        }
209    }
210}
211
212impl terminal_stderr::Host for WasiCliCtxView<'_> {
213    fn get_terminal_stderr(&mut self) -> wasmtime::Result<Option<Resource<TerminalOutput>>> {
214        if self.ctx.stderr.is_terminal() {
215            let fd = self
216                .table
217                .push(TerminalOutput)
218                .context("failed to push terminal stderr resource to table")?;
219            Ok(Some(fd))
220        } else {
221            Ok(None)
222        }
223    }
224}
225
226fn read_stdin(
227    mut store: impl AsContextMut,
228    stdin: Box<dyn AsyncRead + Send + Sync>,
229) -> wasmtime::Result<(StreamReader<u8>, FutureReader<Result<(), ErrorCode>>)> {
230    let mut store = store.as_context_mut();
231    let (result_tx, result_rx) = oneshot::channel();
232    let stream = StreamReader::new(
233        &mut store,
234        InputStreamProducer {
235            rx: Box::into_pin(stdin),
236            result_tx: Some(result_tx),
237        },
238    )?;
239    let future = FutureReader::new(&mut store, async {
240        wasmtime::error::Ok(match result_rx.await {
241            Ok(err) => Err(err),
242            Err(_) => Ok(()),
243        })
244    })?;
245    Ok((stream, future))
246}
247
248impl<U> stdin::HostWithStore<U> for WasiCli {
249    fn read_via_stream(
250        mut store: Access<U, Self>,
251    ) -> wasmtime::Result<(StreamReader<u8>, FutureReader<Result<(), ErrorCode>>)> {
252        let rx = store.get().ctx.stdin.async_stream();
253        read_stdin(&mut store, rx)
254    }
255}
256
257impl stdin::Host for WasiCliCtxView<'_> {}
258
259fn write_output(
260    mut store: impl AsContextMut,
261    data: StreamReader<u8>,
262    writer: Box<dyn AsyncWrite + Send + Sync>,
263) -> wasmtime::Result<FutureReader<Result<(), ErrorCode>>> {
264    let (result_tx, result_rx) = oneshot::channel();
265    data.pipe(
266        &mut store,
267        OutputStreamConsumer {
268            tx: Box::into_pin(writer),
269            result_tx: Some(result_tx),
270            flush_pending: false,
271        },
272    )?;
273    FutureReader::new(&mut store, async {
274        wasmtime::error::Ok(match result_rx.await {
275            Ok(err) => Err(err),
276            Err(_) => Ok(()),
277        })
278    })
279}
280
281impl<U> stdout::HostWithStore<U> for WasiCli {
282    fn write_via_stream(
283        mut store: Access<'_, U, Self>,
284        data: StreamReader<u8>,
285    ) -> wasmtime::Result<FutureReader<Result<(), ErrorCode>>> {
286        let tx = store.get().ctx.stdout.async_stream();
287        write_output(store, data, tx)
288    }
289}
290
291impl stdout::Host for WasiCliCtxView<'_> {}
292
293impl<U> stderr::HostWithStore<U> for WasiCli {
294    fn write_via_stream(
295        mut store: Access<'_, U, Self>,
296        data: StreamReader<u8>,
297    ) -> wasmtime::Result<FutureReader<Result<(), ErrorCode>>> {
298        let tx = store.get().ctx.stderr.async_stream();
299        write_output(store, data, tx)
300    }
301}
302
303impl stderr::Host for WasiCliCtxView<'_> {}
304
305impl environment::Host for WasiCliCtxView<'_> {
306    fn get_environment(&mut self) -> wasmtime::Result<Vec<(String, String)>> {
307        Ok(self.ctx.environment.clone())
308    }
309
310    fn get_arguments(&mut self) -> wasmtime::Result<Vec<String>> {
311        Ok(self.ctx.arguments.clone())
312    }
313
314    fn get_initial_cwd(&mut self) -> wasmtime::Result<Option<String>> {
315        Ok(self.ctx.initial_cwd.clone())
316    }
317}
318
319impl exit::Host for WasiCliCtxView<'_> {
320    fn exit(&mut self, status: Result<(), ()>) -> wasmtime::Result<()> {
321        let status = match status {
322            Ok(()) => 0,
323            Err(()) => 1,
324        };
325        Err(format_err!(I32Exit(status)))
326    }
327
328    fn exit_with_code(&mut self, status_code: u8) -> wasmtime::Result<()> {
329        Err(format_err!(I32Exit(status_code.into())))
330    }
331}
332
333mod named {
334    use crate::cli::{WasiCliNamed, WasiCliNamedView};
335    use crate::p3::bindings::cli::types::ErrorCode;
336    use crate::p3::bindings::named_imports::wasi::cli::{
337        environment, exit, stderr, stdin, stdout, terminal_input, terminal_output, terminal_stderr,
338        terminal_stdin, terminal_stdout,
339    };
340    use crate::p3::cli::{TerminalInput, TerminalOutput};
341    use crate::{NamedId, WasiCtxNamedView};
342    use wasmtime::component::{Access, FutureReader, Resource, StreamReader};
343
344    impl<T> exit::Host for WasiCtxNamedView<'_, T>
345    where
346        T: WasiCliNamedView,
347    {
348        fn exit(&mut self, id: NamedId, status: Result<(), ()>) -> wasmtime::Result<()> {
349            super::exit::Host::exit(&mut self.0.cli(id), status)
350        }
351
352        fn exit_with_code(&mut self, id: NamedId, status_code: u8) -> wasmtime::Result<()> {
353            super::exit::Host::exit_with_code(&mut self.0.cli(id), status_code)
354        }
355    }
356
357    impl<T> terminal_input::Host for WasiCtxNamedView<'_, T> where T: WasiCliNamedView {}
358    impl<T> terminal_output::Host for WasiCtxNamedView<'_, T> where T: WasiCliNamedView {}
359
360    impl<T> terminal_input::HostTerminalInput for WasiCtxNamedView<'_, T>
361    where
362        T: WasiCliNamedView,
363    {
364        fn drop(&mut self, id: NamedId, rep: Resource<TerminalInput>) -> wasmtime::Result<()> {
365            super::terminal_input::HostTerminalInput::drop(&mut self.0.cli(id), rep)
366        }
367    }
368
369    impl<T> terminal_output::HostTerminalOutput for WasiCtxNamedView<'_, T>
370    where
371        T: WasiCliNamedView,
372    {
373        fn drop(&mut self, id: NamedId, rep: Resource<TerminalOutput>) -> wasmtime::Result<()> {
374            super::terminal_output::HostTerminalOutput::drop(&mut self.0.cli(id), rep)
375        }
376    }
377
378    impl<T> terminal_stdin::Host for WasiCtxNamedView<'_, T>
379    where
380        T: WasiCliNamedView,
381    {
382        fn get_terminal_stdin(
383            &mut self,
384            id: NamedId,
385        ) -> wasmtime::Result<Option<Resource<TerminalInput>>> {
386            super::terminal_stdin::Host::get_terminal_stdin(&mut self.0.cli(id))
387        }
388    }
389
390    impl<T> terminal_stdout::Host for WasiCtxNamedView<'_, T>
391    where
392        T: WasiCliNamedView,
393    {
394        fn get_terminal_stdout(
395            &mut self,
396            id: NamedId,
397        ) -> wasmtime::Result<Option<Resource<TerminalOutput>>> {
398            super::terminal_stdout::Host::get_terminal_stdout(&mut self.0.cli(id))
399        }
400    }
401
402    impl<T> terminal_stderr::Host for WasiCtxNamedView<'_, T>
403    where
404        T: WasiCliNamedView,
405    {
406        fn get_terminal_stderr(
407            &mut self,
408            id: NamedId,
409        ) -> wasmtime::Result<Option<Resource<TerminalOutput>>> {
410            super::terminal_stderr::Host::get_terminal_stderr(&mut self.0.cli(id))
411        }
412    }
413
414    impl<T, U> stdin::HostWithStore<U> for WasiCliNamed<T>
415    where
416        T: WasiCliNamedView,
417    {
418        fn read_via_stream(
419            mut store: Access<U, Self>,
420            id: NamedId,
421        ) -> wasmtime::Result<(StreamReader<u8>, FutureReader<Result<(), ErrorCode>>)> {
422            let rx = store.get().0.cli(id).ctx.stdin.async_stream();
423            super::read_stdin(&mut store, rx)
424        }
425    }
426
427    impl<T> stdin::Host for WasiCtxNamedView<'_, T> where T: WasiCliNamedView {}
428
429    impl<T, U> stdout::HostWithStore<U> for WasiCliNamed<T>
430    where
431        T: WasiCliNamedView,
432    {
433        fn write_via_stream(
434            mut store: Access<'_, U, Self>,
435            id: NamedId,
436            data: StreamReader<u8>,
437        ) -> wasmtime::Result<FutureReader<Result<(), ErrorCode>>> {
438            let tx = store.get().0.cli(id).ctx.stdout.async_stream();
439            super::write_output(store, data, tx)
440        }
441    }
442
443    impl<T> stdout::Host for WasiCtxNamedView<'_, T> where T: WasiCliNamedView {}
444
445    impl<T, U> stderr::HostWithStore<U> for WasiCliNamed<T>
446    where
447        T: WasiCliNamedView,
448    {
449        fn write_via_stream(
450            mut store: Access<'_, U, Self>,
451            id: NamedId,
452            data: StreamReader<u8>,
453        ) -> wasmtime::Result<FutureReader<Result<(), ErrorCode>>> {
454            let tx = store.get().0.cli(id).ctx.stderr.async_stream();
455            super::write_output(store, data, tx)
456        }
457    }
458
459    impl<T> stderr::Host for WasiCtxNamedView<'_, T> where T: WasiCliNamedView {}
460
461    impl<T> environment::Host for WasiCtxNamedView<'_, T>
462    where
463        T: WasiCliNamedView,
464    {
465        fn get_environment(&mut self, id: NamedId) -> wasmtime::Result<Vec<(String, String)>> {
466            super::environment::Host::get_environment(&mut self.0.cli(id))
467        }
468
469        fn get_arguments(&mut self, id: NamedId) -> wasmtime::Result<Vec<String>> {
470            super::environment::Host::get_arguments(&mut self.0.cli(id))
471        }
472
473        fn get_initial_cwd(&mut self, id: NamedId) -> wasmtime::Result<Option<String>> {
474            super::environment::Host::get_initial_cwd(&mut self.0.cli(id))
475        }
476    }
477}
wasmtime_wasi/p3/cli/host.rs

wasmtime_wasi/p3/cli/
host.rs
