Skip to main content

wasi_shell/
lib.rs

1use std::collections::HashMap;
2use std::env;
3use std::fs::File;
4use std::io::{self, Read, Write};
5use std::sync::{Arc, Mutex};
6use colored::*;
7pub use wasibox_core::IoContext;
8
9mod readline;
10pub use readline::*;
11
12// ---------------------------------------------------------------------------
13// CommandRegistry
14// ---------------------------------------------------------------------------
15
16/// A function that handles a shell command.
17///
18/// Receives the full argument list (argv\[0\] = command name) and an I/O context
19/// whose stdin/stdout/stderr are already wired to the pipeline.
20pub type CommandFn = Arc<dyn Fn(&[String], &mut IoContext) -> Result<(), String> + Send + Sync>;
21
22/// Registry of named commands.
23///
24/// Commands registered here are used by [`handle_pipeline`] and
25/// [`handle_parallel`] to dispatch each stage of a pipeline.
26///
27/// # Examples
28///
29/// ```ignore
30/// use wasi_shell::{CommandRegistry, handle_pipeline, IoContext};
31/// use std::io::{self, Write};
32///
33/// let mut reg = CommandRegistry::with_builtins();
34///
35/// // Add a custom command
36/// reg.register("greet", |args, ctx| {
37///     let name = args.get(1).map(|s| s.as_str()).unwrap_or("world");
38///     writeln!(ctx.stdout, "Hello, {}!", name).map_err(|e| e.to_string())
39/// });
40///
41/// handle_pipeline("greet Rust", Box::new(io::empty()), Box::new(io::stdout()), &reg, wasibox_core::CancellationToken::new()).unwrap();
42/// ```
43pub struct CommandRegistry {
44    commands: HashMap<String, CommandFn>,
45    fallback: Option<CommandFn>,
46}
47
48impl CommandRegistry {
49    /// Create an empty registry **with** a `wasibox_core` fallback.
50    ///
51    /// Any command name not explicitly registered will be forwarded to
52    /// `wasibox_core::execute_with_context`.
53    pub fn new() -> Self {
54        Self {
55            commands: HashMap::new(),
56            fallback: Some(Arc::new(|args: &[String], ctx: &mut IoContext| {
57                wasibox_core::execute_with_context(args.iter().cloned(), ctx)
58            })),
59        }
60    }
61
62    /// Create a registry pre-loaded with shell built-in commands
63    /// (`cd`, `help`, `exit`, `sl`) and a `wasibox_core` fallback for all
64    /// core utilities (echo, cat, grep, seq, head, …).
65    pub fn with_builtins() -> Self {
66        let mut reg = Self::new();
67
68        reg.register("help", |_args, ctx| {
69            writeln!(ctx.stdout, "{}", "Available Commands:".yellow().bold()).map_err(|e| e.to_string())?;
70
71            let mut shell_builtins = vec!["cd", "help", "exit"];
72            #[cfg(feature = "clear")] shell_builtins.push("clear");
73            shell_builtins.sort();
74            writeln!(ctx.stdout, "  Shell Built-ins: {}", shell_builtins.join(", ")).map_err(|e| e.to_string())?;
75
76            #[cfg(feature = "sl")]
77            writeln!(ctx.stdout, "  Animations: sl").map_err(|e| e.to_string())?;
78
79            let mut utils = Vec::new();
80            #[cfg(feature = "arch")] utils.push("arch");
81            #[cfg(feature = "basename")] utils.push("basename");
82            #[cfg(feature = "cat")] utils.push("cat");
83            #[cfg(feature = "cp")] utils.push("cp");
84            #[cfg(feature = "dir")] utils.push("dir");
85            #[cfg(feature = "dirname")] utils.push("dirname");
86            #[cfg(feature = "echo")] utils.push("echo");
87            #[cfg(feature = "env")] utils.push("env");
88            #[cfg(feature = "false")] utils.push("false");
89            #[cfg(feature = "grep")] utils.push("grep");
90            #[cfg(feature = "head")] utils.push("head");
91            #[cfg(feature = "link")] utils.push("link");
92            #[cfg(feature = "ln")] utils.push("ln");
93            #[cfg(feature = "ls")] utils.push("ls");
94            #[cfg(feature = "mkdir")] utils.push("mkdir");
95            #[cfg(feature = "mv")] utils.push("mv");
96            #[cfg(feature = "pwd")] utils.push("pwd");
97            #[cfg(feature = "rm")] utils.push("rm");
98            #[cfg(feature = "rmdir")] utils.push("rmdir");
99            #[cfg(feature = "seq")] utils.push("seq");
100            #[cfg(feature = "sleep")] utils.push("sleep");
101            #[cfg(feature = "tail")] utils.push("tail");
102            #[cfg(feature = "tee")] utils.push("tee");
103            #[cfg(feature = "touch")] utils.push("touch");
104            #[cfg(feature = "tree")] utils.push("tree");
105            #[cfg(feature = "true")] utils.push("true");
106            #[cfg(feature = "uname")] utils.push("uname");
107            #[cfg(feature = "unlink")] utils.push("unlink");
108            #[cfg(feature = "wc")] utils.push("wc");
109            #[cfg(feature = "whoami")] utils.push("whoami");
110            #[cfg(feature = "yes")] utils.push("yes");
111
112            if !utils.is_empty() {
113                utils.sort();
114                writeln!(ctx.stdout, "  Core Utilities: {}", utils.join(", ")).map_err(|e| e.to_string())?;
115            }
116            Ok(())
117        });
118
119        reg.register("cd", |args, _ctx| {
120            let new_dir = args.get(1).map(|s| s.as_str()).unwrap_or(".");
121            env::set_current_dir(new_dir).map_err(|e| format!("cd: {}", e))
122        });
123
124        reg.register("exit", |_args, _ctx| Ok(()));
125
126        #[cfg(feature = "clear")]
127        reg.register("clear", |_args, ctx| {
128            write!(ctx.stdout, "\x1B[2J\x1B[1;1H").map_err(|e| e.to_string())?;
129            ctx.stdout.flush().map_err(|e| e.to_string())
130        });
131
132        #[cfg(feature = "sl")]
133        reg.register("sl", |args, ctx| {
134            let _ = sl::run_with_token(args.iter().cloned(), Some(ctx.cancel_token.clone()));
135            Ok(())
136        });
137
138        // Register coreutils from wasibox-core
139        #[cfg(feature = "arch")]
140        reg.register("arch", |args, ctx| wasibox_core::utils::arch::execute_with_context(args.iter().cloned(), ctx));
141        #[cfg(feature = "basename")]
142        reg.register("basename", |args, ctx| wasibox_core::utils::basename::execute_with_context(args.iter().cloned(), ctx));
143        #[cfg(feature = "cat")]
144        reg.register("cat", |args, ctx| wasibox_core::utils::cat::execute_with_context(args.iter().cloned(), ctx));
145        #[cfg(feature = "cp")]
146        reg.register("cp", |args, ctx| wasibox_core::utils::cp::execute_with_context(args.iter().cloned(), ctx));
147        #[cfg(feature = "dir")]
148        reg.register("dir", |args, ctx| wasibox_core::utils::dir::execute_with_context(args.iter().cloned(), ctx));
149        #[cfg(feature = "dirname")]
150        reg.register("dirname", |args, ctx| wasibox_core::utils::dirname::execute_with_context(args.iter().cloned(), ctx));
151        #[cfg(feature = "echo")]
152        reg.register("echo", |args, ctx| wasibox_core::utils::echo::execute_with_context(args.iter().cloned(), ctx));
153        #[cfg(feature = "env")]
154        reg.register("env", |args, ctx| wasibox_core::utils::env::execute_with_context(args.iter().cloned(), ctx));
155        #[cfg(feature = "false")]
156        reg.register("false", |args, ctx| wasibox_core::utils::r#false::execute_with_context(args.iter().cloned(), ctx));
157        #[cfg(feature = "grep")]
158        reg.register("grep", |args, ctx| wasibox_core::utils::grep::execute_with_context(args.iter().cloned(), ctx));
159        #[cfg(feature = "head")]
160        reg.register("head", |args, ctx| wasibox_core::utils::head::execute_with_context(args.iter().cloned(), ctx));
161        #[cfg(feature = "link")]
162        reg.register("link", |args, ctx| wasibox_core::utils::link::execute_with_context(args.iter().cloned(), ctx));
163        #[cfg(feature = "ln")]
164        reg.register("ln", |args, ctx| wasibox_core::utils::ln::execute_with_context(args.iter().cloned(), ctx));
165        #[cfg(feature = "ls")]
166        reg.register("ls", |args, ctx| wasibox_core::utils::ls::execute_with_context(args.iter().cloned(), ctx));
167        #[cfg(feature = "mkdir")]
168        reg.register("mkdir", |args, ctx| wasibox_core::utils::mkdir::execute_with_context(args.iter().cloned(), ctx));
169        #[cfg(feature = "mv")]
170        reg.register("mv", |args, ctx| wasibox_core::utils::mv::execute_with_context(args.iter().cloned(), ctx));
171        #[cfg(feature = "pwd")]
172        reg.register("pwd", |args, ctx| wasibox_core::utils::pwd::execute_with_context(args.iter().cloned(), ctx));
173        #[cfg(feature = "rm")]
174        reg.register("rm", |args, ctx| wasibox_core::utils::rm::execute_with_context(args.iter().cloned(), ctx));
175        #[cfg(feature = "rmdir")]
176        reg.register("rmdir", |args, ctx| wasibox_core::utils::rmdir::execute_with_context(args.iter().cloned(), ctx));
177        #[cfg(feature = "seq")]
178        reg.register("seq", |args, ctx| wasibox_core::utils::seq::execute_with_context(args.iter().cloned(), ctx));
179        #[cfg(feature = "sleep")]
180        reg.register("sleep", |args, ctx| wasibox_core::utils::sleep::execute_with_context(args.iter().cloned(), ctx));
181        #[cfg(feature = "tail")]
182        reg.register("tail", |args, ctx| wasibox_core::utils::tail::execute_with_context(args.iter().cloned(), ctx));
183        #[cfg(feature = "tee")]
184        reg.register("tee", |args, ctx| wasibox_core::utils::tee::execute_with_context(args.iter().cloned(), ctx));
185        #[cfg(feature = "touch")]
186        reg.register("touch", |args, ctx| wasibox_core::utils::touch::execute_with_context(args.iter().cloned(), ctx));
187        #[cfg(feature = "tree")]
188        reg.register("tree", |args, ctx| wasibox_core::utils::tree::execute_with_context(args.iter().cloned(), ctx));
189        #[cfg(feature = "true")]
190        reg.register("true", |args, ctx| wasibox_core::utils::r#true::execute_with_context(args.iter().cloned(), ctx));
191        #[cfg(feature = "uname")]
192        reg.register("uname", |args, ctx| wasibox_core::utils::uname::execute_with_context(args.iter().cloned(), ctx));
193        #[cfg(feature = "unlink")]
194        reg.register("unlink", |args, ctx| wasibox_core::utils::unlink::execute_with_context(args.iter().cloned(), ctx));
195        #[cfg(feature = "wc")]
196        reg.register("wc", |args, ctx| wasibox_core::utils::wc::execute_with_context(args.iter().cloned(), ctx));
197        #[cfg(feature = "whoami")]
198        reg.register("whoami", |args, ctx| wasibox_core::utils::whoami::execute_with_context(args.iter().cloned(), ctx));
199        #[cfg(feature = "yes")]
200        reg.register("yes", |args, ctx| wasibox_core::utils::yes::execute_with_context(args.iter().cloned(), ctx));
201
202        reg
203    }
204
205    /// Register (or replace) a command.
206    pub fn register<F>(&mut self, name: impl Into<String>, handler: F)
207    where
208        F: Fn(&[String], &mut IoContext) -> Result<(), String> + Send + Sync + 'static,
209    {
210        self.commands.insert(name.into(), Arc::new(handler));
211    }
212
213    /// Set a fallback handler invoked when no explicit command matches.
214    pub fn set_fallback<F>(&mut self, handler: F)
215    where
216        F: Fn(&[String], &mut IoContext) -> Result<(), String> + Send + Sync + 'static,
217    {
218        self.fallback = Some(Arc::new(handler));
219    }
220
221    /// Remove the fallback handler. Unknown commands will return an error.
222    pub fn remove_fallback(&mut self) {
223        self.fallback = None;
224    }
225
226    /// Look up and execute a command.
227    pub fn execute(&self, args: &[String], ctx: &mut IoContext) -> Result<(), String> {
228        if args.is_empty() {
229            return Ok(());
230        }
231        let cmd = &args[0];
232
233        if let Some(handler) = self.commands.get(cmd.as_str()) {
234            return handler(args, ctx);
235        }
236        if let Some(ref fallback) = self.fallback {
237            return fallback(args, ctx);
238        }
239        Err(format!("command not found: {}", cmd))
240    }
241}
242
243impl Default for CommandRegistry {
244    fn default() -> Self {
245        Self::new()
246    }
247}
248
249// ---------------------------------------------------------------------------
250// Pipeline execution
251// ---------------------------------------------------------------------------
252
253/// Check if an error string represents a BrokenPipe (normal pipeline termination).
254fn is_broken_pipe(err: &str) -> bool {
255    err.contains("Broken pipe")
256        || err.contains("BrokenPipe")
257        || err.contains("broken pipe")
258}
259
260/// Execute a shell pipeline (commands separated by `|`).
261///
262/// Each stage runs in its own thread; stages are connected by in-process
263/// pipes. The last stage supports `>` / `>>` redirection.
264pub fn handle_pipeline(
265    line: &str,
266    initial_stdin: Box<dyn Read + Send + 'static>,
267    final_stdout: Box<dyn Write + Send + 'static>,
268    registry: &CommandRegistry,
269    cancel_token: wasibox_core::CancellationToken,
270) -> Result<(), String> {
271    let stages_str: Vec<&str> = line.split('|').collect();
272
273    let mut final_stdout = Some(final_stdout);
274
275    std::thread::scope(|s| {
276        let mut prev_reader: Box<dyn Read + Send> = initial_stdin;
277        let mut threads = Vec::new();
278
279        for (i, stage_str) in stages_str.iter().enumerate() {
280            let is_last = i == stages_str.len() - 1;
281            let mut tokens = shlex::split(stage_str.trim())
282                .ok_or_else(|| "Error: Invalid input (quoting error)".to_string())?;
283
284            let stdin = std::mem::replace(&mut prev_reader, Box::new(io::empty()));
285            let (stdout, next_reader): (Box<dyn Write + Send>, Option<Box<dyn Read + Send>>) = if is_last {
286                let mut out: Box<dyn Write + Send> = final_stdout.take().unwrap();
287                if let Some(pos) = tokens.iter().position(|t| t == ">" || t == ">>") {
288                    let append = tokens[pos] == ">>";
289                    let filename = tokens.get(pos + 1).ok_or("Error: Missing file for redirection")?;
290                    let file = if append {
291                        std::fs::OpenOptions::new().create(true).append(true).open(filename)
292                    } else {
293                        File::create(filename)
294                    }.map_err(|e| format!("Error opening file: {}", e))?;
295                    out = Box::new(file);
296                    tokens.truncate(pos);
297                }
298                (out, None)
299            } else {
300                let (reader, writer) = create_pipe(cancel_token.clone());
301                (Box::new(writer), Some(Box::new(reader)))
302            };
303
304            if let Some(r) = next_reader {
305                prev_reader = r;
306            }
307
308            let cancel_clone = cancel_token.clone();
309            let thread = s.spawn(move || {
310                if cancel_clone.is_cancelled() {
311                    return Err("Interrupted".to_string());
312                }
313                let wrapped_stdin = Box::new(CancelReader { inner: stdin, token: cancel_clone.clone() });
314                let wrapped_stdout = Box::new(CancelWriter { inner: stdout, token: cancel_clone.clone() });
315                let wrapped_stderr = Box::new(CancelWriter { inner: Box::new(io::stderr()), token: cancel_clone.clone() });
316
317                let mut ctx = IoContext::with_cancel(wrapped_stdin, wrapped_stdout, wrapped_stderr, cancel_clone.clone());
318                let res = registry.execute(&tokens, &mut ctx);
319                if cancel_clone.is_cancelled() {
320                    return Err("Interrupted".to_string());
321                }
322                res
323            });
324            threads.push((i, thread));
325        }
326
327        // Collect results: BrokenPipe in non-last stages is normal
328        // (upstream terminated by downstream closing the pipe).
329        let last_idx = stages_str.len() - 1;
330        let mut final_res = Ok(());
331        for (idx, thread) in threads {
332            match thread.join() {
333                Ok(Ok(())) => {}
334                Ok(Err(e)) => {
335                    if idx != last_idx && is_broken_pipe(&e) {
336                        continue;
337                    }
338                    if final_res.is_ok() {
339                        final_res = Err(e);
340                    }
341                }
342                Err(_) => {
343                    if final_res.is_ok() {
344                        final_res = Err("Thread panicked".to_string());
345                    }
346                }
347            }
348        }
349        final_res
350    })
351}
352
353/// Execute multiple independent command lines in parallel.
354///
355/// The first line receives `initial_stdin` / `final_stdout`; all subsequent
356/// lines use `io::empty()` / `io::stdout()`.
357pub fn handle_parallel(
358    lines: Vec<String>,
359    initial_stdin: Box<dyn Read + Send + 'static>,
360    final_stdout: Box<dyn Write + Send + 'static>,
361    registry: Arc<CommandRegistry>,
362    cancel_token: wasibox_core::CancellationToken,
363) -> Vec<Result<(), String>> {
364    let mut handles = Vec::new();
365    let mut stdin_opt: Option<Box<dyn Read + Send>> = Some(initial_stdin);
366    let mut stdout_opt: Option<Box<dyn Write + Send>> = Some(final_stdout);
367
368    for line in lines {
369        let registry = Arc::clone(&registry);
370        let mut stdin_for_thread: Option<Box<dyn Read + Send>> = Some(stdin_opt.take().unwrap_or_else(|| Box::new(io::empty())));
371        let mut stdout_for_thread: Option<Box<dyn Write + Send>> = Some(stdout_opt.take().unwrap_or_else(|| Box::new(io::stdout())));
372
373        let cancel_clone = cancel_token.clone();
374        let handle = std::thread::spawn(move || {
375            let commands: Vec<&str> = line.split("&&").collect();
376            for cmd in commands {
377                let cmd = cmd.trim();
378                if cmd.is_empty() { continue; }
379                if cancel_clone.is_cancelled() { return Err("Interrupted".to_string()); }
380                let stdin = stdin_for_thread.take().unwrap_or_else(|| Box::new(io::empty()));
381                let stdout = stdout_for_thread.take().unwrap_or_else(|| Box::new(io::stdout()));
382                handle_pipeline(cmd, stdin, stdout, &*registry, cancel_clone.clone())?;
383            }
384            Ok(())
385        });
386        handles.push(handle);
387    }
388
389    handles.into_iter()
390        .map(|h| h.join().unwrap_or(Err("Thread panicked".to_string())))
391        .collect()
392}
393
394/// Execute a complete command line, splitting by `&&` and running sequentially.
395/// Stops execution if any command in the sequence fails.
396pub fn handle_command_line(
397    line: &str,
398    registry: &CommandRegistry,
399    cancel_token: wasibox_core::CancellationToken,
400) -> Result<(), String> {
401    let commands: Vec<&str> = line.split("&&").collect();
402    for cmd in commands {
403        let cmd = cmd.trim();
404        if cmd.is_empty() {
405            continue;
406        }
407        if cancel_token.is_cancelled() { return Err("Interrupted".to_string()); }
408        handle_pipeline(cmd, Box::new(io::stdin()), Box::new(io::stdout()), registry, cancel_token.clone())?;
409    }
410    Ok(())
411}
412
413// ---------------------------------------------------------------------------
414// Pipe implementation
415// ---------------------------------------------------------------------------
416
417struct CancelReader<R: Read> {
418    inner: R,
419    token: wasibox_core::CancellationToken,
420}
421impl<R: Read> Read for CancelReader<R> {
422    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
423        if self.token.is_cancelled() {
424            return Err(io::Error::new(io::ErrorKind::BrokenPipe, "Interrupted"));
425        }
426        self.inner.read(buf)
427    }
428}
429
430struct CancelWriter<W: Write> {
431    inner: W,
432    token: wasibox_core::CancellationToken,
433}
434impl<W: Write> Write for CancelWriter<W> {
435    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
436        if self.token.is_cancelled() {
437            return Err(io::Error::new(io::ErrorKind::BrokenPipe, "Interrupted"));
438        }
439        self.inner.write(buf)
440    }
441    fn flush(&mut self) -> io::Result<()> {
442        if self.token.is_cancelled() {
443            return Err(io::Error::new(io::ErrorKind::BrokenPipe, "Interrupted"));
444        }
445        self.inner.flush()
446    }
447}
448
449struct PipeReader {
450    rx: std::sync::mpsc::Receiver<Vec<u8>>,
451    buffer: Vec<u8>,
452    pos: usize,
453    token: wasibox_core::CancellationToken,
454}
455impl Read for PipeReader {
456    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
457        if self.token.is_cancelled() {
458            return Err(io::Error::new(io::ErrorKind::BrokenPipe, "Interrupted"));
459        }
460        if self.pos >= self.buffer.len() {
461            loop {
462                if self.token.is_cancelled() {
463                    return Err(io::Error::new(io::ErrorKind::BrokenPipe, "Interrupted"));
464                }
465                match self.rx.recv_timeout(std::time::Duration::from_millis(100)) {
466                    Ok(data) => {
467                        self.buffer = data;
468                        self.pos = 0;
469                        break;
470                    }
471                    Err(std::sync::mpsc::RecvTimeoutError::Timeout) => continue,
472                    Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => return Ok(0),
473                }
474            }
475        }
476        let available = self.buffer.len() - self.pos;
477        let to_copy = std::cmp::min(available, buf.len());
478        buf[..to_copy].copy_from_slice(&self.buffer[self.pos..self.pos + to_copy]);
479        self.pos += to_copy;
480        Ok(to_copy)
481    }
482}
483
484struct PipeWriter {
485    tx: std::sync::mpsc::SyncSender<Vec<u8>>,
486    token: wasibox_core::CancellationToken,
487}
488impl Write for PipeWriter {
489    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
490        if self.token.is_cancelled() {
491            return Err(io::Error::new(io::ErrorKind::BrokenPipe, "Interrupted"));
492        }
493        self.tx.send(buf.to_vec()).map_err(|e| io::Error::new(io::ErrorKind::BrokenPipe, e))?;
494        Ok(buf.len())
495    }
496    fn flush(&mut self) -> io::Result<()> { Ok(()) }
497}
498
499fn create_pipe(token: wasibox_core::CancellationToken) -> (PipeReader, PipeWriter) {
500    let (tx, rx) = std::sync::mpsc::sync_channel(1024);
501    (PipeReader { rx, buffer: Vec::new(), pos: 0, token: token.clone() }, PipeWriter { tx, token })
502}
503
504// ---------------------------------------------------------------------------
505// Stdin Multiplexer (for signals on platforms without them)
506// ---------------------------------------------------------------------------
507
508use std::sync::mpsc::{self, Sender, Receiver};
509
510/// A shared stdin reader that can be monitored by a background thread
511/// to intercept control characters (like Ctrl-C) on platforms without signals.
512pub struct StdinMultiplexer {
513    tx_subscribers: Mutex<Vec<Sender<Vec<u8>>>>,
514    cancel_token: wasibox_core::CancellationToken,
515}
516
517impl StdinMultiplexer {
518    pub fn new(cancel_token: wasibox_core::CancellationToken) -> Arc<Self> {
519        let mux = Arc::new(Self {
520            tx_subscribers: Mutex::new(Vec::new()),
521            cancel_token,
522        });
523
524        let mux_clone = Arc::clone(&mux);
525        std::thread::spawn(move || {
526            let mut buf = [0u8; 1024];
527            let mut stdin = io::stdin();
528            loop {
529                match stdin.read(&mut buf) {
530                    Ok(0) => break, // EOF
531                    Ok(n) => {
532                        let data = &buf[..n];
533                        // Check for Ctrl-C (byte 3)
534                        for &b in data {
535                            if b == 3 {
536                                mux_clone.cancel_token.cancel();
537                            }
538                        }
539                        // Broadcast to all active subscribers
540                        let mut subs = mux_clone.tx_subscribers.lock().unwrap();
541                        subs.retain(|tx| {
542                            tx.send(data.to_vec()).is_ok()
543                        });
544                    }
545                    Err(_) => {
546                        // On some platforms, read error might be temporary,
547                        // but for stdin it usually means we should stop.
548                        std::thread::sleep(std::time::Duration::from_millis(100));
549                    }
550                }
551            }
552        });
553
554        mux
555    }
556
557    pub fn subscribe(&self) -> MultiplexedReader {
558        let (tx, rx) = mpsc::channel();
559        self.tx_subscribers.lock().unwrap().push(tx);
560        MultiplexedReader { rx, buffer: Vec::new(), pos: 0 }
561    }
562}
563
564pub struct MultiplexedReader {
565    rx: Receiver<Vec<u8>>,
566    buffer: Vec<u8>,
567    pos: usize,
568}
569
570impl Read for MultiplexedReader {
571    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
572        if self.pos >= self.buffer.len() {
573            match self.rx.recv() {
574                Ok(data) => {
575                    self.buffer = data;
576                    self.pos = 0;
577                }
578                Err(_) => return Ok(0),
579            }
580        }
581        let available = self.buffer.len() - self.pos;
582        let to_copy = std::cmp::min(available, buf.len());
583        buf[..to_copy].copy_from_slice(&self.buffer[self.pos..self.pos + to_copy]);
584        self.pos += to_copy;
585        Ok(to_copy)
586    }
587}
588
589/// A thread-safe writer backed by a shared `Vec<u8>`. Useful for tests.
590pub struct ArcVecWriter {
591    pub inner: Arc<Mutex<Vec<u8>>>,
592}
593impl Write for ArcVecWriter {
594    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
595        let mut inner = self.inner.lock().map_err(|e| io::Error::new(io::ErrorKind::Other, e.to_string()))?;
596        inner.extend_from_slice(buf);
597        Ok(buf.len())
598    }
599    fn flush(&mut self) -> io::Result<()> { Ok(()) }
600}
601
602// ---------------------------------------------------------------------------
603// Tests
604// ---------------------------------------------------------------------------
605
606#[cfg(test)]
607mod tests {
608    use super::*;
609    use std::io::{BufRead, BufReader, Cursor};
610
611    fn get_temp_dir() -> tempfile::TempDir {
612        #[cfg(target_os = "wasi")]
613        {
614            tempfile::Builder::new()
615                .prefix("test_")
616                .tempdir_in(".")
617                .expect("Failed to create temp dir in current directory")
618        }
619        #[cfg(not(target_os = "wasi"))]
620        {
621            tempfile::tempdir().expect("Failed to create system temp dir")
622        }
623    }
624
625    /// Shorthand: registry with all builtins (shell + wasibox-core)
626    fn builtins() -> CommandRegistry {
627        CommandRegistry::with_builtins()
628    }
629
630    // ── basic pipeline tests ────────────────────────────────────────────
631
632    #[test]
633    fn test_simple_echo() {
634        let out = Arc::new(Mutex::new(Vec::new()));
635        handle_pipeline("echo hello", Box::new(Cursor::new("")), Box::new(ArcVecWriter { inner: Arc::clone(&out) }), &builtins(), wasibox_core::CancellationToken::new()).unwrap();
636        let buf = out.lock().unwrap();
637        assert_eq!(String::from_utf8_lossy(&buf).trim(), "hello");
638    }
639
640    #[test]
641    fn test_pipe_echo_cat() {
642        let out = Arc::new(Mutex::new(Vec::new()));
643        handle_pipeline("echo hello | cat", Box::new(Cursor::new("")), Box::new(ArcVecWriter { inner: Arc::clone(&out) }), &builtins(), wasibox_core::CancellationToken::new()).unwrap();
644        let buf = out.lock().unwrap();
645        assert_eq!(String::from_utf8_lossy(&buf).trim(), "hello");
646    }
647
648    #[test]
649    fn test_pipe_grep() {
650        let out = Arc::new(Mutex::new(Vec::new()));
651        handle_pipeline("echo world | grep world", Box::new(Cursor::new("")), Box::new(ArcVecWriter { inner: Arc::clone(&out) }), &builtins(), wasibox_core::CancellationToken::new()).unwrap();
652        let buf = out.lock().unwrap();
653        assert_eq!(String::from_utf8_lossy(&buf).trim(), "world");
654    }
655
656    #[test]
657    fn test_redirection_create() {
658        let dir = get_temp_dir();
659        let file_path = dir.path().join("test_create.txt");
660        let cmd = format!("echo hello > \"{}\"", file_path.display());
661        handle_pipeline(&cmd, Box::new(Cursor::new("")), Box::new(io::sink()), &builtins(), wasibox_core::CancellationToken::new()).unwrap();
662        let content = std::fs::read_to_string(file_path).unwrap();
663        assert_eq!(content.trim(), "hello");
664    }
665
666    #[test]
667    fn test_redirection_append() {
668        let dir = get_temp_dir();
669        let file_path = dir.path().join("test_append.txt");
670        let cmd1 = format!("echo hello > \"{}\"", file_path.display());
671        handle_pipeline(&cmd1, Box::new(Cursor::new("")), Box::new(io::sink()), &builtins(), wasibox_core::CancellationToken::new()).unwrap();
672        let cmd2 = format!("echo world >> \"{}\"", file_path.display());
673        handle_pipeline(&cmd2, Box::new(Cursor::new("")), Box::new(io::sink()), &builtins(), wasibox_core::CancellationToken::new()).unwrap();
674        let content = std::fs::read_to_string(file_path).unwrap();
675        assert!(content.contains("hello"));
676        assert!(content.contains("world"));
677    }
678
679    #[test]
680    fn test_complex_pipeline() {
681        let out = Arc::new(Mutex::new(Vec::new()));
682        handle_pipeline("echo hello | grep h | wc -c", Box::new(Cursor::new("")), Box::new(ArcVecWriter { inner: Arc::clone(&out) }), &builtins(), wasibox_core::CancellationToken::new()).unwrap();
683        let buf = out.lock().unwrap();
684        assert_eq!(String::from_utf8_lossy(&buf).trim(), "6");
685    }
686
687    // ── custom command tests ────────────────────────────────────────────
688
689    #[test]
690    fn test_custom_command() {
691        let mut reg = CommandRegistry::with_builtins();
692        reg.register("magic", |_args, ctx| {
693            write!(ctx.stdout, "magic happen").map_err(|e| e.to_string())
694        });
695
696        let out = Arc::new(Mutex::new(Vec::new()));
697        handle_pipeline("magic", Box::new(io::empty()), Box::new(ArcVecWriter { inner: Arc::clone(&out) }), &reg, wasibox_core::CancellationToken::new()).unwrap();
698        let buf = out.lock().unwrap();
699        assert_eq!(String::from_utf8_lossy(&buf), "magic happen");
700    }
701
702    #[test]
703    fn test_custom_command_in_pipeline() {
704        // Custom "double" command that duplicates each input line
705        let mut reg = CommandRegistry::with_builtins();
706        reg.register("double", |_args, ctx| {
707            for line in BufReader::new(&mut ctx.stdin).lines() {
708                let line = match line {
709                    Ok(l) => l,
710                    Err(_) => break,
711                };
712                if writeln!(ctx.stdout, "{}", line).is_err() { break; }
713                if writeln!(ctx.stdout, "{}", line).is_err() { break; }
714            }
715            Ok(())
716        });
717
718        let out = Arc::new(Mutex::new(Vec::new()));
719        handle_pipeline("echo hello | double", Box::new(io::empty()), Box::new(ArcVecWriter { inner: Arc::clone(&out) }), &reg, wasibox_core::CancellationToken::new()).unwrap();
720        let buf = out.lock().unwrap();
721        assert_eq!(String::from_utf8_lossy(&buf).trim(), "hello\nhello");
722    }
723
724    #[test]
725    fn test_override_builtin() {
726        // Override "echo" with a custom version
727        let mut reg = CommandRegistry::with_builtins();
728        reg.register("echo", |args, ctx| {
729            let msg = args[1..].join(" ").to_uppercase();
730            writeln!(ctx.stdout, "{}", msg).map_err(|e| e.to_string())
731        });
732
733        let out = Arc::new(Mutex::new(Vec::new()));
734        handle_pipeline("echo hello", Box::new(io::empty()), Box::new(ArcVecWriter { inner: Arc::clone(&out) }), &reg, wasibox_core::CancellationToken::new()).unwrap();
735        let buf = out.lock().unwrap();
736        assert_eq!(String::from_utf8_lossy(&buf).trim(), "HELLO");
737    }
738
739    // ── fully external pipeline (no builtins at all) ────────────────────
740
741    #[test]
742    fn test_all_external_pipeline() {
743        let mut reg = CommandRegistry::new();
744        reg.remove_fallback(); // no wasibox-core fallback
745
746        // "count" — infinite counter
747        reg.register("count", |_args, ctx| {
748            for i in 1u64.. {
749                if writeln!(ctx.stdout, "{}", i).is_err() { break; }
750            }
751            Ok(())
752        });
753
754        // "filter2" — keep lines containing '2'
755        reg.register("filter2", |_args, ctx| {
756            for line in BufReader::new(&mut ctx.stdin).lines() {
757                let line = match line {
758                    Ok(l) => l,
759                    Err(_) => break,
760                };
761                if line.contains('2') {
762                    if writeln!(ctx.stdout, "{}", line).is_err() { break; }
763                }
764            }
765            Ok(())
766        });
767
768        // "take N" — first N lines
769        reg.register("take", |args, ctx| {
770            let n: usize = args.get(1).and_then(|s| s.parse().ok()).unwrap_or(5);
771            for (i, line) in BufReader::new(&mut ctx.stdin).lines().enumerate() {
772                if i >= n { break; }
773                let line = match line {
774                    Ok(l) => l,
775                    Err(_) => break,
776                };
777                if writeln!(ctx.stdout, "{}", line).is_err() { break; }
778            }
779            Ok(())
780        });
781
782        let out = Arc::new(Mutex::new(Vec::new()));
783        handle_pipeline(
784            "count | filter2 | take 3",
785            Box::new(io::empty()),
786            Box::new(ArcVecWriter { inner: Arc::clone(&out) }),
787            &reg,
788            wasibox_core::CancellationToken::new(),
789        ).unwrap();
790
791        let buf = out.lock().unwrap();
792        let result = String::from_utf8_lossy(&buf);
793        let lines: Vec<&str> = result.trim().lines().collect();
794        assert_eq!(lines.len(), 3);
795        for line in &lines {
796            assert!(line.contains('2'), "expected '2' in line, got: {}", line);
797        }
798        assert_eq!(lines[0], "2");
799        assert_eq!(lines[1], "12");
800        assert_eq!(lines[2], "20");
801    }
802
803    // ── parallel tests ──────────────────────────────────────────────────
804
805    #[test]
806    fn test_parallel_execution() {
807        let mut reg = CommandRegistry::with_builtins();
808        reg.register("slow", |_args, _ctx| {
809            std::thread::sleep(std::time::Duration::from_millis(100));
810            Ok(())
811        });
812        let registry = Arc::new(reg);
813
814        let lines = vec!["slow".to_string(), "slow".to_string(), "slow".to_string()];
815        let start = std::time::Instant::now();
816        let results = handle_parallel(lines, Box::new(io::empty()), Box::new(io::sink()), registry, wasibox_core::CancellationToken::new());
817        let duration = start.elapsed();
818
819        assert_eq!(results.len(), 3);
820        for res in results {
821            assert!(res.is_ok());
822        }
823        assert!(duration < std::time::Duration::from_millis(250));
824    }
825
826    // ── streaming / infinite pipeline tests ─────────────────────────────
827
828    #[test]
829    fn test_streaming_pipeline() {
830        let out = Arc::new(Mutex::new(Vec::new()));
831        handle_pipeline("yes | head -n 2", Box::new(io::empty()), Box::new(ArcVecWriter { inner: Arc::clone(&out) }), &builtins(), wasibox_core::CancellationToken::new()).unwrap();
832        let buf = out.lock().unwrap();
833        assert_eq!(String::from_utf8_lossy(&buf).trim(), "y\ny");
834    }
835
836    #[test]
837    fn test_seq_pipeline_head() {
838        let out = Arc::new(Mutex::new(Vec::new()));
839        handle_pipeline("seq | head -n 5", Box::new(io::empty()), Box::new(ArcVecWriter { inner: Arc::clone(&out) }), &builtins(), wasibox_core::CancellationToken::new()).unwrap();
840        let buf = out.lock().unwrap();
841        assert_eq!(String::from_utf8_lossy(&buf).trim(), "1\n2\n3\n4\n5");
842    }
843
844    #[test]
845    fn test_seq_pipeline_grep_head() {
846        let out = Arc::new(Mutex::new(Vec::new()));
847        handle_pipeline("seq | grep 2 | head -n 3", Box::new(io::empty()), Box::new(ArcVecWriter { inner: Arc::clone(&out) }), &builtins(), wasibox_core::CancellationToken::new()).unwrap();
848        let buf = out.lock().unwrap();
849        let result = String::from_utf8_lossy(&buf);
850        let lines: Vec<&str> = result.trim().lines().collect();
851        assert_eq!(lines.len(), 3);
852        for line in &lines {
853            assert!(line.contains('2'));
854        }
855        assert_eq!(lines[0], "2");
856        assert_eq!(lines[1], "12");
857        assert_eq!(lines[2], "20");
858    }
859
860    #[test]
861    fn test_seq_pipeline_grep_head_5() {
862        let out = Arc::new(Mutex::new(Vec::new()));
863        handle_pipeline("seq | grep 2 | head -n 5", Box::new(io::empty()), Box::new(ArcVecWriter { inner: Arc::clone(&out) }), &builtins(), wasibox_core::CancellationToken::new()).unwrap();
864        let buf = out.lock().unwrap();
865        let result = String::from_utf8_lossy(&buf);
866        let lines: Vec<&str> = result.trim().lines().collect();
867        assert_eq!(lines.len(), 5);
868        assert_eq!(lines[0], "2");
869        assert_eq!(lines[1], "12");
870        assert_eq!(lines[2], "20");
871        assert_eq!(lines[3], "21");
872        assert_eq!(lines[4], "22");
873    }
874
875    #[test]
876    fn test_seq_finite() {
877        let out = Arc::new(Mutex::new(Vec::new()));
878        handle_pipeline("seq 3", Box::new(io::empty()), Box::new(ArcVecWriter { inner: Arc::clone(&out) }), &builtins(), wasibox_core::CancellationToken::new()).unwrap();
879        let buf = out.lock().unwrap();
880        assert_eq!(String::from_utf8_lossy(&buf).trim(), "1\n2\n3");
881    }
882
883    #[test]
884    fn test_seq_range() {
885        let out = Arc::new(Mutex::new(Vec::new()));
886        handle_pipeline("seq 5 8", Box::new(io::empty()), Box::new(ArcVecWriter { inner: Arc::clone(&out) }), &builtins(), wasibox_core::CancellationToken::new()).unwrap();
887        let buf = out.lock().unwrap();
888        assert_eq!(String::from_utf8_lossy(&buf).trim(), "5\n6\n7\n8");
889    }
890
891    #[test]
892    fn test_seq_step() {
893        let out = Arc::new(Mutex::new(Vec::new()));
894        handle_pipeline("seq 1 2 10", Box::new(io::empty()), Box::new(ArcVecWriter { inner: Arc::clone(&out) }), &builtins(), wasibox_core::CancellationToken::new()).unwrap();
895        let buf = out.lock().unwrap();
896        assert_eq!(String::from_utf8_lossy(&buf).trim(), "1\n3\n5\n7\n9");
897    }
898    #[test]
899    fn test_cd_and_ls() {
900        let dir = get_temp_dir();
901        // Create files inside the temp dir
902        std::fs::write(dir.path().join("aaa.txt"), "hello").unwrap();
903        std::fs::write(dir.path().join("bbb.txt"), "world").unwrap();
904
905        let original_cwd = env::current_dir().unwrap();
906        let reg = builtins();
907
908        // cd into the temp directory
909        let cd_cmd = format!("cd \"{}\"", dir.path().display());
910        handle_pipeline(&cd_cmd, Box::new(io::empty()), Box::new(io::sink()), &reg, wasibox_core::CancellationToken::new()).unwrap();
911
912        // ls the current directory (should now be the temp dir)
913        let out = Arc::new(Mutex::new(Vec::new()));
914        handle_pipeline("ls", Box::new(io::empty()), Box::new(ArcVecWriter { inner: Arc::clone(&out) }), &reg, wasibox_core::CancellationToken::new()).unwrap();
915
916        let buf = out.lock().unwrap();
917        let output = String::from_utf8_lossy(&buf);
918        assert!(output.contains("aaa.txt"), "expected aaa.txt in ls output, got: {}", output);
919        assert!(output.contains("bbb.txt"), "expected bbb.txt in ls output, got: {}", output);
920
921        // Restore original cwd
922        env::set_current_dir(original_cwd).unwrap();
923    }
924
925    // ── sequence (&&) tests ─────────────────────────────────────────────
926
927    #[test]
928    fn test_command_line_sequence() {
929        let mut reg = CommandRegistry::with_builtins();
930        let counter = Arc::new(Mutex::new(0));
931        let c1 = Arc::clone(&counter);
932        reg.register("inc", move |_args, _ctx| {
933            *c1.lock().unwrap() += 1;
934            Ok(())
935        });
936
937        // inc && inc && inc
938        crate::handle_command_line("inc && inc && inc", &reg, wasibox_core::CancellationToken::new()).unwrap();
939        assert_eq!(*counter.lock().unwrap(), 3);
940    }
941
942    #[test]
943    fn test_command_line_short_circuit() {
944        let mut reg = CommandRegistry::with_builtins();
945        let counter = Arc::new(Mutex::new(0));
946        let c1 = Arc::clone(&counter);
947        reg.register("inc", move |_args, _ctx| {
948            *c1.lock().unwrap() += 1;
949            Ok(())
950        });
951        reg.register("fail", |_args, _ctx| Err("simulated failure".to_string()));
952
953        // inc && fail && inc
954        // Should stop after "fail"
955        let res = crate::handle_command_line("inc && fail && inc", &reg, wasibox_core::CancellationToken::new());
956        assert!(res.is_err());
957        assert_eq!(res.unwrap_err(), "simulated failure");
958        assert_eq!(*counter.lock().unwrap(), 1); // Only the first "inc" should execute
959    }
960
961    #[test]
962    fn test_handle_parallel_cancel() {
963        println!("STARTING TEST");
964        let registry = Arc::new(builtins());
965        let cancel_token = wasibox_core::CancellationToken::new();
966        let cancel_clone = cancel_token.clone();
967
968        let thread = std::thread::spawn(move || {
969            println!("THREAD SPAWNED");
970            super::handle_parallel(
971                vec!["yes".to_string()],
972                Box::new(std::io::empty()),
973                Box::new(std::io::sink()),
974                registry,
975                cancel_clone,
976            )
977        });
978
979        std::thread::sleep(std::time::Duration::from_millis(10));
980        println!("CANCELLING TOKEN");
981        cancel_token.cancel();
982
983        println!("JOINING THREAD");
984        let results = thread.join().unwrap();
985        println!("RESULTS: {:?}", results);
986        assert_eq!(results.len(), 1);
987        assert!(results[0].is_err());
988        assert_eq!(results[0].as_ref().unwrap_err(), "Interrupted");
989    }
990
991    #[test]
992    fn test_pipeline_cancel() {
993        let registry = builtins();
994        let cancel_token = wasibox_core::CancellationToken::new();
995        let cancel_clone = cancel_token.clone();
996
997        let thread = std::thread::spawn(move || {
998            super::handle_pipeline(
999                "yes | grep y",
1000                Box::new(std::io::empty()),
1001                Box::new(std::io::sink()),
1002                &registry,
1003                cancel_clone,
1004            )
1005        });
1006
1007        std::thread::sleep(std::time::Duration::from_millis(10));
1008        cancel_token.cancel();
1009
1010        let result = thread.join().unwrap();
1011        assert!(result.is_err());
1012        assert_eq!(result.unwrap_err(), "Interrupted");
1013    }
1014
1015    #[test]
1016    fn test_seq_redirection_parallel_cancel() {
1017        let dir = get_temp_dir();
1018        let file_path = dir.path().join("i.txt");
1019        let cmd = format!("seq > \"{}\"", file_path.display());
1020        
1021        let registry = Arc::new(builtins());
1022        let cancel_token = wasibox_core::CancellationToken::new();
1023        let cancel_clone = cancel_token.clone();
1024
1025        let thread = std::thread::spawn(move || {
1026            super::handle_parallel(
1027                vec![cmd],
1028                Box::new(std::io::empty()),
1029                Box::new(std::io::sink()),
1030                registry,
1031                cancel_clone,
1032            )
1033        });
1034
1035        // Give it some time to start writing
1036        std::thread::sleep(std::time::Duration::from_millis(50));
1037        cancel_token.cancel();
1038
1039        let results = thread.join().unwrap();
1040        assert_eq!(results.len(), 1);
1041        assert!(results[0].is_err());
1042        assert_eq!(results[0].as_ref().unwrap_err(), "Interrupted");
1043        
1044        // Verify that the file was created and contains some data
1045        assert!(file_path.exists());
1046        let content = std::fs::read_to_string(&file_path).unwrap();
1047        assert!(!content.is_empty(), "File should not be empty");
1048        
1049        // Ensure it stopped
1050        println!("Final sequence number written: {}", content.lines().last().unwrap_or("none"));
1051    }
1052}
1053