wasmtime_wasi/cli/
file.rs

1use crate::cli::{IsTerminal, StdinStream, StdoutStream};
2use crate::p2::{InputStream, OutputStream, Pollable, StreamError, StreamResult};
3use bytes::Bytes;
4use std::io::{Read, Write};
5use std::pin::Pin;
6use std::sync::Arc;
7use std::task::{Context, Poll};
8use tokio::io::{self, AsyncRead, AsyncWrite};
9
10/// This implementation will yield output streams that block on writes, and
11/// output directly to a file. If truly async output is required,
12/// [`AsyncStdoutStream`] should be used instead.
13///
14/// [`AsyncStdoutStream`]: crate::cli::AsyncStdoutStream
15#[derive(Clone)]
16pub struct OutputFile {
17    file: Arc<std::fs::File>,
18}
19
20impl OutputFile {
21    pub fn new(file: std::fs::File) -> Self {
22        Self {
23            file: Arc::new(file),
24        }
25    }
26}
27
28impl IsTerminal for OutputFile {
29    fn is_terminal(&self) -> bool {
30        false
31    }
32}
33
34impl StdoutStream for OutputFile {
35    fn p2_stream(&self) -> Box<dyn OutputStream> {
36        Box::new(self.clone())
37    }
38
39    fn async_stream(&self) -> Box<dyn AsyncWrite + Send + Sync> {
40        Box::new(self.clone())
41    }
42}
43
44#[async_trait::async_trait]
45impl Pollable for OutputFile {
46    async fn ready(&mut self) {}
47}
48
49impl OutputStream for OutputFile {
50    fn write(&mut self, bytes: Bytes) -> StreamResult<()> {
51        if bytes.len() > crate::MAX_READ_SIZE_ALLOC {
52            return Err(StreamError::Trap(wasmtime::format_err!(
53                "write exceeded budget"
54            )));
55        }
56        (&*self.file)
57            .write_all(&bytes)
58            .map_err(|e| StreamError::LastOperationFailed(wasmtime::format_err!(e)))
59    }
60
61    fn flush(&mut self) -> StreamResult<()> {
62        use std::io::Write;
63        self.file
64            .flush()
65            .map_err(|e| StreamError::LastOperationFailed(wasmtime::format_err!(e)))
66    }
67
68    fn check_write(&mut self) -> StreamResult<usize> {
69        Ok(crate::MAX_READ_SIZE_ALLOC)
70    }
71}
72
73impl AsyncWrite for OutputFile {
74    fn poll_write(
75        self: Pin<&mut Self>,
76        _cx: &mut Context<'_>,
77        buf: &[u8],
78    ) -> Poll<io::Result<usize>> {
79        match (&*self.file).write_all(buf) {
80            Ok(()) => Poll::Ready(Ok(buf.len())),
81            Err(e) => Poll::Ready(Err(e)),
82        }
83    }
84    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
85        Poll::Ready((&*self.file).flush())
86    }
87    fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
88        Poll::Ready(Ok(()))
89    }
90}
91
92/// This implementation will yield input streams that block on reads, and
93/// reads directly from a file. If truly async input is required,
94/// [`AsyncStdinStream`] should be used instead.
95///
96/// [`AsyncStdinStream`]: crate::cli::AsyncStdinStream
97#[derive(Clone)]
98pub struct InputFile {
99    file: Arc<std::fs::File>,
100}
101
102impl InputFile {
103    pub fn new(file: std::fs::File) -> Self {
104        Self {
105            file: Arc::new(file),
106        }
107    }
108}
109
110impl StdinStream for InputFile {
111    fn p2_stream(&self) -> Box<dyn InputStream> {
112        Box::new(self.clone())
113    }
114    fn async_stream(&self) -> Box<dyn AsyncRead + Send + Sync> {
115        Box::new(self.clone())
116    }
117}
118
119impl IsTerminal for InputFile {
120    fn is_terminal(&self) -> bool {
121        false
122    }
123}
124
125#[async_trait::async_trait]
126impl Pollable for InputFile {
127    async fn ready(&mut self) {}
128}
129
130impl InputStream for InputFile {
131    fn read(&mut self, size: usize) -> StreamResult<Bytes> {
132        let mut buf = bytes::BytesMut::zeroed(size.min(crate::MAX_READ_SIZE_ALLOC));
133        let bytes_read = self
134            .file
135            .read(&mut buf)
136            .map_err(|e| StreamError::LastOperationFailed(wasmtime::format_err!(e)))?;
137        if bytes_read == 0 {
138            return Err(StreamError::Closed);
139        }
140        buf.truncate(bytes_read);
141        StreamResult::Ok(buf.into())
142    }
143}
144
145impl AsyncRead for InputFile {
146    fn poll_read(
147        self: Pin<&mut Self>,
148        _cx: &mut Context<'_>,
149        buf: &mut io::ReadBuf<'_>,
150    ) -> Poll<io::Result<()>> {
151        match (&*self.file).read(buf.initialize_unfilled()) {
152            Ok(n) => {
153                buf.advance(n);
154                Poll::Ready(Ok(()))
155            }
156            Err(e) => Poll::Ready(Err(e)),
157        }
158    }
159}
160
161#[cfg(test)]
162mod tests {
163    use super::*;
164
165    /// A write larger than the budget `check_write` reports has to trap instead
166    /// of being written out.
167    #[test]
168    fn write_larger_than_budget_traps() {
169        let path = std::env::temp_dir().join("wasmtime-output-file-write-budget");
170        let file = std::fs::File::create(&path).unwrap();
171        let mut stream = OutputFile::new(file);
172        assert_eq!(stream.check_write().unwrap(), crate::MAX_READ_SIZE_ALLOC);
173        let err = stream
174            .write(Bytes::from(vec![0; crate::MAX_READ_SIZE_ALLOC + 1]))
175            .unwrap_err();
176        assert!(matches!(err, StreamError::Trap(_)), "{err:?}");
177    }
178}
wasmtime_wasi/cli/file.rs

wasmtime_wasi/cli/
file.rs
