From 42b41200551502969a99d29d70c3453746b6c538 Mon Sep 17 00:00:00 2001 From: pagedmov Date: Fri, 6 Mar 2026 11:07:46 -0500 Subject: [PATCH] Add ulimit builtin and optimize `shed -c` to exec single commands directly without forking --- Cargo.toml | 15 +++- src/builtin/mod.rs | 5 +- src/builtin/resource.rs | 169 ++++++++++++++++++++++++++++++++++++++++ src/main.rs | 4 +- src/parse/execute.rs | 134 +++++++++++++++++++------------ src/signal.rs | 8 +- 6 files changed, 279 insertions(+), 56 deletions(-) create mode 100644 src/builtin/resource.rs diff --git a/Cargo.toml b/Cargo.toml index 3ff04c8..dbf5717 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -17,7 +17,20 @@ env_logger = "0.11.9" glob = "0.3.2" itertools = "0.14.0" log = "0.4.29" -nix = { version = "0.29.0", features = ["uio", "term", "user", "hostname", "fs", "default", "signal", "process", "event", "ioctl", "poll"] } +nix = { version = "0.29.0", features = [ + "uio", + "term", + "user", + "resource", + "hostname", + "fs", + "default", + "signal", + "process", + "event", + "ioctl", + "poll" + ] } rand = "0.10.0" regex = "1.11.1" scopeguard = "1.2.0" diff --git a/src/builtin/mod.rs b/src/builtin/mod.rs index 54ee5cf..3a8b5d5 100644 --- a/src/builtin/mod.rs +++ b/src/builtin/mod.rs @@ -24,13 +24,14 @@ pub mod test; // [[ ]] thing pub mod trap; pub mod varcmds; pub mod zoltraak; +pub mod resource; -pub const BUILTINS: [&str; 47] = [ +pub const BUILTINS: [&str; 49] = [ "echo", "cd", "read", "export", "local", "pwd", "source", "shift", "jobs", "fg", "bg", "disown", "alias", "unalias", "return", "break", "continue", "exit", "zoltraak", "shopt", "builtin", "command", "trap", "pushd", "popd", "dirs", "exec", "eval", "true", "false", ":", "readonly", "unset", "complete", "compgen", "map", "pop", "fpop", "push", "fpush", "rotate", "wait", "type", - "getopts", "keymap", "read_key", "autocmd", + "getopts", "keymap", "read_key", "autocmd", "ulimit", "umask" ]; pub fn true_builtin() -> ShResult<()> { diff --git a/src/builtin/resource.rs b/src/builtin/resource.rs new file mode 100644 index 0000000..ca43f5c --- /dev/null +++ b/src/builtin/resource.rs @@ -0,0 +1,169 @@ +use ariadne::Fmt; +use nix::sys::resource::{Resource, getrlimit, setrlimit}; +use yansi::Color; + +use crate::{ + getopt::{Opt, OptSpec, get_opts_from_tokens}, libsh::error::{ShErr, ShErrKind, ShResult, ShResultExt, next_color}, parse::{NdRule, Node, execute::prepare_argv}, prelude::*, state::{self} +}; + +fn ulimit_opt_spec() -> [OptSpec;5] { + [ + OptSpec { + opt: Opt::Short('n'), // file descriptors + takes_arg: true, + }, + OptSpec { + opt: Opt::Short('u'), // max user processes + takes_arg: true, + }, + OptSpec { + opt: Opt::Short('s'), // stack size + takes_arg: true, + }, + OptSpec { + opt: Opt::Short('c'), // core dump file size + takes_arg: true, + }, + OptSpec { + opt: Opt::Short('v'), // virtual memory + takes_arg: true, + } + ] +} + +struct UlimitOpts { + fds: Option, + procs: Option, + stack: Option, + core: Option, + vmem: Option, +} + +fn get_ulimit_opts(opt: &[Opt]) -> ShResult { + let mut opts = UlimitOpts { + fds: None, + procs: None, + stack: None, + core: None, + vmem: None, + }; + + for o in opt { + match o { + Opt::ShortWithArg('n', arg) => { + opts.fds = Some(arg.parse().map_err(|_| ShErr::simple( + ShErrKind::ParseErr, + format!("invalid argument for -n: {}", arg.fg(next_color())), + ))?); + }, + Opt::ShortWithArg('u', arg) => { + opts.procs = Some(arg.parse().map_err(|_| ShErr::simple( + ShErrKind::ParseErr, + format!("invalid argument for -u: {}", arg.fg(next_color())), + ))?); + }, + Opt::ShortWithArg('s', arg) => { + opts.stack = Some(arg.parse().map_err(|_| ShErr::simple( + ShErrKind::ParseErr, + format!("invalid argument for -s: {}", arg.fg(next_color())), + ))?); + }, + Opt::ShortWithArg('c', arg) => { + opts.core = Some(arg.parse().map_err(|_| ShErr::simple( + ShErrKind::ParseErr, + format!("invalid argument for -c: {}", arg.fg(next_color())), + ))?); + }, + Opt::ShortWithArg('v', arg) => { + opts.vmem = Some(arg.parse().map_err(|_| ShErr::simple( + ShErrKind::ParseErr, + format!("invalid argument for -v: {}", arg.fg(next_color())), + ))?); + }, + o => return Err(ShErr::simple( + ShErrKind::ParseErr, + format!("invalid option: {}", o.fg(next_color())), + )), + } + } + + Ok(opts) +} + +pub fn ulimit(node: Node) -> ShResult<()> { + let span = node.get_span(); + let NdRule::Command { + assignments: _, + argv, + } = node.class + else { + unreachable!() + }; + + let (_, opts) = get_opts_from_tokens(argv, &ulimit_opt_spec()).promote_err(span.clone())?; + let ulimit_opts = get_ulimit_opts(&opts).promote_err(span.clone())?; + + if let Some(fds) = ulimit_opts.fds { + let (_, hard) = getrlimit(Resource::RLIMIT_NOFILE).map_err(|e| ShErr::at( + ShErrKind::ExecFail, + span.clone(), + format!("failed to get file descriptor limit: {}", e), + ))?; + setrlimit(Resource::RLIMIT_NOFILE, fds, hard).map_err(|e| ShErr::at( + ShErrKind::ExecFail, + span.clone(), + format!("failed to set file descriptor limit: {}", e), + ))?; + } + if let Some(procs) = ulimit_opts.procs { + let (_, hard) = getrlimit(Resource::RLIMIT_NPROC).map_err(|e| ShErr::at( + ShErrKind::ExecFail, + span.clone(), + format!("failed to get process limit: {}", e), + ))?; + setrlimit(Resource::RLIMIT_NPROC, procs, hard).map_err(|e| ShErr::at( + ShErrKind::ExecFail, + span.clone(), + format!("failed to set process limit: {}", e), + ))?; + } + if let Some(stack) = ulimit_opts.stack { + let (_, hard) = getrlimit(Resource::RLIMIT_STACK).map_err(|e| ShErr::at( + ShErrKind::ExecFail, + span.clone(), + format!("failed to get stack size limit: {}", e), + ))?; + setrlimit(Resource::RLIMIT_STACK, stack, hard).map_err(|e| ShErr::at( + ShErrKind::ExecFail, + span.clone(), + format!("failed to set stack size limit: {}", e), + ))?; + } + if let Some(core) = ulimit_opts.core { + let (_, hard) = getrlimit(Resource::RLIMIT_CORE).map_err(|e| ShErr::at( + ShErrKind::ExecFail, + span.clone(), + format!("failed to get core dump size limit: {}", e), + ))?; + setrlimit(Resource::RLIMIT_CORE, core, hard).map_err(|e| ShErr::at( + ShErrKind::ExecFail, + span.clone(), + format!("failed to set core dump size limit: {}", e), + ))?; + } + if let Some(vmem) = ulimit_opts.vmem { + let (_, hard) = getrlimit(Resource::RLIMIT_AS).map_err(|e| ShErr::at( + ShErrKind::ExecFail, + span.clone(), + format!("failed to get virtual memory limit: {}", e), + ))?; + setrlimit(Resource::RLIMIT_AS, vmem, hard).map_err(|e| ShErr::at( + ShErrKind::ExecFail, + span.clone(), + format!("failed to set virtual memory limit: {}", e), + ))?; + } + + state::set_status(0); + Ok(()) +} diff --git a/src/main.rs b/src/main.rs index 55f0c60..6e8a61f 100644 --- a/src/main.rs +++ b/src/main.rs @@ -30,7 +30,7 @@ use crate::builtin::trap::TrapTarget; use crate::libsh::error::{self, ShErr, ShErrKind, ShResult}; use crate::libsh::sys::TTY_FILENO; use crate::libsh::utils::AutoCmdVecUtils; -use crate::parse::execute::exec_input; +use crate::parse::execute::{exec_dash_c, exec_input}; use crate::prelude::*; use crate::procio::borrow_fd; use crate::readline::term::{LineWriter, RawModeGuard, raw_mode}; @@ -130,7 +130,7 @@ fn main() -> ExitCode { if let Err(e) = if let Some(path) = args.script { run_script(path, args.script_args) } else if let Some(cmd) = args.command { - exec_input(cmd, None, false, None) + exec_dash_c(cmd) } else { let res = shed_interactive(args); write(borrow_fd(*TTY_FILENO), b"\x1b[?2004l").ok(); // disable bracketed paste mode on exit diff --git a/src/parse/execute.rs b/src/parse/execute.rs index ed8dd31..81d4566 100644 --- a/src/parse/execute.rs +++ b/src/parse/execute.rs @@ -5,31 +5,11 @@ use std::{ }; use ariadne::Fmt; +use nix::sys::resource; use crate::{ builtin::{ - alias::{alias, unalias}, - arrops::{arr_fpop, arr_fpush, arr_pop, arr_push, arr_rotate}, - autocmd::autocmd, - cd::cd, - complete::{compgen_builtin, complete_builtin}, - dirstack::{dirs, popd, pushd}, - echo::echo, - eval, exec, - flowctl::flowctl, - getopts::getopts, - intro, - jobctl::{self, JobBehavior, continue_job, disown, jobs}, - keymap, map, - pwd::pwd, - read::{self, read_builtin}, - shift::shift, - shopt::shopt, - source::source, - test::double_bracket_test, - trap::{TrapTarget, trap}, - varcmds::{export, local, readonly, unset}, - zoltraak::zoltraak, + alias::{alias, unalias}, arrops::{arr_fpop, arr_fpush, arr_pop, arr_push, arr_rotate}, autocmd::autocmd, cd::cd, complete::{compgen_builtin, complete_builtin}, dirstack::{dirs, popd, pushd}, echo::echo, eval, exec, flowctl::flowctl, getopts::getopts, intro, jobctl::{self, JobBehavior, continue_job, disown, jobs}, keymap, map, pwd::pwd, read::{self, read_builtin}, resource::ulimit, shift::shift, shopt::shopt, source::source, test::double_bracket_test, trap::{TrapTarget, trap}, varcmds::{export, local, readonly, unset}, zoltraak::zoltraak }, expand::{expand_aliases, expand_case_pattern, glob_to_regex}, jobs::{ChildProc, JobStack, attach_tty, dispatch_job}, @@ -132,6 +112,64 @@ impl ExecArgs { } } +/// Execute a `-c` command string, optimizing single simple commands to exec +/// directly without forking. This avoids process group issues where grandchild +/// processes (e.g. nvim spawning opencode) lose their controlling terminal. +pub fn exec_dash_c(input: String) -> ShResult<()> { + let log_tab = read_logic(|l| l.clone()); + let expanded = expand_aliases(input, HashSet::new(), &log_tab); + let source_name = "".to_string(); + let mut parser = ParsedSrc::new(Arc::new(expanded)) + .with_lex_flags(super::lex::LexFlags::empty()) + .with_name(source_name.clone()); + if let Err(errors) = parser.parse_src() { + for error in errors { + error.print_error(); + } + return Ok(()); + } + + let mut nodes = parser.extract_nodes(); + + // Single simple command: exec directly without forking. + // The parser wraps single commands as Conjunction → Pipeline → Command. + // Unwrap all layers to check, then set NO_FORK on the inner Command. + if nodes.len() == 1 { + let is_single_cmd = match &nodes[0].class { + NdRule::Command { .. } => true, + NdRule::Pipeline { cmds } => cmds.len() == 1 && matches!(cmds[0].class, NdRule::Command { .. }), + NdRule::Conjunction { elements } => { + elements.len() == 1 && match &elements[0].cmd.class { + NdRule::Pipeline { cmds } => cmds.len() == 1 && matches!(cmds[0].class, NdRule::Command { .. }), + NdRule::Command { .. } => true, + _ => false, + } + } + _ => false, + }; + if is_single_cmd { + // Unwrap to the inner Command node + let mut node = nodes.remove(0); + loop { + match node.class { + NdRule::Conjunction { mut elements } => { node = *elements.remove(0).cmd; } + NdRule::Pipeline { mut cmds } => { node = cmds.remove(0); } + NdRule::Command { .. } => break, + _ => break, + } + } + node.flags |= NdFlags::NO_FORK; + nodes.push(node); + } + } + + let mut dispatcher = Dispatcher::new(nodes, false, source_name); + // exec_cmd expects a job on the stack (normally set up by exec_pipeline). + // For the NO_FORK exec-in-place path, create one so it doesn't panic. + dispatcher.job_stack.new_job(); + dispatcher.begin_dispatch() +} + pub fn exec_input( input: String, io_stack: Option, @@ -909,6 +947,7 @@ impl Dispatcher { "keymap" => keymap::keymap(cmd), "read_key" => read::read_key(cmd), "autocmd" => autocmd(cmd), + "ulimit" => ulimit(cmd), "true" | ":" => { state::set_status(0); Ok(()) @@ -946,7 +985,6 @@ impl Dispatcher { } let no_fork = cmd.flags.contains(NdFlags::NO_FORK); - if argv.is_empty() { return Ok(()); } @@ -959,36 +997,34 @@ impl Dispatcher { let existing_pgid = job.pgid(); let fg_job = self.fg_job; + let interactive = self.interactive; let child_logic = |pgid: Option| -> ! { - // Put ourselves in the correct process group before exec. - // For the first child in a pipeline pgid is None, so we - // become our own group leader (setpgid(0,0)). For later - // children we join the leader's group. - let our_pgid = pgid.unwrap_or(Pid::from_raw(0)); - let _ = setpgid(Pid::from_raw(0), our_pgid); + // For non-interactive exec-in-place (e.g. shed -c), skip process group + // and terminal setup — just transparently replace the current process. + if interactive || !no_fork { + // Put ourselves in the correct process group before exec. + // For the first child in a pipeline pgid is None, so we + // become our own group leader (setpgid(0,0)). For later + // children we join the leader's group. + let our_pgid = pgid.unwrap_or(Pid::from_raw(0)); + let _ = setpgid(Pid::from_raw(0), our_pgid); - // For foreground jobs, take the terminal BEFORE resetting - // signals. SIGTTOU is still SIG_IGN (inherited from the shell), - // so tcsetpgrp won't stop us. This prevents a race - // where the child exec's and tries to read stdin before the - // parent has called tcsetpgrp — which would deliver SIGTTIN - // (now SIG_DFL after reset_signals) and stop the child. - if fg_job { - let tty_pgid = if our_pgid == Pid::from_raw(0) { - nix::unistd::getpid() - } else { - our_pgid - }; - let _ = tcsetpgrp( - unsafe { BorrowedFd::borrow_raw(*crate::libsh::sys::TTY_FILENO) }, - tty_pgid, - ); + if fg_job { + let tty_pgid = if our_pgid == Pid::from_raw(0) { + nix::unistd::getpid() + } else { + our_pgid + }; + let _ = tcsetpgrp( + unsafe { BorrowedFd::borrow_raw(*crate::libsh::sys::TTY_FILENO) }, + tty_pgid, + ); + } } - // Reset signal dispositions before exec. SIG_IGN is preserved - // across execvpe, so the shell's ignored SIGTTIN/SIGTTOU would - // leak into child processes. - crate::signal::reset_signals(); + if interactive || !no_fork { + crate::signal::reset_signals(fg_job); + } let cmd = &exec_args.cmd.0; let span = exec_args.cmd.1; diff --git a/src/signal.rs b/src/signal.rs index 3a5655d..2c33418 100644 --- a/src/signal.rs +++ b/src/signal.rs @@ -154,10 +154,10 @@ pub fn sig_setup(is_login: bool) { } } -/// Reset all signal dispositions to SIG_DFL. +/// Reset signal dispositions to SIG_DFL. /// Called in child processes before exec so that the shell's custom /// handlers and SIG_IGN dispositions don't leak into child programs. -pub fn reset_signals() { +pub fn reset_signals(is_fg: bool) { let default = SigAction::new(SigHandler::SigDfl, SaFlags::empty(), SigSet::empty()); unsafe { for sig in Signal::iterator() { @@ -165,6 +165,10 @@ pub fn reset_signals() { if sig == Signal::SIGKILL || sig == Signal::SIGSTOP { continue; } + if is_fg && (sig == Signal::SIGTTIN || sig == Signal::SIGTTOU) { + log::debug!("Not resetting SIGTTIN/SIGTTOU in foreground child"); + continue; + } let _ = sigaction(sig, &default); } }