wasm_sandbox/communication/
io.rs1use std::sync::{Arc, Mutex};
4use std::io::{self, Read, Write};
5
6use crate::error::Result;
7use crate::communication::CommunicationChannel;
8
9pub struct StdioRedirection {
11 stdin_channel: Arc<dyn CommunicationChannel>,
13
14 stdout_channel: Arc<dyn CommunicationChannel>,
16
17 stderr_channel: Arc<dyn CommunicationChannel>,
19
20 stdin_buffer: Mutex<Vec<u8>>,
22
23 closed: Mutex<bool>,
25}
26
27impl StdioRedirection {
28 pub fn new(
30 stdin_channel: Arc<dyn CommunicationChannel>,
31 stdout_channel: Arc<dyn CommunicationChannel>,
32 stderr_channel: Arc<dyn CommunicationChannel>,
33 ) -> Self {
34 Self {
35 stdin_channel,
36 stdout_channel,
37 stderr_channel,
38 stdin_buffer: Mutex::new(Vec::new()),
39 closed: Mutex::new(false),
40 }
41 }
42
43 pub fn write_stdout(&self, data: &[u8]) -> Result<()> {
45 self.stdout_channel.send_to_guest(data)
46 }
47
48 pub fn write_stderr(&self, data: &[u8]) -> Result<()> {
50 self.stderr_channel.send_to_guest(data)
51 }
52
53 pub fn read_stdin(&self, buf: &mut [u8]) -> Result<usize> {
55 if *self.closed.lock().unwrap() {
57 return Ok(0);
58 }
59
60 let mut stdin_buffer = self.stdin_buffer.lock().unwrap();
62 if stdin_buffer.is_empty() {
63 if let Ok(data) = self.stdin_channel.receive_from_guest() {
64 stdin_buffer.extend_from_slice(&data);
65 }
66 }
67
68 if !stdin_buffer.is_empty() {
70 let n = std::cmp::min(buf.len(), stdin_buffer.len());
71 buf[..n].copy_from_slice(&stdin_buffer[..n]);
72 stdin_buffer.drain(..n);
73 Ok(n)
74 } else {
75 Ok(0)
77 }
78 }
79
80 pub fn close(&self) -> Result<()> {
82 let mut closed = self.closed.lock().unwrap();
83 *closed = true;
84
85 self.stdin_channel.close()?;
86 self.stdout_channel.close()?;
87 self.stderr_channel.close()?;
88
89 Ok(())
90 }
91}
92
93pub struct StdioInput {
95 redirection: Arc<StdioRedirection>,
97}
98
99impl StdioInput {
100 pub fn new(redirection: Arc<StdioRedirection>) -> Self {
102 Self {
103 redirection,
104 }
105 }
106}
107
108impl Read for StdioInput {
109 fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
110 self.redirection.read_stdin(buf)
111 .map_err(|e| io::Error::other(format!("{:?}", e)))
112 }
113}
114
115pub struct StdioOutput {
117 redirection: Arc<StdioRedirection>,
119
120 is_stderr: bool,
122}
123
124impl StdioOutput {
125 pub fn new_stdout(redirection: Arc<StdioRedirection>) -> Self {
127 Self {
128 redirection,
129 is_stderr: false,
130 }
131 }
132
133 pub fn new_stderr(redirection: Arc<StdioRedirection>) -> Self {
135 Self {
136 redirection,
137 is_stderr: true,
138 }
139 }
140}
141
142impl Write for StdioOutput {
143 fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
144 let result = if self.is_stderr {
145 self.redirection.write_stderr(buf)
146 } else {
147 self.redirection.write_stdout(buf)
148 };
149
150 match result {
151 Ok(_) => Ok(buf.len()),
152 Err(e) => Err(io::Error::other(format!("{:?}", e))),
153 }
154 }
155
156 fn flush(&mut self) -> io::Result<()> {
157 Ok(())
158 }
159}
160
161pub struct StdioFactory {
163 channel_factory: Arc<dyn crate::communication::CommunicationFactory>,
165}
166
167impl StdioFactory {
168 pub fn new(channel_factory: Arc<dyn crate::communication::CommunicationFactory>) -> Self {
170 Self {
171 channel_factory,
172 }
173 }
174
175 pub fn create_redirection(&self) -> Result<Arc<StdioRedirection>> {
177 let stdin_channel = self.channel_factory.create_channel()?;
179 let stdout_channel = self.channel_factory.create_channel()?;
180 let stderr_channel = self.channel_factory.create_channel()?;
181
182 let redirection = StdioRedirection::new(stdin_channel, stdout_channel, stderr_channel);
184
185 Ok(Arc::new(redirection))
186 }
187}