use std::future::Future;
use std::io::Write;
use std::path::{Path, PathBuf};
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use tokio::sync::mpsc::UnboundedSender;
use tokio_util::sync::CancellationToken;
use crate::ast::{AndOrList, AndOrOp, Command, Pipeline, Program, Redirect};
use crate::builtins::{self, Builtin};
use crate::expand::{self, ExpandCtx};
use crate::io::{CommandIo, InTarget, OutTarget};
use crate::spawn;
use crate::state::{ShellState, Var};
pub struct ShellEnv {
pub env: Vec<(String, String)>,
pub cwd: PathBuf,
pub output: UnboundedSender<ShellOutputLine>,
pub cancel: CancellationToken,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ShellOutputLine {
pub stream: ShellStream,
pub text: String,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ShellStream {
Stdout,
Stderr,
}
pub struct ShellResult {
pub exit_code: i32,
}
pub(crate) enum Flow {
Next(i32),
Exit(i32),
}
pub async fn execute(program: &Program, env: ShellEnv) -> ShellResult {
let mut state = ShellState::new(env.env, env.cwd);
let stdout = OutTarget::Lines {
tx: env.output.clone(),
stream: ShellStream::Stdout,
};
let flow = exec_program(program, &mut state, &stdout, &env.output, &env.cancel).await;
let exit_code = match flow {
Flow::Next(code) | Flow::Exit(code) => code,
};
ShellResult { exit_code }
}
pub(crate) fn execute_captured<'a>(
program: &'a Program,
state: &'a mut ShellState,
ctx: &'a ExpandCtx<'a>,
) -> Pin<Box<dyn Future<Output = (i32, String)> + Send + 'a>> {
Box::pin(async move {
let buffer: Arc<Mutex<Vec<u8>>> = Arc::new(Mutex::new(Vec::new()));
let stdout = OutTarget::Capture(buffer.clone());
let flow = exec_program(program, state, &stdout, ctx.output, ctx.cancel).await;
let code = match flow {
Flow::Next(code) | Flow::Exit(code) => code,
};
let bytes = buffer.lock().expect("capture buffer poisoned").clone();
let text = String::from_utf8_lossy(&bytes).replace("\r\n", "\n");
(code, text)
})
}
async fn exec_program(
program: &Program,
state: &mut ShellState,
stdout: &OutTarget,
output: &UnboundedSender<ShellOutputLine>,
cancel: &CancellationToken,
) -> Flow {
let mut code = 0;
for item in &program.items {
if cancel.is_cancelled() {
return Flow::Exit(130);
}
match exec_and_or_list(item, state, stdout, output, cancel).await {
Flow::Exit(code) => return Flow::Exit(code),
Flow::Next(next) => code = next,
}
}
Flow::Next(code)
}
async fn exec_and_or_list(
list: &AndOrList,
state: &mut ShellState,
stdout: &OutTarget,
output: &UnboundedSender<ShellOutputLine>,
cancel: &CancellationToken,
) -> Flow {
let mut code = match exec_pipeline(&list.first, state, stdout, output, cancel).await {
Flow::Exit(code) => return Flow::Exit(code),
Flow::Next(code) => code,
};
for (op, pipeline) in &list.rest {
let should_run = match op {
AndOrOp::And => code == 0,
AndOrOp::Or => code != 0,
};
if !should_run {
continue;
}
code = match exec_pipeline(pipeline, state, stdout, output, cancel).await {
Flow::Exit(code) => return Flow::Exit(code),
Flow::Next(code) => code,
};
}
Flow::Next(code)
}
async fn exec_pipeline(
pipeline: &Pipeline,
state: &mut ShellState,
stdout: &OutTarget,
output: &UnboundedSender<ShellOutputLine>,
cancel: &CancellationToken,
) -> Flow {
if cancel.is_cancelled() {
return Flow::Exit(130);
}
let flow = if pipeline.commands.len() == 1 {
match CommandIo::defaults(stdout, output) {
Ok(io) => exec_command(&pipeline.commands[0], state, io, output, cancel).await,
Err(error) => io_setup_failure(&error, output),
}
} else {
exec_stages(pipeline, state, stdout, output, cancel).await
};
match flow {
Flow::Exit(code) => Flow::Exit(code),
Flow::Next(code) => {
let code = if pipeline.negated { negate(code) } else { code };
state.last_exit = code;
Flow::Next(code)
}
}
}
async fn exec_stages(
pipeline: &Pipeline,
state: &ShellState,
stdout: &OutTarget,
output: &UnboundedSender<ShellOutputLine>,
cancel: &CancellationToken,
) -> Flow {
let ios = match build_stage_io(pipeline.commands.len(), stdout, output) {
Ok(ios) => ios,
Err(error) => return io_setup_failure(&error, output),
};
let mut stages = Vec::with_capacity(ios.len());
for (command, io) in pipeline.commands.iter().zip(ios) {
let command = command.clone();
let mut state = state.clone();
let output = output.clone();
let cancel = cancel.clone();
stages.push(tokio::spawn(async move {
exec_command(&command, &mut state, io, &output, &cancel).await
}));
}
let mut code = 0;
for stage in stages {
let flow = tokio::select! {
joined = stage => joined.expect("pipeline stage panicked"),
() = cancel.cancelled() => return Flow::Exit(130),
};
code = match flow {
Flow::Next(stage_code) | Flow::Exit(stage_code) => stage_code,
};
}
if cancel.is_cancelled() {
return Flow::Exit(130);
}
Flow::Next(code)
}
fn build_stage_io(
stages: usize,
stdout: &OutTarget,
output: &UnboundedSender<ShellOutputLine>,
) -> std::io::Result<Vec<CommandIo>> {
let mut ios = Vec::with_capacity(stages);
let mut upstream: Option<os_pipe::PipeReader> = None;
for stage in 0..stages {
let mut io = CommandIo::defaults(stdout, output)?;
if let Some(reader) = upstream.take() {
io.stdin = InTarget::Pipe(reader);
}
if stage + 1 < stages {
let (reader, writer) = os_pipe::pipe()?;
upstream = Some(reader);
io.stdout = OutTarget::Pipe(writer);
}
ios.push(io);
}
Ok(ios)
}
fn io_setup_failure(error: &std::io::Error, output: &UnboundedSender<ShellOutputLine>) -> Flow {
let _ = output.send(ShellOutputLine {
stream: ShellStream::Stderr,
text: format!("alba-shell: cannot set up command io: {error}"),
});
Flow::Next(1)
}
fn negate(code: i32) -> i32 {
if code == 0 { 1 } else { 0 }
}
async fn exec_command(
command: &Command,
state: &mut ShellState,
mut io: CommandIo,
output: &UnboundedSender<ShellOutputLine>,
cancel: &CancellationToken,
) -> Flow {
let ctx = ExpandCtx { output, cancel };
let mut assignments = Vec::with_capacity(command.assignments.len());
for assignment in &command.assignments {
let value = match expand::expand_word_single(&assignment.value, state, &ctx).await {
Ok(value) => value,
Err(flow) => return flow,
};
assignments.push((assignment.name.clone(), value));
}
let words = match expand::expand_words(&command.words, state, &ctx).await {
Ok(words) => words,
Err(flow) => return flow,
};
let redirects = match resolve_redirects(&command.redirects, state, &ctx).await {
Ok(redirects) => redirects,
Err(flow) => return flow,
};
if let Err(message) = apply_redirects(&redirects, &state.cwd, &mut io) {
let mut stderr = io.stderr.writer();
let _ = writeln!(stderr, "{message}");
return Flow::Next(1);
}
if words.is_empty() {
for (name, value) in assignments {
state.set(name, value);
}
return Flow::Next(0);
}
let overlay = AssignmentOverlay::apply(state, &assignments);
let name = words[0].as_str();
let args = &words[1..];
let flow = if !has_path_separator(name)
&& let Some(builtin) = builtins::find(name)
{
run_builtin(builtin, args, state, io, cancel).await
} else {
run_external_command(name, args, state, io, cancel).await
};
overlay.restore(state);
flow
}
async fn run_builtin(
builtin: Builtin,
args: &[String],
state: &mut ShellState,
io: CommandIo,
cancel: &CancellationToken,
) -> Flow {
if builtin == Builtin::Sleep {
return builtins::run_sleep(args, io, cancel).await;
}
if !io.can_block() {
return builtins::run(builtin, args, state, io);
}
let args = args.to_vec();
let mut owned = state.clone();
let (tx, rx) = tokio::sync::oneshot::channel();
std::thread::spawn(move || {
let flow = builtins::run(builtin, &args, &mut owned, io);
let _ = tx.send((flow, owned));
});
let (flow, owned) = rx.await.expect("builtin panicked");
*state = owned;
flow
}
async fn run_external_command(
name: &str,
args: &[String],
state: &ShellState,
io: CommandIo,
cancel: &CancellationToken,
) -> Flow {
let resolved = if has_path_separator(name) {
let candidate = state.cwd.join(name);
candidate.exists().then_some(candidate)
} else {
which::which_in(name, state.path(), &state.cwd).ok()
};
let Some(path) = resolved else {
let mut stderr = io.stderr.writer();
let _ = writeln!(stderr, "{}", not_found_message(name, state.path()));
return Flow::Next(127);
};
let env = state.exported_env();
let code = spawn::run_external(&path, args, &env, &state.cwd, io, cancel).await;
if cancel.is_cancelled() {
Flow::Exit(130)
} else {
Flow::Next(code)
}
}
enum ResolvedRedirect {
Out {
stderr: bool,
append: bool,
target: String,
},
In {
target: String,
},
StderrToStdout,
}
async fn resolve_redirects(
redirects: &[Redirect],
state: &mut ShellState,
ctx: &ExpandCtx<'_>,
) -> Result<Vec<ResolvedRedirect>, Flow> {
let mut resolved = Vec::with_capacity(redirects.len());
for redirect in redirects {
resolved.push(match redirect {
Redirect::Out {
stderr,
append,
target,
} => ResolvedRedirect::Out {
stderr: *stderr,
append: *append,
target: expand::expand_word_single(target, state, ctx).await?,
},
Redirect::In { target } => ResolvedRedirect::In {
target: expand::expand_word_single(target, state, ctx).await?,
},
Redirect::StderrToStdout => ResolvedRedirect::StderrToStdout,
});
}
Ok(resolved)
}
fn apply_redirects(
redirects: &[ResolvedRedirect],
cwd: &Path,
io: &mut CommandIo,
) -> Result<(), String> {
for redirect in redirects {
match redirect {
ResolvedRedirect::Out {
stderr,
append,
target,
} => {
let file = open_for_write(&cwd.join(target), *append)
.map_err(|error| cannot_open(target, &error))?;
if *stderr {
io.stderr = OutTarget::File(file);
} else {
io.stdout = OutTarget::File(file);
}
}
ResolvedRedirect::In { target } => {
let file = std::fs::File::open(cwd.join(target))
.map_err(|error| cannot_open(target, &error))?;
io.stdin = InTarget::File(file);
}
ResolvedRedirect::StderrToStdout => {
io.stderr = io
.stdout
.try_clone()
.map_err(|error| cannot_open("&1", &error))?;
}
}
}
Ok(())
}
fn open_for_write(path: &Path, append: bool) -> std::io::Result<std::fs::File> {
std::fs::OpenOptions::new()
.write(true)
.create(true)
.append(append)
.truncate(!append)
.open(path)
}
fn cannot_open(target: &str, error: &std::io::Error) -> String {
format!("alba-shell: cannot open {target}: {error}")
}
fn not_found_message(name: &str, path: Option<&str>) -> String {
match did_you_mean(name, path) {
Some(candidate) => {
format!("alba-shell: command not found: {name} (did you mean `{candidate}`?)")
}
None => format!("alba-shell: command not found: {name}"),
}
}
fn did_you_mean(name: &str, path: Option<&str>) -> Option<String> {
let builtins = builtins::NAMES.iter().map(|name| (*name).to_string());
builtins
.chain(path_command_names(path))
.map(|candidate| {
let distance = distance(name, &candidate);
(candidate, distance)
})
.filter(|(candidate, distance)| is_plausible_typo(name, candidate, *distance))
.min_by(|left, right| left.1.cmp(&right.1).then_with(|| left.0.cmp(&right.0)))
.map(|(candidate, _)| candidate)
}
fn is_plausible_typo(name: &str, candidate: &str, distance: usize) -> bool {
match distance {
1 => true,
2 => name.chars().count() >= 6 && candidate.chars().count() >= 6,
_ => false,
}
}
fn path_command_names(path: Option<&str>) -> impl Iterator<Item = String> {
let dirs: Vec<PathBuf> = path
.map(|path| std::env::split_paths(path).collect())
.unwrap_or_default();
dirs.into_iter()
.filter_map(|dir| std::fs::read_dir(dir).ok())
.flatten()
.filter_map(|entry| command_name(&entry.ok()?.file_name()))
}
fn command_name(file_name: &std::ffi::OsStr) -> Option<String> {
let name = file_name.to_str()?;
if cfg!(windows) {
return Some(Path::new(name).file_stem()?.to_string_lossy().into_owned());
}
Some(name.to_string())
}
fn distance(a: &str, b: &str) -> usize {
let a: Vec<char> = a.chars().collect();
let b: Vec<char> = b.chars().collect();
let mut dp = vec![vec![0usize; b.len() + 1]; a.len() + 1];
for (i, row) in dp.iter_mut().enumerate() {
row[0] = i;
}
if let Some(first_row) = dp.first_mut() {
for (j, cell) in first_row.iter_mut().enumerate() {
*cell = j;
}
}
for i in 1..=a.len() {
for j in 1..=b.len() {
let cost = if a[i - 1] == b[j - 1] { 0 } else { 1 };
dp[i][j] = (dp[i - 1][j] + 1)
.min(dp[i][j - 1] + 1)
.min(dp[i - 1][j - 1] + cost);
if i > 1 && j > 1 && a[i - 1] == b[j - 2] && a[i - 2] == b[j - 1] {
dp[i][j] = dp[i][j].min(dp[i - 2][j - 2] + 1);
}
}
}
dp[a.len()][b.len()]
}
fn has_path_separator(name: &str) -> bool {
name.contains('/') || name.contains(std::path::MAIN_SEPARATOR)
}
struct AssignmentOverlay {
saved: Vec<(String, Option<Var>)>,
}
impl AssignmentOverlay {
fn apply(state: &mut ShellState, assignments: &[(String, String)]) -> Self {
let mut saved = Vec::with_capacity(assignments.len());
for (name, value) in assignments {
saved.push((name.clone(), state.vars.get(name).cloned()));
state.export(name, Some(value.clone()));
}
Self { saved }
}
fn restore(self, state: &mut ShellState) {
for (name, prior) in self.saved.into_iter().rev() {
match prior {
Some(var) => {
state.vars.insert(name, var);
}
None => {
state.vars.remove(&name);
}
}
}
}
}