diff --git a/Cargo.lock b/Cargo.lock index 6c5d360..d261b65 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -14,6 +14,12 @@ version = "0.7.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d3fb67a6e08acf24fdeccbac2cb6ac4305825bd1f117462e0e6f2f193345ad56" +[[package]] +name = "base64" +version = "0.22.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" + [[package]] name = "bitflags" version = "1.3.2" @@ -157,6 +163,7 @@ dependencies = [ name = "fleetcom" version = "0.0.0" dependencies = [ + "base64", "crossterm", "dirs", "jzon", diff --git a/Cargo.toml b/Cargo.toml index 5876b70..0a47d94 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -11,6 +11,7 @@ repository = "https://github.com/ReagentX/fleetcom" version = "0.0.0" [dependencies] +base64 = "=0.22.1" crossterm = "=0.29.0" dirs = "=6.0.0" jzon = "=0.12.5" diff --git a/src/app.rs b/src/app.rs index e47bbe8..aa4c4a5 100644 --- a/src/app.rs +++ b/src/app.rs @@ -1222,9 +1222,12 @@ mod tests { impl App { /// A synchronous App: the supervisor ticks inline on `poll`, so `send` /// then `pump` is deterministic with no core-thread timing to race. + /// Uses this process's launch context. fn new_local(rows: u16, cols: u16) -> App { App::assemble(rows, cols, |pr, c, _wait_tx| { - Box::new(LocalTransport::new(Supervisor::new(pr, c))) + let mut sup = Supervisor::new(pr, c); + sup.set_launch_context(crate::protocol::LaunchContext::here()); + Box::new(LocalTransport::new(sup)) }) } diff --git a/src/core.rs b/src/core.rs index cfc9450..3eb3192 100644 --- a/src/core.rs +++ b/src/core.rs @@ -199,6 +199,7 @@ mod tests { let cwd = std::env::current_dir().unwrap(); let mut sup = Supervisor::new(24, 80); + sup.set_launch_context(crate::protocol::LaunchContext::here()); let (wake_tx, wake_rx) = channel::(); let (evt_tx, evt_rx) = channel::(); sup.set_waker(wake_tx.clone()); @@ -259,6 +260,7 @@ mod tests { fn stop_flag_ends_loop_with_shutdown() { let cwd = std::env::current_dir().unwrap(); let mut sup = Supervisor::new(24, 80); + sup.set_launch_context(crate::protocol::LaunchContext::here()); let (wake_tx, wake_rx) = std::sync::mpsc::channel::(); sup.set_waker(wake_tx); // Spawn synchronously so a live task exists *before* the loop runs: the diff --git a/src/daemon.rs b/src/daemon.rs index 4ed317c..c5ac65f 100644 --- a/src/daemon.rs +++ b/src/daemon.rs @@ -36,8 +36,8 @@ use nix::unistd::Pid; use crate::core::{LoopExit, Wake, run_loop}; use crate::frame::{read_frame, write_frame}; use crate::protocol::{ - Command, Event, PROTOCOL_VERSION, decode_command, decode_event, encode_command, encode_event, - hello_here, + Command, Event, LaunchContext, PROTOCOL_VERSION, decode_command, decode_event, decode_hello, + encode_command, encode_event, encode_hello, }; use crate::supervisor::Supervisor; @@ -193,7 +193,7 @@ pub fn connect_ready() -> io::Result { }); } - let (kind, payload) = encode_command(&hello_here()); + let (kind, payload) = encode_hello(&LaunchContext::here()); write_frame(&mut stream, kind, &payload)?; let reply = read_frame(&mut stream); done.store(true, Ordering::Relaxed); @@ -202,28 +202,34 @@ pub fn connect_ready() -> io::Result { Ok(stream) } +/// Convert handshake timeouts to a busy-daemon error; preserve other errors. +fn busy_daemon_error(e: io::Error) -> io::Error { + if is_timeout(&e) { + io::Error::new( + ErrorKind::TimedOut, + "the daemon is serving another client; retry after it detaches", + ) + } else { + e + } +} + /// The handshake for `reconnect`: called from inside the live UI (raw mode, /// alternate screen), where an unbounded wait would freeze the client and a /// printed notice would land on the alternate screen. A busy daemon surfaces /// as a status-line error instead; the user retries once the other client /// detaches. Write is bounded too: a full send buffer (large env, unaccepted -/// connection) must not wedge the UI either. +/// connection) must not wedge the UI either. A timed-out write drops the +/// connection, so a partial frame is never read. pub fn connect_ready_bounded() -> io::Result { let mut stream = connect_or_autostart()?; stream.set_write_timeout(Some(HANDSHAKE_TIMEOUT))?; - let (kind, payload) = encode_command(&hello_here()); - write_frame(&mut stream, kind, &payload)?; + let (kind, payload) = encode_hello(&LaunchContext::here()); + write_frame(&mut stream, kind, &payload).map_err(busy_daemon_error)?; stream.set_write_timeout(None)?; - let (kind, payload) = match read_frame_bounded(&mut stream, HANDSHAKE_TIMEOUT) { - Err(e) if is_timeout(&e) => { - return Err(io::Error::new( - ErrorKind::TimedOut, - "the daemon is serving another client; retry after it detaches", - )); - } - other => other.map_err(hello_read_error)?, - }; + let (kind, payload) = read_frame_bounded(&mut stream, HANDSHAKE_TIMEOUT) + .map_err(|e| hello_read_error(busy_daemon_error(e)))?; check_hello_ack(kind, &payload)?; Ok(stream) } @@ -336,7 +342,7 @@ fn kill_via_socket() -> io::Result<()> { let path = socket_path(); match UnixStream::connect(&path) { Ok(mut s) => { - let (kind, payload) = encode_command(&hello_here()); + let (kind, payload) = encode_hello(&LaunchContext::here()); write_frame(&mut s, kind, &payload)?; let (kind, payload) = encode_command(&Command::Shutdown); write_frame(&mut s, kind, &payload)?; @@ -388,7 +394,7 @@ pub fn run_daemon() -> io::Result<()> { fs::set_permissions(&path, fs::Permissions::from_mode(0o600))?; // 24x80 until the first client's Resize, which arrives before any Spawn. - // Launch context (env, session base dir) arrives per-connection via Hello. + // Each connection supplies its launch context in the hello frame. let mut sup = Supervisor::new(24, 80); // A signalled daemon shuts down *cleanly*: TERM each job's group with a @@ -474,33 +480,45 @@ enum ServeOutcome { Shutdown, } -/// Read and validate the connection-opening `Hello`. `Ok` carries the decoded -/// command for `apply` (it sets the client's launch context); `Err` carries the -/// refusal text for the client's status line. Bounded read: a peer that -/// connects and sends nothing must not wedge the daemon — accept, reap, and -/// `--kill` all wait behind this. -fn handshake(stream: &mut UnixStream) -> Result { +/// Return the version from a v2 control-frame hello. +fn is_v2_hello(kind: u8, payload: &[u8]) -> Option { + if kind != crate::frame::KIND_CONTROL { + return None; + } + let v = jzon::parse(std::str::from_utf8(payload).ok()?).ok()?; + if v["t"].as_str()? == "hello" { + v["v"].as_u32() + } else { + None + } +} + +/// Read and validate the connection-opening hello frame. +/// The bounded read prevents an idle peer from blocking the daemon. +fn handshake(stream: &mut UnixStream) -> Result { let (kind, payload) = read_frame_bounded(stream, HANDSHAKE_TIMEOUT) .map_err(|e| format!("no valid hello received: {e}"))?; - match decode_command(kind, &payload) { - Some( - hello @ Command::Hello { - version: PROTOCOL_VERSION, - .. - }, - ) => Ok(hello), - Some(Command::Hello { version, .. }) => Err(format!( + match decode_hello(kind, &payload) { + Some((PROTOCOL_VERSION, ctx)) => Ok(ctx), + Some((version, _)) => Err(format!( "protocol mismatch: daemon {} speaks v{PROTOCOL_VERSION}, client speaks \ v{version}; run 'fleetcom --kill' and retry", env!("CARGO_PKG_VERSION"), )), - // Reject commands received before `Hello`. - // Refuse requests before the required handshake. - _ => Err(format!( - "daemon {} requires a hello handshake (older client?); upgrade the \ - client or run 'fleetcom --kill' and retry", - env!("CARGO_PKG_VERSION"), - )), + // Report a v2 hello as a version mismatch. + None => match is_v2_hello(kind, &payload) { + Some(version) => Err(format!( + "protocol mismatch: daemon {} speaks v{PROTOCOL_VERSION}, client speaks \ + v{version}; run 'fleetcom --kill' and retry", + env!("CARGO_PKG_VERSION"), + )), + // Refuse anything else sent before the required handshake. + None => Err(format!( + "daemon {} requires a hello handshake (older client?); upgrade the \ + client or run 'fleetcom --kill' and retry", + env!("CARGO_PKG_VERSION"), + )), + }, } } @@ -519,8 +537,8 @@ fn handshake(stream: &mut UnixStream) -> Result { fn serve_client(sup: &mut Supervisor, stream: UnixStream, stop: &AtomicBool) -> ServeOutcome { let mut stream = stream; match handshake(&mut stream) { - Ok(hello) => { - sup.apply(hello); + Ok(ctx) => { + sup.set_launch_context(ctx); let (kind, payload) = encode_event(&Event::HelloOk); if write_frame(&mut stream, kind, &payload).is_err() { return ServeOutcome::Disconnected; diff --git a/src/frame.rs b/src/frame.rs index 18f138b..79763b3 100644 --- a/src/frame.rs +++ b/src/frame.rs @@ -11,6 +11,9 @@ pub const KIND_CONTROL: u8 = 1; /// A `Screen` event: a jzon header (id, cursor, lines) followed by the raw /// `contents_formatted` bytes, spliced by [`crate::protocol`]. pub const KIND_SCREEN: u8 = 2; +/// Connection-opening handshake containing the protocol version and launch +/// context. Handshakes are not command frames. +pub const KIND_HELLO: u8 = 3; /// Reject an absurd length prefix (corrupt or hostile peer) before allocating. /// 64 MiB is far above any real frame. A full 8K screen's formatted bytes are diff --git a/src/protocol.rs b/src/protocol.rs index 9fb7e1f..24d35ff 100644 --- a/src/protocol.rs +++ b/src/protocol.rs @@ -5,25 +5,39 @@ use std::os::unix::ffi::{OsStrExt, OsStringExt}; use std::path::{Path, PathBuf}; use std::time::Duration; -use crate::frame::{KIND_CONTROL, KIND_SCREEN}; +use base64::Engine as _; +use base64::engine::general_purpose::STANDARD as B64; + +use crate::frame::{KIND_CONTROL, KIND_HELLO, KIND_SCREEN}; use crate::task::Lifecycle; /// Wire-protocol version; the handshake rejects mismatched peers. -pub const PROTOCOL_VERSION: u32 = 2; +pub const PROTOCOL_VERSION: u32 = 3; + +/// Environment and working directory supplied by the launching client. +#[derive(Debug, Clone, PartialEq)] +pub struct LaunchContext { + pub env: Vec<(OsString, OsString)>, + /// Base directory for relative session-recipe paths. + pub cwd: PathBuf, +} + +impl LaunchContext { + /// Capture this process's environment and current directory. + pub fn here() -> LaunchContext { + LaunchContext { + env: std::env::vars_os().collect(), + cwd: std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")), + } + } +} /// A client→core request. Every mutation of the task set is one of these; the /// client never touches a `Task` directly. Fire-and-forget: results come back -/// as `Event`s, never as return values. +/// as `Event`s, never as return values. The handshake uses `KIND_HELLO`, not a +/// command. #[derive(Debug, Clone, PartialEq)] pub enum Command { - /// The client's protocol version, environment, and working directory. - Hello { - version: u32, - env: Vec<(OsString, OsString)>, - /// Base for resolving a session recipe's relative dirs: the client's - /// invocation dir, not the daemon's (frozen, first-client) cwd. - cwd: PathBuf, - }, /// Run `command` under `$SHELL -c` in `cwd`. Spawn { command: String, cwd: PathBuf }, /// Signal-kill a live task's process group; it reaps into Completed. @@ -101,7 +115,7 @@ pub enum MouseKind { /// the watched screen, updated only by these. #[derive(Debug, Clone, PartialEq)] pub enum Event { - /// The daemon accepted a compatible `Hello` and stored its launch context. + /// The daemon accepted a compatible hello frame and stored its launch context. HelloOk, /// Full task-set snapshot; replaces the client's mirror wholesale. Tasks(Vec), @@ -159,39 +173,14 @@ fn ps(p: &Path) -> String { p.to_string_lossy().into_owned() } -/// The `Hello` for this process: its own env and cwd at `PROTOCOL_VERSION`. -/// The socket client sends it as the handshake; the in-process transport -/// applies it directly — either way a launch context always comes from the -/// process that asked for the launch. -pub fn hello_here() -> Command { - Command::Hello { - version: PROTOCOL_VERSION, - env: std::env::vars_os().collect(), - cwd: std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")), - } -} - -/// An `OsStr` as a jzon byte array. JSON strings are UTF-8, so non-UTF-8 env -/// bytes can't ride as a string; a number array is lossless and paid once per -/// connection (the `Hello`), never per spawn. -fn os_arr(s: &OsStr) -> jzon::JsonValue { - let mut a = jzon::JsonValue::new_array(); - for b in s.as_bytes() { - let _ = a.push(*b as u64); - } - a +/// Encode an `OsStr` as lossless base64 for a JSON string. +fn os_b64(s: &OsStr) -> String { + B64.encode(s.as_bytes()) } -/// Decode a JSON byte array as an `OsString`. -fn os_from(v: &jzon::JsonValue) -> Option { - if !v.is_array() { - return None; - } - let mut bytes = Vec::with_capacity(v.len()); - for m in v.members() { - bytes.push(u8::try_from(m.as_u64()?).ok()?); - } - Some(OsString::from_vec(bytes)) +/// Decode a strictly valid base64 JSON string as an `OsString`. +fn os_from_b64(v: &jzon::JsonValue) -> Option { + Some(OsString::from_vec(B64.decode(v.as_str()?).ok()?)) } fn lifecycle_str(l: Lifecycle) -> &'static str { @@ -213,24 +202,47 @@ fn lifecycle_from(s: &str) -> Option { } } +/// Serialize a launch context as a `KIND_HELLO` frame. +pub fn encode_hello(ctx: &LaunchContext) -> (u8, Vec) { + let mut o = jzon::JsonValue::new_object(); + let _ = o.insert("v", PROTOCOL_VERSION); + let _ = o.insert("cwd", ps(&ctx.cwd)); + let mut pairs = jzon::JsonValue::new_array(); + for (k, v) in &ctx.env { + let mut pair = jzon::JsonValue::new_array(); + let _ = pair.push(os_b64(k)); + let _ = pair.push(os_b64(v)); + let _ = pairs.push(pair); + } + let _ = o.insert("env", pairs); + (KIND_HELLO, o.dump().into_bytes()) +} + +/// Parse a `KIND_HELLO` frame into `(version, context)`. +/// Returns `None` for malformed frames or environment entries. +pub fn decode_hello(kind: u8, payload: &[u8]) -> Option<(u32, LaunchContext)> { + if kind != KIND_HELLO { + return None; + } + let v = jzon::parse(std::str::from_utf8(payload).ok()?).ok()?; + let mut env = Vec::new(); + for pair in v["env"].members() { + env.push((os_from_b64(&pair[0])?, os_from_b64(&pair[1])?)); + } + Some(( + v["v"].as_u32()?, + LaunchContext { + env, + cwd: PathBuf::from(v["cwd"].as_str()?), + }, + )) +} + /// Serialize a command to `(kind, payload)` for [`crate::frame::write_frame`]. /// Every command is a jzon control frame tagged by a `"t"` discriminant. pub fn encode_command(cmd: &Command) -> (u8, Vec) { let mut o = jzon::JsonValue::new_object(); match cmd { - Command::Hello { version, env, cwd } => { - let _ = o.insert("t", "hello"); - let _ = o.insert("v", *version); - let _ = o.insert("cwd", ps(cwd)); - let mut pairs = jzon::JsonValue::new_array(); - for (k, v) in env { - let mut pair = jzon::JsonValue::new_array(); - let _ = pair.push(os_arr(k)); - let _ = pair.push(os_arr(v)); - let _ = pairs.push(pair); - } - let _ = o.insert("env", pairs); - } Command::Spawn { command, cwd } => { let _ = o.insert("t", "spawn"); let _ = o.insert("command", command.as_str()); @@ -342,17 +354,6 @@ pub fn decode_command(kind: u8, payload: &[u8]) -> Option { } let v = jzon::parse(std::str::from_utf8(payload).ok()?).ok()?; let cmd = match v["t"].as_str()? { - "hello" => { - let mut env = Vec::new(); - for pair in v["env"].members() { - env.push((os_from(&pair[0])?, os_from(&pair[1])?)); - } - Command::Hello { - version: v["v"].as_u32()?, - env, - cwd: PathBuf::from(v["cwd"].as_str()?), - } - } "spawn" => Command::Spawn { command: v["command"].as_str()?.to_string(), cwd: PathBuf::from(v["cwd"].as_str()?), @@ -566,24 +567,6 @@ mod tests { #[test] fn command_round_trips() { let cases = [ - Command::Hello { - version: PROTOCOL_VERSION, - // 0xFF/0xFE are invalid UTF-8 anywhere in a sequence: the env - // encoding must be byte-exact, not string-shaped. - env: vec![ - ("PATH".into(), "/usr/bin:/bin".into()), - ( - OsString::from_vec(b"BAD\xff\xfe".to_vec()), - OsString::from_vec(b"v\xff".to_vec()), - ), - ], - cwd: PathBuf::from("/home/x"), - }, - Command::Hello { - version: 0, - env: Vec::new(), - cwd: PathBuf::from("/"), - }, Command::Spawn { command: "echo hi".into(), cwd: PathBuf::from("/tmp"), @@ -670,26 +653,59 @@ mod tests { assert_eq!(decode_event(k, &p), Some(ack)); } - /// A malformed env pair (a non-byte element) rejects the whole `Hello`: - /// the daemon must never launch jobs under a guessed environment. + /// Handshake environment entries round-trip byte-for-byte. #[test] - fn hello_with_malformed_env_is_rejected() { - let json = r#"{"t":"hello","v":1,"cwd":"/","env":[[[300],[65]]]}"#; - assert_eq!(decode_command(KIND_CONTROL, json.as_bytes()), None); + fn hello_round_trips() { + let ctx = LaunchContext { + env: vec![ + ("PATH".into(), "/usr/bin:/bin".into()), + ( + OsString::from_vec(b"BAD\xff\xfe".to_vec()), + OsString::from_vec(b"v\xff".to_vec()), + ), + ], + cwd: PathBuf::from("/home/x"), + }; + let (k, p) = encode_hello(&ctx); + assert_eq!(k, KIND_HELLO); + assert_eq!(decode_hello(k, &p), Some((PROTOCOL_VERSION, ctx))); + + let empty = LaunchContext { + env: Vec::new(), + cwd: PathBuf::from("/"), + }; + let (k, p) = encode_hello(&empty); + assert_eq!(decode_hello(k, &p), Some((PROTOCOL_VERSION, empty))); } - /// Wrong-*shape* pairs must reject too, not decode as empty strings: a - /// client that encodes env entries as JSON strings would otherwise pass - /// the handshake and spawn PATH-less jobs with nothing pointing back at - /// the malformed hello. + /// A hello payload is valid only in a hello frame. #[test] - fn hello_with_string_env_pairs_is_rejected() { - // Pair elements as strings instead of byte arrays. - let strings = r#"{"t":"hello","v":1,"cwd":"/","env":[["PATH","/bin"]]}"#; - assert_eq!(decode_command(KIND_CONTROL, strings.as_bytes()), None); - // Pair itself as a string: indexing it yields Null for both elements. - let flat = r#"{"t":"hello","v":1,"cwd":"/","env":["PATH=/bin"]}"#; - assert_eq!(decode_command(KIND_CONTROL, flat.as_bytes()), None); + fn hello_requires_its_own_frame_kind() { + let (_, p) = encode_hello(&LaunchContext { + env: Vec::new(), + cwd: PathBuf::from("/"), + }); + assert_eq!(decode_hello(KIND_CONTROL, &p), None); + assert_eq!(decode_command(KIND_CONTROL, &p), None); + } + + /// Malformed environment entries reject the entire hello frame. + #[test] + fn hello_with_malformed_env_is_rejected() { + for env in [ + r#"[["P@TH","L2Jpbg=="]]"#, // invalid base64 character + r#"[["QUFBQUE","L2Jpbg=="]]"#, // truncated: missing padding + r#"[["UEFUSA==","AAAA="]]"#, // bad padding length + r#"[[[80],[65]]]"#, // v2 number arrays are not v3 + r#"["PATH=/bin"]"#, // flat string pair + ] { + let json = format!(r#"{{"v":3,"cwd":"/","env":{env}}}"#); + assert_eq!( + decode_hello(KIND_HELLO, json.as_bytes()), + None, + "should reject env {env}" + ); + } } #[test] diff --git a/src/supervisor.rs b/src/supervisor.rs index 4d96a37..6b5449b 100644 --- a/src/supervisor.rs +++ b/src/supervisor.rs @@ -4,7 +4,6 @@ //! a task snapshot plus the watched screen), and `drain` (take the queued //! `Event`s). -use std::ffi::OsString; use std::path::PathBuf; use std::sync::mpsc::Sender; use std::sync::{Arc, Mutex}; @@ -12,7 +11,7 @@ use std::time::{Duration, Instant}; use crate::core::{Wake, Waker}; use crate::path; -use crate::protocol::{Command, Event, ScreenView, ScrollAction, TaskView}; +use crate::protocol::{Command, Event, LaunchContext, ScreenView, ScrollAction, TaskView}; use crate::session::{self, SessionConfig}; use crate::task::Task; @@ -49,6 +48,9 @@ pub struct Supervisor { /// removal, escalated to KILL by `reap` at grace end, and dropped once the /// leader's zombie is collected. Invisible to `tick` snapshots, so the row /// disappears instantly while the sweep runs behind it. + /// + /// Entries remain through `kill_grace` because group emptiness cannot be + /// reliably observed before escalation. `shutdown_all` waits for them. graveyard: Vec, next_id: u64, /// PTY content size (rows already minus the client's status bar). Every task @@ -62,14 +64,9 @@ pub struct Supervisor { /// `None` whenever `watched` changes, so re-attaching always gets a fresh /// full screen (the client cleared its copy on detach). last_screen: Option, - /// The current client's launch context, from its `Hello`: every spawn - /// (including rerun and session load) uses this env, and session-recipe - /// dirs resolve against this cwd. Initialized to this process's own - /// context, which is the real thing for `--foreground` (client and core - /// share the process) and unreachable in the daemon: the handshake applies - /// the client's `Hello` before any command is served. - client_env: Vec<(OsString, OsString)>, - client_cwd: PathBuf, + /// The current client's launch context, used for spawns and session paths. + /// Spawning is refused until one is installed. + launch: Option, events: Vec, /// Handed to every `Task` so its reader thread can wake the core loop when the /// PTY produces output. The serving loop installs its sender on connect @@ -91,14 +88,18 @@ impl Supervisor { cols, watched: None, last_screen: None, - client_env: std::env::vars_os().collect(), - client_cwd: std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")), + launch: None, events: Vec::new(), waker: Arc::new(Mutex::new(None)), kill_grace: KILL_GRACE, } } + /// Install the launch context used by subsequent spawns. + pub fn set_launch_context(&mut self, ctx: LaunchContext) { + self.launch = Some(ctx); + } + /// Shrink the TERM→KILL grace so escalation tests run in milliseconds. #[cfg(test)] pub fn set_kill_grace(&mut self, grace: Duration) { @@ -135,13 +136,6 @@ impl Supervisor { /// notice, a spawn failure) is queued as `Event::Status`, never returned. pub fn apply(&mut self, cmd: Command) { match cmd { - // Version checking happens at the connection seam (the daemon's - // handshake, before anything reaches `apply`); here a Hello is - // purely the launch context taking effect. - Command::Hello { env, cwd, .. } => { - self.client_env = env; - self.client_cwd = cwd; - } Command::Spawn { command, cwd } => self.spawn(&command, cwd), Command::Kill { id } => { if let Some(t) = self.by_id_mut(id) { @@ -153,6 +147,7 @@ impl Supervisor { if let Some(i) = self.index_of(id) { let mut t = self.tasks.remove(i); t.terminate(); + t.shed_writer(); self.graveyard.push(t); } } @@ -227,10 +222,9 @@ impl Supervisor { } /// Kill every task for the quit path: TERM all groups at once, wait out one - /// shared grace (early exit as soon as every leader has exited), SIGKILL - /// the stragglers via `Task::drop`. Blocking here is fine (the core is - /// exiting), and the wait is bounded by the grace, paid only by jobs that - /// ignore their TERM. Anything the final KILLs don't collect (a leader in + /// shared grace (exiting early after all leaders and graveyard entries are + /// collected), then SIGKILL the stragglers. Blocking is bounded by the + /// grace. Anything the final KILLs don't collect (a leader in /// uninterruptible sleep) reparents to init when the daemon exits moments /// later; blocking on it here could wedge shutdown forever. fn shutdown_all(&mut self) { @@ -238,7 +232,9 @@ impl Supervisor { t.terminate(); } let deadline = Instant::now() + self.kill_grace; - while self.tasks.iter().any(|t| t.finished.is_none()) && Instant::now() < deadline { + while (self.tasks.iter().any(|t| t.finished.is_none()) || !self.graveyard.is_empty()) + && Instant::now() < deadline + { std::thread::sleep(Duration::from_millis(25)); self.reap(); } @@ -315,6 +311,16 @@ impl Supervisor { self.tasks.iter_mut().find(|t| t.id == id) } + /// Return the launch context, or report that spawning is unavailable. + fn launch_or_refuse(&mut self) -> Option { + if self.launch.is_none() { + self.events.push(Event::Status( + "no launch context; reconnect and retry".into(), + )); + } + self.launch.clone() + } + fn spawn(&mut self, command: &str, cwd: PathBuf) { if self.tasks.len() >= MAX_TASKS { self.events.push(Event::Status(format!( @@ -322,13 +328,16 @@ impl Supervisor { ))); return; } + let Some(launch) = self.launch_or_refuse() else { + return; + }; match Task::spawn( self.next_id, command, &cwd, self.rows, self.cols, - &self.client_env, + &launch.env, Arc::clone(&self.waker), ) { Ok(task) => { @@ -353,6 +362,9 @@ impl Supervisor { .push(Event::Status("rerun: task is still running".into())); return; } + let Some(launch) = self.launch_or_refuse() else { + return; + }; // Preserve the finished task if its replacement cannot start. match Task::spawn( id, @@ -360,7 +372,7 @@ impl Supervisor { &self.tasks[i].cwd, self.rows, self.cols, - &self.client_env, + &launch.env, Arc::clone(&self.waker), ) { Ok(mut fresh) => { @@ -370,6 +382,7 @@ impl Supervisor { // would straight-SIGKILL stragglers of the old run. let mut old = std::mem::replace(&mut self.tasks[i], fresh); old.terminate(); + old.shed_writer(); self.graveyard.push(old); // Reset the fingerprint for the replacement task's screen. if self.watched == Some(id) { @@ -419,9 +432,12 @@ impl Supervisor { return; } }; + let Some(launch) = self.launch_or_refuse() else { + return; + }; let (mut spawned, mut skipped) = (0usize, 0usize); for (dir, cmds) in &cfg { - let resolved = path::resolve(&self.client_cwd, dir); + let resolved = path::resolve(&launch.cwd, dir); if !resolved.is_dir() { skipped += cmds.len(); continue; @@ -437,7 +453,7 @@ impl Supervisor { &resolved, self.rows, self.cols, - &self.client_env, + &launch.env, Arc::clone(&self.waker), ) { self.next_id += 1; @@ -467,11 +483,18 @@ mod tests { std::env::current_dir().unwrap() } + /// Build a supervisor with this process's launch context. + fn sup(rows: u16, cols: u16) -> Supervisor { + let mut s = Supervisor::new(rows, cols); + s.set_launch_context(LaunchContext::here()); + s + } + /// The recipe groups commands by dir and preserves spawn order within a dir. /// `a`/`c` share the invocation dir; `b` is off in `/tmp`. #[test] fn session_config_groups_by_dir_in_spawn_order() { - let mut s = Supervisor::new(24, 80); + let mut s = sup(24, 80); s.apply(Command::Spawn { command: "a".into(), cwd: here(), @@ -498,7 +521,7 @@ mod tests { /// the client's render loop depends on. #[test] fn tick_emits_snapshot_and_watched_screen() { - let mut s = Supervisor::new(24, 80); + let mut s = sup(24, 80); s.apply(Command::Spawn { command: "sleep 30".into(), cwd: here(), @@ -530,7 +553,7 @@ mod tests { /// every tick: the send-on-change that kills idle attach churn. #[test] fn watched_screen_not_resent_when_unchanged() { - let mut s = Supervisor::new(24, 80); + let mut s = sup(24, 80); s.apply(Command::Spawn { command: "sleep 30".into(), cwd: here(), @@ -622,7 +645,7 @@ mod tests { use crate::task::Lifecycle; let dir = scratch("term_first"); let (ready, trapped) = (dir.join("ready"), dir.join("trapped")); - let mut s = Supervisor::new(24, 80); + let mut s = sup(24, 80); let id = spawn_ready( &mut s, format!( @@ -646,7 +669,7 @@ mod tests { use crate::task::Lifecycle; let dir = scratch("escalate"); let ready = dir.join("ready"); - let mut s = Supervisor::new(24, 80); + let mut s = sup(24, 80); s.set_kill_grace(Duration::from_millis(150)); let id = spawn_ready( &mut s, @@ -666,7 +689,7 @@ mod tests { /// grace, not after it. #[test] fn shutdown_returns_early_when_jobs_respect_term() { - let mut s = Supervisor::new(24, 80); + let mut s = sup(24, 80); s.apply(Command::Spawn { command: "sleep 300".into(), cwd: here(), @@ -691,7 +714,7 @@ mod tests { fn shutdown_is_bounded_by_grace() { let dir = scratch("shutdown_bound"); let ready = dir.join("ready"); - let mut s = Supervisor::new(24, 80); + let mut s = sup(24, 80); s.set_kill_grace(Duration::from_millis(200)); spawn_ready( &mut s, @@ -723,7 +746,7 @@ mod tests { /// fresh full screen instead of being skipped as "unchanged". #[test] fn clear_watch_stops_screen_stream_and_resets_dedup() { - let mut s = Supervisor::new(24, 80); + let mut s = sup(24, 80); s.apply(Command::Spawn { command: "sleep 30".into(), cwd: here(), @@ -767,7 +790,7 @@ mod tests { use crate::task::Lifecycle; let dir = scratch("restart"); let marker = dir.join("marker"); - let mut s = Supervisor::new(24, 80); + let mut s = sup(24, 80); s.apply(Command::Spawn { command: format!("echo run >> {}", marker.display()), cwd: dir.clone(), @@ -799,7 +822,7 @@ mod tests { #[test] fn restart_refuses_running_task_and_unknown_id() { use crate::task::Lifecycle; - let mut s = Supervisor::new(24, 80); + let mut s = sup(24, 80); s.apply(Command::Spawn { command: "sleep 30".into(), cwd: here(), @@ -840,7 +863,7 @@ mod tests { #[test] fn restart_watched_task_resends_screen() { use crate::task::Lifecycle; - let mut s = Supervisor::new(24, 80); + let mut s = sup(24, 80); s.apply(Command::Spawn { command: "true".into(), cwd: here(), @@ -873,7 +896,7 @@ mod tests { /// without a panic/OOM is the assertion. #[test] fn resize_clamps_hostile_dimensions() { - let mut s = Supervisor::new(24, 80); + let mut s = sup(24, 80); s.apply(Command::Spawn { command: "sleep 30".into(), cwd: here(), @@ -908,19 +931,12 @@ mod tests { pred(s) } - /// Pin the launch shell to `/bin/sh` via a `Hello`: the straggler tests - /// assert POSIX group mechanics, and zsh kills a `-c` shell's background - /// jobs on exit (even under `trap '' HUP`), which would end the straggler - /// before the sweep under test ran. + /// Use `/bin/sh` so background-process tests have consistent semantics. fn hello_with_sh(s: &mut Supervisor, cwd: PathBuf) { let mut env: Vec<(std::ffi::OsString, std::ffi::OsString)> = std::env::vars_os().collect(); env.retain(|(k, _)| k != "SHELL"); env.push(("SHELL".into(), "/bin/sh".into())); - s.apply(Command::Hello { - version: crate::protocol::PROTOCOL_VERSION, - env, - cwd, - }); + s.set_launch_context(LaunchContext { env, cwd }); } /// Read a pid a test job wrote, waiting for the write to land. @@ -947,7 +963,7 @@ mod tests { use nix::sys::signal::kill; let dir = scratch("remove_sweep"); let (spid, ready) = (dir.join("spid"), dir.join("ready")); - let mut s = Supervisor::new(24, 80); + let mut s = sup(24, 80); hello_with_sh(&mut s, dir.clone()); let id = spawn_ready( &mut s, @@ -988,7 +1004,7 @@ mod tests { use nix::sys::signal::kill; let dir = scratch("restart_sweep"); let (spid, ready) = (dir.join("spid"), dir.join("ready")); - let mut s = Supervisor::new(24, 80); + let mut s = sup(24, 80); hello_with_sh(&mut s, dir.clone()); let id = spawn_ready( &mut s, @@ -1030,7 +1046,7 @@ mod tests { use nix::sys::signal::kill; let dir = scratch("kill_escalate_straggler"); let (spid, ready) = (dir.join("spid"), dir.join("ready")); - let mut s = Supervisor::new(24, 80); + let mut s = sup(24, 80); s.set_kill_grace(Duration::from_millis(150)); hello_with_sh(&mut s, dir.clone()); // The leader ignores HUP (inherited by the `&` child, so it survives @@ -1060,13 +1076,53 @@ mod tests { let _ = std::fs::remove_dir_all(&dir); } - /// A spawn runs under the `Hello` client's environment, not the daemon - /// process's: the client's marker is visible, and a var only this process - /// has (`USER` — chosen because no shell synthesizes it, unlike `HOME`, - /// which zsh fills from passwd when unset) is absent because the builder's - /// captured base env is cleared. + /// Shutdown after removal preserves the removed task's TERM grace. + #[test] + fn shutdown_waits_for_graveyard_grace() { + use nix::sys::signal::kill; + let dir = scratch("shutdown_graveyard"); + let (spid, ready) = (dir.join("spid"), dir.join("ready")); + let mut s = sup(24, 80); + s.set_kill_grace(Duration::from_millis(400)); + hello_with_sh(&mut s, dir.clone()); + // The background process ignores HUP and TERM. + let id = spawn_ready( + &mut s, + format!( + "trap '' HUP; (trap '' TERM; exec sleep 300) & echo $! > {sp}; echo r > {r}", + sp = spid.display(), + r = ready.display() + ), + dir.clone(), + &ready, + ); + let straggler = read_pid(&spid); + assert!(reap_until(&mut s, Duration::from_secs(5), |s| { + s.tasks.iter().all(|t| t.id != id || t.finished.is_some()) + })); + + s.apply(Command::Remove { id }); // graveyard: TERM sent, grace running + // Check that the background process remains alive during the grace. + let alive_mid_grace = std::thread::spawn(move || { + std::thread::sleep(Duration::from_millis(200)); + kill(straggler, None).is_ok() + }); + s.apply(Command::Shutdown); + assert!( + alive_mid_grace.join().unwrap(), + "straggler was KILLed before its grace elapsed" + ); + assert!( + reap_until(&mut s, Duration::from_secs(5), |_| kill(straggler, None) + .is_err()), + "straggler survived shutdown" + ); + let _ = std::fs::remove_dir_all(&dir); + } + + /// Spawns inherit only the installed launch-context environment. #[test] - fn spawn_uses_the_hello_env_not_the_process_env() { + fn spawn_uses_the_launch_context_env_not_the_process_env() { assert!( std::env::var_os("USER").is_some(), "test needs USER set in the process env to prove it doesn't leak" @@ -1074,8 +1130,7 @@ mod tests { let dir = scratch("hello_env"); let out = dir.join("out"); let mut s = Supervisor::new(24, 80); - s.apply(Command::Hello { - version: crate::protocol::PROTOCOL_VERSION, + s.set_launch_context(LaunchContext { env: vec![("FLEETCOM_MARKER".into(), "xyzzy".into())], cwd: dir.clone(), }); @@ -1093,4 +1148,36 @@ mod tests { assert_eq!(std::fs::read_to_string(&out).unwrap(), "xyzzy:unset"); let _ = std::fs::remove_dir_all(&dir); } + + /// A supervisor with no launch context refuses to launch — spawn, rerun, + /// and session load alike — with a status notice. + #[test] + fn launch_without_context_is_refused() { + let mut s = Supervisor::new(24, 80); + s.apply(Command::Spawn { + command: "true".into(), + cwd: here(), + }); + assert!( + s.drain() + .iter() + .any(|e| matches!(e, Event::Status(m) if m.contains("no launch context"))), + "context-less spawn must be refused with a status notice" + ); + s.tick(); + assert!( + s.drain() + .iter() + .any(|e| matches!(e, Event::Tasks(v) if v.is_empty())), + "no task may exist after a refused spawn" + ); + + s.apply(Command::LoadSession { name: "any".into() }); + let evs = s.drain(); + assert!( + evs.iter().any(|e| matches!(e, Event::Status(m) + if m.contains("no launch context") || m.contains("not found"))), + "context-less load must not spawn; got {evs:?}" + ); + } } diff --git a/src/task.rs b/src/task.rs index 7cf252b..6e2afe3 100644 --- a/src/task.rs +++ b/src/task.rs @@ -9,7 +9,7 @@ use std::time::{Duration, Instant}; use nix::sys::signal::{Signal, killpg}; use nix::unistd::Pid; -use portable_pty::{Child, CommandBuilder, MasterPty, PtySize, native_pty_system}; +use portable_pty::{CommandBuilder, MasterPty, PtySize, native_pty_system}; use rustix::process::{WaitId, WaitIdOptions, waitid}; use crate::core::{Wake, Waker}; @@ -162,10 +162,7 @@ pub struct Task { /// Kept for resize (`TIOCSWINSZ`); `try_clone_reader`/`take_writer` borrow it. master: Box, writer: Box, - child: Box, - /// The leader's pid, cached at spawn. portable-pty `setsid`s the child, so - /// this is also the job's pgid: the target for group signals and the - /// `WNOWAIT` status latch. + /// Session-leader PID, also used as the process-group ID. pid: Option, /// Shared with the reader thread: it writes (process bytes), the UI reads /// (render/preview). Contention is trivial: writes are per output chunk. @@ -205,6 +202,14 @@ fn grid(parser: &Mutex) -> std::sync::MutexGuard<'_, vt100::Parse .unwrap_or_else(std::sync::PoisonError::into_inner) } +/// Convert a wait status to a shell-style exit code. +fn wait_code(status: &rustix::process::WaitIdStatus) -> i32 { + status + .exit_status() + .or_else(|| status.terminating_signal().map(|s| 128 + s)) + .unwrap_or(1) +} + impl Task { /// Spawn `command` under `$SHELL -c` in `cwd`, in a fresh PTY sized `rows`×`cols`, /// with exactly `env` as the environment (the launching client's; the caller @@ -299,14 +304,15 @@ impl Task { }) }; + // Process-group signalling and `waitid` use the leader PID directly. let pid = child.process_id(); + drop(child); Ok(Task { id, command: command.to_string(), cwd: cwd.to_path_buf(), master: pair.master, writer, - child, pid, parser, last_activity, @@ -338,13 +344,7 @@ impl Task { }; let flags = WaitIdOptions::EXITED | WaitIdOptions::NOWAIT | WaitIdOptions::NOHANG; if let Some(status) = waitid(WaitId::Pid(pid), flags)? { - // 128+signal mirrors the shell convention, so a KILLed job reads as - // 137 in the dashboard rather than masquerading as a clean exit. - let code = status - .exit_status() - .or_else(|| status.terminating_signal().map(|s| 128 + s)) - .unwrap_or(1); - self.exit_code = Some(code); + self.exit_code = Some(wait_code(&status)); self.finished = Some(Instant::now()); } Ok(()) @@ -355,20 +355,33 @@ impl Task { if self.reaped { return; } - if let Ok(Some(status)) = self.child.try_wait() { + let Some(pid) = self + .pid + .and_then(|p| rustix::process::Pid::from_raw(p as i32)) + else { + // No pid was ever known: nothing waitable or signalable exists. self.reaped = true; - if self.finished.is_none() { - self.exit_code = Some(status.exit_code() as i32); - self.finished = Some(Instant::now()); + return; + }; + match waitid( + WaitId::Pid(pid), + WaitIdOptions::EXITED | WaitIdOptions::NOHANG, + ) { + Ok(Some(status)) => { + self.reaped = true; + if self.finished.is_none() { + self.exit_code = Some(wait_code(&status)); + self.finished = Some(Instant::now()); + } } + // Treat an already-reaped leader as collected. + Err(rustix::io::Errno::CHILD) => self.reaped = true, + // Still running, or a transient failure: retry next reap pass. + Ok(None) | Err(_) => {} } } - /// Graveyard step: once the KILL has gone out, try to collect the leader's - /// zombie. Returns true when collected — nothing left to signal, so the - /// caller can drop this task silently. Never blocks: a leader wedged in - /// uninterruptible sleep (dead NFS/FUSE) stays uncollected and the caller - /// retries next reap pass instead of hanging the daemon. + /// After SIGKILL, try to collect the leader without blocking. pub fn try_collect(&mut self) -> bool { if self.kill_sent { self.collect(); @@ -376,6 +389,12 @@ impl Task { self.reaped } + /// Release the PTY writer after removal while retaining the master for + /// process-group teardown. + pub fn shed_writer(&mut self) { + self.writer = Box::new(io::sink()); + } + pub fn lifecycle(&self, now: Instant, idle_after: Duration) -> Lifecycle { if self.finished.is_some() { return if self.exit_code == Some(0) { @@ -704,6 +723,19 @@ mod tests { let _ = std::fs::remove_dir_all(&dir); } + /// `collect` reports SIGKILL as shell exit code 137. + #[test] + fn killed_leader_latches_137_via_collect() { + let mut t = spawn(8, "sleep 300"); + t.force_kill(); // sets kill_sent, so try_collect may reap + let deadline = Instant::now() + Duration::from_secs(5); + while !t.try_collect() && Instant::now() < deadline { + thread::sleep(Duration::from_millis(10)); + } + assert!(t.try_collect(), "KILLed leader was never collected"); + assert_eq!(t.exit_code, Some(137)); + } + /// Paste encoding follows the child's DECSET 2004 opt-in: markers only /// when asked for, newline→CR conversion only when not. #[test] diff --git a/src/transport.rs b/src/transport.rs index 06e5208..6869bcf 100644 --- a/src/transport.rs +++ b/src/transport.rs @@ -74,10 +74,8 @@ impl ThreadTransport { pub fn spawn(mut sup: Supervisor, wait_tx: Sender<()>) -> ThreadTransport { let (wake_tx, wake_rx) = channel::(); let (evt_tx, evt_rx) = channel::(); - // The in-process "connection": client and core are the same process, so - // its own env/cwd *is* the launch context. The daemon path receives the - // same Hello over the socket handshake instead. - sup.apply(crate::protocol::hello_here()); + // The in-process core uses this process's launch context. + sup.set_launch_context(crate::protocol::LaunchContext::here()); // Install the waker before the thread starts, so tasks spawned on the core // can signal output back to this same loop. sup.set_waker(wake_tx.clone()); diff --git a/tests/common/mod.rs b/tests/common/mod.rs index 76e165f..6b893df 100644 --- a/tests/common/mod.rs +++ b/tests/common/mod.rs @@ -15,15 +15,45 @@ use std::time::{Duration, Instant}; /// The protocol version this test suite speaks; must track /// `protocol::PROTOCOL_VERSION` (drift fails the handshake, loudly). -pub const PROTOCOL_VERSION: u32 = 2; +pub const PROTOCOL_VERSION: u32 = 3; + +/// One frame of the given kind: `[u32 len][kind][payload]`. +pub fn frame(kind: u8, payload: &[u8]) -> Vec { + let mut f = Vec::with_capacity(5 + payload.len()); + f.extend_from_slice(&(u32::try_from(payload.len()).unwrap()).to_be_bytes()); + f.push(kind); + f.extend_from_slice(payload); + f +} /// One `KIND_CONTROL` frame: `[u32 len][kind=1][jzon payload]`. pub fn control_frame(json: &str) -> Vec { - let mut f = Vec::with_capacity(5 + json.len()); - f.extend_from_slice(&(u32::try_from(json.len()).unwrap()).to_be_bytes()); - f.push(1); - f.extend_from_slice(json.as_bytes()); - f + frame(1, json.as_bytes()) +} + +/// Encode bytes as padded standard base64. +pub fn b64(bytes: &[u8]) -> String { + const ALPHABET: &[u8; 64] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"; + let mut out = String::with_capacity(bytes.len().div_ceil(3) * 4); + for chunk in bytes.chunks(3) { + let n = (u32::from(chunk[0]) << 16) + | (u32::from(*chunk.get(1).unwrap_or(&0)) << 8) + | u32::from(*chunk.get(2).unwrap_or(&0)); + let idx = [(n >> 18) & 63, (n >> 12) & 63, (n >> 6) & 63, n & 63]; + out.push(ALPHABET[idx[0] as usize] as char); + out.push(ALPHABET[idx[1] as usize] as char); + out.push(if chunk.len() > 1 { + ALPHABET[idx[2] as usize] as char + } else { + '=' + }); + out.push(if chunk.len() > 2 { + ALPHABET[idx[3] as usize] as char + } else { + '=' + }); + } + out } /// Read one frame, blocking: `(kind, payload)`. Mirrors `frame::read_frame`. @@ -38,22 +68,17 @@ pub fn read_frame(stream: &mut UnixStream) -> std::io::Result<(u8, Vec)> { Ok((kind[0], payload)) } -/// A `hello` control frame carrying `env` (byte-exact key/value pairs) and the -/// client cwd. Env values ride as JSON byte arrays: they need not be UTF-8. +/// A `KIND_HELLO` frame carrying the client environment and working directory. pub fn hello_frame(version: u32, env: &[(&[u8], &[u8])], cwd: &str) -> Vec { - let arr = |bytes: &[u8]| { - let nums: Vec = bytes.iter().map(|b| b.to_string()).collect(); - format!("[{}]", nums.join(",")) - }; let pairs: Vec = env .iter() - .map(|(k, v)| format!("[{},{}]", arr(k), arr(v))) + .map(|(k, v)| format!(r#"["{}","{}"]"#, b64(k), b64(v))) .collect(); let json = format!( - r#"{{"t":"hello","v":{version},"cwd":"{cwd}","env":[{}]}}"#, + r#"{{"v":{version},"cwd":"{cwd}","env":[{}]}}"#, pairs.join(",") ); - control_frame(&json) + frame(3, json.as_bytes()) } /// Complete the client side of the handshake: send a well-formed hello (this diff --git a/tests/daemon_handshake.rs b/tests/daemon_handshake.rs index 539bd9a..c4a790e 100644 --- a/tests/daemon_handshake.rs +++ b/tests/daemon_handshake.rs @@ -39,6 +39,18 @@ fn version_mismatch_is_refused_with_both_versions_named() { let _ = std::fs::remove_dir_all(&dir); } +/// A v2 control-frame hello is reported as a version mismatch. +#[test] +fn v2_hello_is_refused_as_a_version_mismatch() { + let (dir, daemon, mut stream) = start_daemon_raw("v2hello", |_| {}); + let cwd = dir.display().to_string(); + let v2 = format!(r#"{{"t":"hello","v":2,"cwd":"{cwd}","env":[[[65],[66]]]}}"#); + stream.write_all(&control_frame(&v2)).unwrap(); + expect_refusal(&mut stream, "v2"); + drop(daemon); + let _ = std::fs::remove_dir_all(&dir); +} + #[test] fn pre_handshake_command_is_refused() { let (dir, daemon, mut stream) = start_daemon_raw("nohello", |_| {});