Skip to main content

wasm_sandbox/communication/
io.rs

1//! I/O redirection for sandboxed processes
2
3use std::sync::{Arc, Mutex};
4use std::io::{self, Read, Write};
5
6use crate::error::Result;
7use crate::communication::CommunicationChannel;
8
9/// Standard input/output redirection
10pub struct StdioRedirection {
11    /// Standard input channel
12    stdin_channel: Arc<dyn CommunicationChannel>,
13    
14    /// Standard output channel
15    stdout_channel: Arc<dyn CommunicationChannel>,
16    
17    /// Standard error channel
18    stderr_channel: Arc<dyn CommunicationChannel>,
19    
20    /// Buffered stdin data
21    stdin_buffer: Mutex<Vec<u8>>,
22    
23    /// Is closed
24    closed: Mutex<bool>,
25}
26
27impl StdioRedirection {
28    /// Create a new stdio redirection
29    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    /// Write to stdout
44    pub fn write_stdout(&self, data: &[u8]) -> Result<()> {
45        self.stdout_channel.send_to_guest(data)
46    }
47    
48    /// Write to stderr
49    pub fn write_stderr(&self, data: &[u8]) -> Result<()> {
50        self.stderr_channel.send_to_guest(data)
51    }
52    
53    /// Read from stdin
54    pub fn read_stdin(&self, buf: &mut [u8]) -> Result<usize> {
55        // If closed, return EOF
56        if *self.closed.lock().unwrap() {
57            return Ok(0);
58        }
59        
60        // Try to fill the buffer if it's empty
61        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 we have data, copy it to the buffer
69        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            // No data available
76            Ok(0)
77        }
78    }
79    
80    /// Close the redirection
81    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
93/// Standard input implementation
94pub struct StdioInput {
95    /// Redirection
96    redirection: Arc<StdioRedirection>,
97}
98
99impl StdioInput {
100    /// Create a new stdin implementation
101    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
115/// Standard output implementation
116pub struct StdioOutput {
117    /// Redirection
118    redirection: Arc<StdioRedirection>,
119    
120    /// Whether this is stderr
121    is_stderr: bool,
122}
123
124impl StdioOutput {
125    /// Create a new stdout implementation
126    pub fn new_stdout(redirection: Arc<StdioRedirection>) -> Self {
127        Self {
128            redirection,
129            is_stderr: false,
130        }
131    }
132    
133    /// Create a new stderr implementation
134    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
161/// Factory for creating stdio redirections
162pub struct StdioFactory {
163    /// Channel factory
164    channel_factory: Arc<dyn crate::communication::CommunicationFactory>,
165}
166
167impl StdioFactory {
168    /// Create a new stdio factory
169    pub fn new(channel_factory: Arc<dyn crate::communication::CommunicationFactory>) -> Self {
170        Self {
171            channel_factory,
172        }
173    }
174    
175    /// Create a new stdio redirection
176    pub fn create_redirection(&self) -> Result<Arc<StdioRedirection>> {
177        // Create the channels
178        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        // Create the redirection
183        let redirection = StdioRedirection::new(stdin_channel, stdout_channel, stderr_channel);
184        
185        Ok(Arc::new(redirection))
186    }
187}