diff --git a/libshpool/src/attach.rs b/libshpool/src/attach.rs index 6693e6a7..98e580ac 100644 --- a/libshpool/src/attach.rs +++ b/libshpool/src/attach.rs @@ -35,6 +35,7 @@ use crate::{ protocol::{ClientResult, PipeBytesResult}, template, test_hooks, tty::TtySizeExt as _, + ErrorStatusMode, }; const MAX_FORCE_RETRIES: usize = 20; @@ -46,6 +47,7 @@ pub fn run( name: String, force: bool, background: bool, + error_status_mode: ErrorStatusMode, ttl: Option, cmd: Option, dir: Option, @@ -83,7 +85,8 @@ pub fn run( None => None, }; - let attach = Attach { config_manager, force, background, ttl, tmpls, socket }; + let attach = + Attach { config_manager, force, background, error_status_mode, ttl, tmpls, socket }; attach.run() } @@ -92,6 +95,7 @@ struct Attach { config_manager: config::Manager, force: bool, background: bool, + error_status_mode: ErrorStatusMode, ttl: Option, tmpls: Templates, socket: PathBuf, @@ -140,19 +144,19 @@ impl Attach { pub fn attach_resolved(&self, resolved: ResolvedTemplates) -> anyhow::Result { if resolved.session_name.is_empty() { eprintln!("blank session names are not allowed"); - return Ok(AttachResult::Done); + return self.attach_error(anyhow!("blank session name")); } if resolved.session_name.contains(char::is_whitespace) { eprintln!("session name '{}' may not have whitespace", resolved.session_name); - return Ok(AttachResult::Done); + return self.attach_error(anyhow!("session name contains whitespace")); } if resolved.session_name.chars().any(|c| '/' == c) { eprintln!("session names may not contain slashes"); - return Ok(AttachResult::Done); + return self.attach_error(anyhow!("session name contains a slash")); } if resolved.session_name == "." || resolved.session_name == ".." { eprintln!("session names may not be special directory names"); - return Ok(AttachResult::Done); + return self.attach_error(anyhow!("session name is a special directory name")); } let mut detached = false; @@ -166,7 +170,7 @@ impl Attach { "session '{}' already has a terminal attached", resolved.session_name ); - return Ok(AttachResult::Done); + return self.attach_error(BusyError.into()); } Ok(BusyError) => { if !detached { @@ -232,12 +236,26 @@ impl Attach { let var_map: HashMap = maybe_switch.vars.iter().cloned().collect(); session_name_tmpl.apply(&var_map) != resolved.session_name }) { - Ok(PipeBytesResult::Exit(exit_status)) => std::process::exit(exit_status), + Ok(PipeBytesResult::Exit(exit_status)) => { + std::process::exit(if self.error_status_mode == ErrorStatusMode::Attach { + 0 + } else { + exit_status + }) + } Ok(PipeBytesResult::MaybeSwitch(s)) => Ok(AttachResult::Switch(s)), Err(e) => Err(e), } } + fn attach_error(&self, err: anyhow::Error) -> anyhow::Result { + if self.error_status_mode == ErrorStatusMode::Attach { + Err(err) + } else { + Ok(AttachResult::Done) + } + } + /// Attach to a session and return the connected client without piping /// stdio. fn dial_attach(&self, resolved: &ResolvedTemplates) -> anyhow::Result { diff --git a/libshpool/src/lib.rs b/libshpool/src/lib.rs index ee9e2ac1..18170c9b 100644 --- a/libshpool/src/lib.rs +++ b/libshpool/src/lib.rs @@ -21,7 +21,7 @@ use std::{ }; use anyhow::{anyhow, Context}; -use clap::{Parser, Subcommand}; +use clap::{Parser, Subcommand, ValueEnum}; pub use hooks::Hooks; use parking_lot::{Mutex, MutexGuard}; use tracing::error; @@ -116,6 +116,16 @@ the daemon is launched by systemd." pub __non_exhaustive: (), } +/// Selects which operation determines the exit status of `shpool attach`. +#[derive(Copy, Clone, Debug, Default, Eq, PartialEq, ValueEnum)] +pub enum ErrorStatusMode { + /// Report whether the attach operation succeeded. + Attach, + /// Report the exit status of the attached shell. + #[default] + Shell, +} + /// The subcommds that shpool supports. #[derive(Subcommand, Debug, Default)] #[non_exhaustive] @@ -138,6 +148,17 @@ pub enum Commands { help = "Create/attach the session and immediately detach (use with --force to detach any existing client first)" )] background: bool, + #[clap( + long, + value_enum, + default_value = "shell", + long_help = "Choose which operation determines the exit status + +In 'shell' mode, preserve the existing behavior: report the attached shell's +exit status and treat a busy session as a successful no-op. In 'attach' mode, +report whether the attach operation succeeded instead." + )] + error_status_mode: ErrorStatusMode, #[clap( long, long_help = "Automatically kill the session after the given time @@ -468,9 +489,27 @@ pub unsafe fn run( log_level_handle, socket, ), - Commands::Attach { force, background, ttl, cmd, dir, start_cmd, name } => { - attach::run(socket, config_manager, name, force, background, ttl, cmd, dir, start_cmd) - } + Commands::Attach { + force, + background, + error_status_mode, + ttl, + cmd, + dir, + start_cmd, + name, + } => attach::run( + socket, + config_manager, + name, + force, + background, + error_status_mode, + ttl, + cmd, + dir, + start_cmd, + ), Commands::Detach { sessions } => detach::run(sessions, socket), Commands::Kill { sessions } => kill::run(sessions, socket), Commands::List { json } => list::run(socket, json), diff --git a/shpool/tests/attach.rs b/shpool/tests/attach.rs index 2154535b..0efef785 100644 --- a/shpool/tests/attach.rs +++ b/shpool/tests/attach.rs @@ -560,6 +560,37 @@ fn busy_background() -> anyhow::Result<()> { Ok(()) } +#[test] +#[timeout(30000)] +fn busy_background_reports_attach_failure() -> anyhow::Result<()> { + let mut daemon_proc = support::daemon::Proc::new( + "norc.toml", + DaemonArgs { listen_events: false, ..DaemonArgs::default() }, + ) + .context("starting daemon proc")?; + + let mut tty1 = daemon_proc.attach("sh1", Default::default()).context("attaching from tty1")?; + let mut line_matcher1 = tty1.line_matcher()?; + tty1.run_cmd("echo foo")?; + line_matcher1.scan_until_re("foo$")?; + + let mut tty2 = daemon_proc + .attach( + "sh1", + AttachArgs { + background: true, + error_status_mode: Some("attach"), + ..Default::default() + }, + ) + .context("background attaching from tty2")?; + let mut line_matcher2 = tty2.stderr_line_matcher()?; + line_matcher2.scan_until_re("already has a terminal attached$")?; + assert_eq!(tty2.proc.wait()?.code(), Some(1), "busy attach should report failure"); + + Ok(()) +} + #[test] #[timeout(30000)] fn busy_background_force() -> anyhow::Result<()> { @@ -1119,6 +1150,25 @@ fn exits_with_same_status_as_shell() -> anyhow::Result<()> { Ok(()) } +#[test] +#[timeout(30000)] +fn attach_error_status_ignores_shell_status() -> anyhow::Result<()> { + let mut daemon_proc = support::daemon::Proc::new("norc.toml", DaemonArgs::default()) + .context("starting daemon proc")?; + let mut attach_proc = daemon_proc + .attach("sh", AttachArgs { error_status_mode: Some("attach"), ..Default::default() }) + .context("starting attach proc")?; + + attach_proc.run_cmd("exit 19")?; + + assert!( + attach_proc.proc.wait().context("waiting for attach proc to exit")?.success(), + "successful attach should exit successfully" + ); + + Ok(()) +} + #[test] #[timeout(30000)] fn ttl_hangup() -> anyhow::Result<()> { diff --git a/shpool/tests/support/daemon.rs b/shpool/tests/support/daemon.rs index c2715409..4e464f41 100644 --- a/shpool/tests/support/daemon.rs +++ b/shpool/tests/support/daemon.rs @@ -51,6 +51,7 @@ pub struct AttachArgs { pub config: Option, pub force: bool, pub background: bool, + pub error_status_mode: Option<&'static str>, pub extra_env: Vec<(String, String)>, pub ttl: Option, pub cmd: Option, @@ -363,6 +364,10 @@ impl Proc { if args.background { cmd.arg("-b"); } + if let Some(error_status_mode) = args.error_status_mode { + cmd.arg("--error-status-mode"); + cmd.arg(error_status_mode); + } if let Some(ttl) = args.ttl { cmd.arg("--ttl"); cmd.arg(format!("{}s", ttl.as_secs()));