wasi_shell/
lib.rs

1use std::env;
2use std::fs::File;
3use std::io::{self, Read, Write};
4use std::sync::{Arc, Mutex};
5use colored::*;
6pub use wasibox_core::IoContext;
7
8pub fn handle_pipeline<R, W, F>(
9    line: &str, 
10    initial_stdin: Box<R>, 
11    final_stdout: Box<W>,
12    handler: &mut F,
13) -> Result<(), String> 
14where 
15    R: Read + Send + 'static,
16    W: Write + Send + 'static,
17    F: FnMut(&[String], &mut IoContext) -> Option<Result<(), String>> + Send,
18{
19    let stages_str: Vec<&str> = line.split('|').collect();
20    let handler = Arc::new(Mutex::new(handler));
21
22    let mut final_stdout = Some(final_stdout);
23
24    std::thread::scope(|s| {
25        let mut prev_reader: Box<dyn Read + Send> = Box::new(initial_stdin);
26        let mut threads = Vec::new();
27
28        for (i, stage_str) in stages_str.iter().enumerate() {
29            let is_last = i == stages_str.len() - 1;
30            let mut tokens = shlex::split(stage_str.trim())
31                .ok_or_else(|| "Error: Invalid input (quoting error)".to_string())?;
32
33            let stdin = std::mem::replace(&mut prev_reader, Box::new(io::empty()));
34            let (stdout, next_reader): (Box<dyn Write + Send>, Option<Box<dyn Read + Send>>) = if is_last {
35                let mut out: Box<dyn Write + Send> = Box::new(final_stdout.take().unwrap());
36                if let Some(pos) = tokens.iter().position(|t| t == ">" || t == ">>") {
37                    let append = tokens[pos] == ">>";
38                    let filename = tokens.get(pos + 1).ok_or("Error: Missing file for redirection")?;
39                    let file = if append {
40                        std::fs::OpenOptions::new().create(true).append(true).open(filename)
41                    } else {
42                        File::create(filename)
43                    }.map_err(|e| format!("Error opening file: {}", e))?;
44                    out = Box::new(file);
45                    tokens.truncate(pos);
46                }
47                (out, None)
48            } else {
49                let (reader, writer) = create_pipe();
50                (Box::new(writer), Some(Box::new(reader)))
51            };
52
53            if let Some(r) = next_reader {
54                prev_reader = r;
55            }
56
57            let handler = Arc::clone(&handler);
58            let thread = s.spawn(move || {
59                let mut ctx = IoContext::new(stdin, stdout, Box::new(io::stderr()));
60                let mut h = |args: &[String], ctx: &mut IoContext| {
61                    let mut h_lock = handler.lock().unwrap();
62                    (h_lock)(args, ctx)
63                };
64                execute_command(tokens, &mut ctx, &mut h)
65            });
66            threads.push(thread);
67        }
68
69        let mut final_res = Ok(());
70        for thread in threads {
71            if let Ok(res) = thread.join() {
72                if res.is_err() && final_res.is_ok() {
73                    final_res = res;
74                }
75            }
76        }
77        final_res
78    })
79}
80
81struct PipeReader {
82    rx: std::sync::mpsc::Receiver<Vec<u8>>,
83    buffer: Vec<u8>,
84    pos: usize,
85}
86impl Read for PipeReader {
87    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
88        if self.pos >= self.buffer.len() {
89            match self.rx.recv() {
90                Ok(data) => {
91                    self.buffer = data;
92                    self.pos = 0;
93                }
94                Err(_) => return Ok(0),
95            }
96        }
97        let available = self.buffer.len() - self.pos;
98        let to_copy = std::cmp::min(available, buf.len());
99        buf[..to_copy].copy_from_slice(&self.buffer[self.pos..self.pos + to_copy]);
100        self.pos += to_copy;
101        Ok(to_copy)
102    }
103}
104
105struct PipeWriter {
106    tx: std::sync::mpsc::Sender<Vec<u8>>,
107}
108impl Write for PipeWriter {
109    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
110        self.tx.send(buf.to_vec()).map_err(|e| io::Error::new(io::ErrorKind::BrokenPipe, e))?;
111        Ok(buf.len())
112    }
113    fn flush(&mut self) -> io::Result<()> { Ok(()) }
114}
115
116fn create_pipe() -> (PipeReader, PipeWriter) {
117    let (tx, rx) = std::sync::mpsc::channel();
118    (PipeReader { rx, buffer: Vec::new(), pos: 0 }, PipeWriter { tx })
119}
120
121pub struct ArcVecWriter {
122    pub inner: Arc<Mutex<Vec<u8>>>,
123}
124impl Write for ArcVecWriter {
125    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
126        let mut inner = self.inner.lock().map_err(|e| io::Error::new(io::ErrorKind::Other, e.to_string()))?;
127        inner.extend_from_slice(buf);
128        Ok(buf.len())
129    }
130    fn flush(&mut self) -> io::Result<()> { Ok(()) }
131}
132
133pub fn execute_command<F>(
134    args: Vec<String>, 
135    ctx: &mut IoContext,
136    handler: &mut F,
137) -> Result<(), String> 
138where
139    F: FnMut(&[String], &mut IoContext) -> Option<Result<(), String>>,
140{
141    if args.is_empty() {
142        return Ok(());
143    }
144
145    if let Some(res) = handler(&args, ctx) {
146        return res;
147    }
148
149    let cmd = &args[0];
150
151    match cmd.as_str() {
152        "help" => {
153            writeln!(ctx.stdout, "{}", "Available Commands:".yellow().bold()).map_err(|e| e.to_string())?;
154            writeln!(ctx.stdout, "  Shell Built-ins: cd, help, exit").map_err(|e| e.to_string())?;
155            writeln!(ctx.stdout, "  Animations: sl").map_err(|e| e.to_string())?;
156            writeln!(ctx.stdout, "  Core Utilities: arch, cat, cp, dir, echo, grep, head, ls, mkdir, mv, pwd, rm, rmdir, sleep, tail, tee, touch, tree, uname, wc, whoami, yes, etc.").map_err(|e| e.to_string())?;
157            Ok(())
158        }
159        "cd" => {
160            let new_dir = args.get(1).map(|s| s.as_str()).unwrap_or(".");
161            env::set_current_dir(new_dir).map_err(|e| format!("cd: {}", e))
162        }
163        "exit" => Ok(()),
164        "sl" => {
165            let _ = sl::run(args.iter().cloned());
166            Ok(())
167        }
168        _ => wasibox_core::execute_with_context(args.iter().cloned(), ctx),
169    }
170}
171
172pub fn handle_parallel<F>(
173    lines: Vec<String>,
174    handler: Arc<F>,
175) -> Vec<Result<(), String>>
176where
177    F: Fn(&[String], &mut IoContext) -> Option<Result<(), String>> + Send + Sync + 'static,
178{
179    let mut handles = Vec::new();
180    for line in lines {
181        let handler = Arc::clone(&handler);
182        let handle = std::thread::spawn(move || {
183            let mut h = |args: &[String], ctx: &mut IoContext| handler(args, ctx);
184            handle_pipeline(
185                &line, 
186                Box::new(io::empty()), 
187                Box::new(io::sink()), 
188                &mut h
189            )
190        });
191        handles.push(handle);
192    }
193
194    handles.into_iter()
195        .map(|h| h.join().unwrap_or(Err("Thread panicked".to_string())))
196        .collect()
197}
198
199#[cfg(test)]
200mod tests {
201    use super::*;
202    use std::io::Cursor;
203
204    fn get_temp_dir() -> tempfile::TempDir {
205        #[cfg(target_os = "wasi")]
206        {
207            // Avoid any logic that might indirectly call env::temp_dir()
208            tempfile::Builder::new()
209                .prefix("test_")
210                .tempdir_in(".")
211                .expect("Failed to create temp dir in current directory")
212        }
213        #[cfg(not(target_os = "wasi"))]
214        {
215            tempfile::tempdir().expect("Failed to create system temp dir")
216        }
217    }
218
219    #[test]
220    fn test_simple_echo() {
221        let out = Arc::new(Mutex::new(Vec::new()));
222        handle_pipeline("echo hello", Box::new(Cursor::new("")), Box::new(ArcVecWriter { inner: Arc::clone(&out) }), &mut |_, _| None).unwrap();
223        let buf = out.lock().unwrap();
224        assert_eq!(String::from_utf8_lossy(&buf).trim(), "hello");
225    }
226
227    #[test]
228    fn test_pipe_echo_cat() {
229        let out = Arc::new(Mutex::new(Vec::new()));
230        handle_pipeline("echo hello | cat", Box::new(Cursor::new("")), Box::new(ArcVecWriter { inner: Arc::clone(&out) }), &mut |_, _| None).unwrap();
231        let buf = out.lock().unwrap();
232        assert_eq!(String::from_utf8_lossy(&buf).trim(), "hello");
233    }
234
235    #[test]
236    fn test_pipe_grep() {
237        let out = Arc::new(Mutex::new(Vec::new()));
238        // Our 'echo' likely doesn't support -e. Let's use a simpler check.
239        handle_pipeline("echo world | grep world", Box::new(Cursor::new("")), Box::new(ArcVecWriter { inner: Arc::clone(&out) }), &mut |_, _| None).unwrap();
240        let buf = out.lock().unwrap();
241        assert_eq!(String::from_utf8_lossy(&buf).trim(), "world");
242    }
243
244    #[test]
245    fn test_redirection_create() {
246        let dir = get_temp_dir();
247        let file_path = dir.path().join("test_create.txt");
248        let cmd = format!("echo hello > \"{}\"", file_path.display());
249        
250        handle_pipeline(&cmd, Box::new(Cursor::new("")), Box::new(io::sink()), &mut |_, _| None).unwrap();
251        
252        let content = std::fs::read_to_string(file_path).unwrap();
253        assert_eq!(content.trim(), "hello");
254    }
255
256    #[test]
257    fn test_redirection_append() {
258        let dir = get_temp_dir();
259        let file_path = dir.path().join("test_append.txt");
260        
261        let cmd1 = format!("echo hello > \"{}\"", file_path.display());
262        handle_pipeline(&cmd1, Box::new(Cursor::new("")), Box::new(io::sink()), &mut |_, _| None).unwrap();
263        
264        let cmd2 = format!("echo world >> \"{}\"", file_path.display());
265        handle_pipeline(&cmd2, Box::new(Cursor::new("")), Box::new(io::sink()), &mut |_, _| None).unwrap();
266        
267        let content = std::fs::read_to_string(file_path).unwrap();
268        assert!(content.contains("hello"));
269        assert!(content.contains("world"));
270    }
271
272    #[test]
273    fn test_complex_pipeline() {
274        let out = Arc::new(Mutex::new(Vec::new()));
275        handle_pipeline("echo hello | grep h | wc -c", Box::new(Cursor::new("")), Box::new(ArcVecWriter { inner: Arc::clone(&out) }), &mut |_, _| None).unwrap();
276        let buf = out.lock().unwrap();
277        let result = String::from_utf8_lossy(&buf).trim().to_string();
278        assert_eq!(result, "6");
279    }
280    #[test]
281    fn test_custom_handler() {
282        let out = Arc::new(Mutex::new(Vec::new()));
283        let mut handler = |args: &[String], ctx: &mut IoContext| {
284            if args[0] == "magic" {
285                write!(ctx.stdout, "magic happen").unwrap();
286                return Some(Ok(()));
287            }
288            None
289        };
290        handle_pipeline("magic", Box::new(io::empty()), Box::new(ArcVecWriter { inner: Arc::clone(&out) }), &mut handler).unwrap();
291        let buf = out.lock().unwrap();
292        assert_eq!(String::from_utf8_lossy(&buf), "magic happen");
293    }
294
295    #[test]
296    fn test_parallel_execution() {
297        let handler = Arc::new(|args: &[String], _ctx: &mut IoContext| {
298            if args[0] == "slow" {
299                std::thread::sleep(std::time::Duration::from_millis(100));
300                return Some(Ok(()));
301            }
302            None
303        });
304        let lines = vec!["slow".to_string(), "slow".to_string(), "slow".to_string()];
305        let start = std::time::Instant::now();
306        let results = handle_parallel(lines, handler);
307        let duration = start.elapsed();
308        
309        assert_eq!(results.len(), 3);
310        for res in results {
311            assert!(res.is_ok());
312        }
313        // If parallel, it should take less than 300ms.
314        assert!(duration < std::time::Duration::from_millis(250));
315    }
316
317    #[test]
318    fn test_streaming_pipeline() {
319        let out = Arc::new(Mutex::new(Vec::new()));
320        handle_pipeline("yes | head -n 2", Box::new(io::empty()), Box::new(ArcVecWriter { inner: Arc::clone(&out) }), &mut |_, _| None).unwrap();
321        let buf = out.lock().unwrap();
322        let result = String::from_utf8_lossy(&buf);
323        assert_eq!(result.trim(), "y\ny");
324    }
325}
wasi_shell/lib.rs

wasi_shell/
lib.rs
