diff --git a/Cargo.lock b/Cargo.lock index 4ad1905d..36881b9a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5022,6 +5022,7 @@ dependencies = [ name = "sush-tests" version = "0.1.0" dependencies = [ + "blake3", "bytes", "camino", "chrono", diff --git a/client/src/cli.rs b/client/src/cli.rs index a21c1ec2..54b2988c 100644 --- a/client/src/cli.rs +++ b/client/src/cli.rs @@ -353,18 +353,36 @@ impl CommandContext for Cli { ) { let mut progress = self.progress.lock().unwrap(); if matches!(self.get_output_format(), OutputFormat::Text) && progress.is_none() { - let bar = ProgressBar::new(total_length); + // With no known total, show a spinner and a running byte count. + let bar = if total_length == 0 { + let bar = ProgressBar::new_spinner(); + bar.set_style( + ProgressStyle::with_template( + "{spinner} \ + {prefix} \ + [{elapsed_precise}] \ + {decimal_bytes:>7} \ + {msg}", + ) + .unwrap(), + ); + bar.enable_steady_tick(Duration::from_millis(100)); + bar + } else { + let bar = ProgressBar::new(total_length); + bar.set_style( + ProgressStyle::with_template( + "{prefix} \ + [{elapsed_precise}] \ + {bar:40.cyan/blue} \ + {decimal_bytes:>7}/{decimal_total_bytes:7} \ + {msg}", + ) + .unwrap(), + ); + bar + }; bar.set_prefix(format!("{stage} {stream}")); - bar.set_style( - ProgressStyle::with_template( - "{prefix} \ - [{elapsed_precise}] \ - {bar:40.cyan/blue} \ - {decimal_bytes:>7}/{decimal_total_bytes:7} \ - {msg}", - ) - .unwrap(), - ); *progress = Some(bar); } } @@ -384,7 +402,7 @@ impl CommandContext for Cli { if let Some(progress) = self.progress.lock().unwrap().take() { progress.finish_and_clear(); if let Some(stage) = stage { - let length = ByteSize::b(progress.length().unwrap_or(0)); + let length = ByteSize::b(progress.length().unwrap_or(progress.position())); let elapsed = progress.elapsed(); println!( "{stage} {stream}:\t{} in {} ({:.0} MB/s)", diff --git a/client/src/commands.rs b/client/src/commands.rs index c2407945..8f04422b 100644 --- a/client/src/commands.rs +++ b/client/src/commands.rs @@ -7,6 +7,8 @@ //! May be executed via either the main CLI or the interactive REPL. use std::fs::{File, OpenOptions, read}; +#[cfg(feature = "permslip")] +use std::io::ErrorKind; use std::io::{Read as _, Seek as _, SeekFrom, Write as _, stdin}; use std::num::{NonZeroU8, NonZeroU64}; use std::path::{Path, PathBuf}; @@ -42,13 +44,13 @@ use sush_api::JobWait; use sush_common::authn::{ AuthnError, Challenge, ChallengeResponse, Credentials, Identity, RequestKey, }; -use sush_common::interactive::InteractiveJobError; +use sush_common::interactive::{InteractiveJobError, InteractiveJobMessage}; use sush_common::jobs::JobOutputStream::{self, Stderr, Stdout}; #[cfg(feature = "permslip")] use sush_common::jobs::JobStartRequest; use sush_common::jobs::{ Access, JobId, JobLimits, JobOutputHash, JobOutputState, JobStatus, JobStatusMap, Session, - SessionId, SignedJob, job_status_try_from_json_map, + SessionId, SignedJob, Streaming, job_status_try_from_json_map, }; use sush_common::keys::{KeyError, KeyId, Signer as _}; use sush_common::targets::{SledId, Target}; @@ -247,37 +249,60 @@ fn watch_settling() { time_started: Utc::now(), }; - let mut quiet = 0; + let all = Target::All; + let mut settling = Settling::default(); let mut status = JobStatusMap::new(); - assert!( - !watch_done(1, 0, &status, &mut quiet), - "empty never settles" - ); + assert!(!settling.done(&all, None, &status), "empty never settles"); status.insert(sled("A"), terminal.clone()); - assert!(!watch_done(2, 0, &status, &mut quiet), "new sled resets"); - assert!(!watch_done(3, 1, &status, &mut quiet), "first quiet poll"); + assert!(!settling.done(&all, None, &status), "new sled resets"); + assert!(!settling.done(&all, None, &status), "first quiet poll"); assert!( - !watch_done(4, 1, &status, &mut quiet), + !settling.done(&all, None, &status), "quiet but too young for gossip" ); - assert!(watch_done(5, 1, &status, &mut quiet), "old enough, quiet"); + assert!(settling.done(&all, None, &status), "old enough, quiet"); - let mut quiet = 0; - status.insert(sled("B"), running); - assert!(!watch_done(6, 2, &status, &mut quiet), "running sled holds"); - assert_eq!(quiet, 0); + let mut settling = Settling::default(); + status.insert(sled("B"), running.clone()); + assert!(!settling.done(&all, None, &status), "running sled holds"); + status.insert(sled("B"), terminal.clone()); + assert!( + !settling.done(&all, Some(2), &status), + "the rack count must hold" + ); + assert!( + settling.done(&all, Some(2), &status), + "the rack count held twice" + ); + let named: Target = format!("{},{}", sled("A"), sled("B")).parse().unwrap(); + let mut settling = Settling::default(); + assert!( + settling.done(&named, None, &status), + "named sleds settle on the first terminal poll" + ); + let duplicates: Target = format!("{},{}", sled("A"), sled("A")).parse().unwrap(); + let mut settling = Settling::default(); + assert!( + settling.done(&duplicates, None, &status), + "duplicate baseboards collapse" + ); + status.insert(sled("B"), running); + assert!( + !settling.done(&named, None, &status), + "a running named sled holds" + ); status.insert(sled("B"), terminal); - assert!(!watch_done(7, 2, &status, &mut quiet)); + let missing: Target = format!("{},{}", sled("A"), sled("C")).parse().unwrap(); + let mut settling = Settling::default(); assert!( - watch_done(8, 2, &status, &mut quiet), - "settles once terminal" + !settling.done(&missing, None, &status), + "a missing named sled holds" ); } -/// Anything the target grammar accepts stays a target; bare serials -/// fall through; everything else still fails. +/// Anything the target grammar accepts stays a target. #[test] fn target_arg() { assert!(matches!("*".parse(), Ok(TargetArg::Target(_)))); @@ -286,8 +311,23 @@ fn target_arg() { "913-0000019:BRM42220030".parse(), Ok(TargetArg::Target(_)) )); - assert!(matches!("brm42220030".parse(), Ok(TargetArg::Serial(s)) if s == "brm42220030")); - assert!("not,a:target!".parse::().is_err()); + assert!(matches!( + "brm42220030".parse(), + Ok(TargetArg::Abbreviated(sleds)) + if matches!(sleds.as_slice(), [SledArg::Serial(s)] if s == "brm42220030") + )); + assert!(matches!( + "14,brm42220030".parse(), + Ok(TargetArg::Abbreviated(sleds)) + if matches!(sleds.as_slice(), [SledArg::Sled(SledId::Cubby(14)), SledArg::Serial(_)]) + )); + assert!(matches!( + "913-0000019:BRM42220036,brm42220030".parse(), + Ok(TargetArg::Abbreviated(sleds)) + if matches!(sleds.as_slice(), [SledArg::Sled(SledId::Baseboard(_)), SledArg::Serial(_)]) + )); + assert!("n!ot,a:target".parse::().is_err()); + assert!("*,brm42220030".parse::().is_err()); } /// [`ClientArgs`] must satisfy Clap's internal consistency asserts @@ -571,10 +611,22 @@ pub struct JobStartArgs { #[arg(short, long)] interactive: bool, + /// Stream the job's output to an attached client instead of recording it. + #[arg(short = 'S', long, conflicts_with = "interactive")] + streaming: bool, + + /// File that streamed output should be written to. + #[arg(short, long)] + file: Option, + + /// Overwrite the output file if it exists. + #[arg(long, requires = "file")] + force: bool, + /// Where the job runs: every sled (`*`), or a comma-separated list - /// of cubby numbers and baseboard IDs. + /// of cubby numbers, baseboard IDs, or bare serial numbers. #[arg(short = 'T', long, default_value = "*")] - target: Target, + target: TargetArg, /// Use `permslip` to sign job requests with this key name. #[cfg(feature = "permslip")] @@ -1049,6 +1101,7 @@ async fn job( ref job_id, ref permslip_url, ref interactive, + ref streaming, ref target, .. }, @@ -1088,15 +1141,38 @@ async fn job( let Some(permslip_url) = permslip_url else { return Err(CommandError::MissingPermslipUrl); }; - // Sign an interactive job for the sled its attachment - // will land on. - let target = if *interactive && target.single_baseboard().is_none() { + let target = match target { + TargetArg::Target(target) => target.clone(), + abbreviated => { + let Some(client) = client.as_ref() else { + return Err(CommandError::Offline); + }; + resolve_target_arg(client, abbreviated).await? + } + }; + // Sign an interactive or streaming job for the sled its + // attachment will land on. + let target = if (*interactive || *streaming) && target.single_baseboard().is_none() { let Some(client) = client.as_ref() else { return Err(CommandError::InteractiveTarget); }; - Target::from(resolve_target(client, target).await?) + Target::from(resolve_target(client, &target).await?) } else { - target.clone() + target + }; + // Catch output file problems before the signing ceremony. + if *streaming { + let Some(path) = &start_args.file else { + return Err(CommandError::StreamingNeedsFile); + }; + if !start_args.force && path.exists() { + return Err(CommandError::io(path, ErrorKind::AlreadyExists.into())); + } + } + let streaming = if *streaming { + Streaming::Output + } else { + Streaming::None }; let mut signer = PermslipSigner::new(key_name, permslip_url).await?; let mut interval = interval(SIGNING_UPDATE_INTERVAL); @@ -1105,6 +1181,7 @@ async fn job( job_id.to_owned(), command, *interactive, + streaming, target, )); pin!(sign); @@ -1146,7 +1223,7 @@ async fn job( StatusDisplayStyle::Short }; if wait { - let status = job_watch(ctx, client, &job_id).await?; + let status = job_watch(ctx, client, &job_id, &Target::All).await?; ctx.job_status(&job_id, &status, style); Ok(()) } else { @@ -1164,8 +1241,16 @@ async fn job( (JobCommand::Attach { job_id, target }, Some(client)) => { let target = match &target { + TargetArg::Abbreviated(sleds) => match sleds.as_slice() { + [SledArg::Serial(serial)] => { + resolve_serial(ctx, client, &job_id, serial).await? + } + _ => { + let target = resolve_target_arg(client, &target).await?; + resolve_target(client, &target).await? + } + }, TargetArg::Target(target) => resolve_target(client, target).await?, - TargetArg::Serial(serial) => resolve_serial(ctx, client, &job_id, serial).await?, }; job_attach(ctx, client, &job_id, &target).await?; Ok(()) @@ -1213,10 +1298,13 @@ async fn job_start( limits, term, wait, + file, + force, .. } = start_args; let interactive = job.payload().interactive; - let wait = if interactive { + let streaming = matches!(job.payload().streaming, Streaming::Output); + let wait = if interactive || streaming { JobWait::Start } else if wait { JobWait::Stop @@ -1260,16 +1348,60 @@ async fn job_start( } Err(error) => return Err(error), } + } else if streaming { + let Some(path) = file else { + return Err(CommandError::StreamingNeedsFile); + }; + let mut options = OpenOptions::new(); + options.write(true); + if force { + options.create(true); + } else { + options.create_new(true); + } + let file = options + .open(&path) + .map_err(|error| CommandError::io(&path, error))?; + start.await?; + ctx.job_started(&job); + let hasher = job_stream(ctx, client, &job_id, &target, &path, file).await?; + let status = job_watch(ctx, client, &job_id, &target.clone().into()).await?; + ctx.job_status(&job_id, &status, StatusDisplayStyle::Short); + let Some(JobStatus::Stopped { + result: Ok(0), + output: + JobOutputState { + stdout_len, + stdout_hash, + .. + }, + .. + }) = status.get(&target) + else { + return Err(CommandError::StreamUnverified(job_id)); + }; + if *stdout_len != hasher.count() { + return Err(CommandError::LengthMismatch { + expected: *stdout_len, + received: hasher.count(), + }); + } + let received = JobOutputHash::from(hasher.finalize()); + if received != *stdout_hash { + return Err(CommandError::OutputHashMismatch { + expected: stdout_hash.to_owned(), + received, + }); + } } else if wait.is_some() { // Watch the whole rack while the start request runs, and keep // watching until the job settles everywhere. ctx.job_watch_started(&job_id); + let job_target = job.payload().target().clone(); let mut ticker = interval(WATCH_INTERVAL); ticker.set_missed_tick_behavior(MissedTickBehavior::Skip); let mut last = JobStatusMap::new(); - let mut sleds = 0; - let mut quiet = 0; - let mut polls = 0; + let mut settling = Settling::default(); let mut started = false; let mut stopped = false; let mut sigint = signal(SignalKind::interrupt())?; @@ -1285,6 +1417,7 @@ async fn job_start( } } started = true; + ticker.reset_immediately(); } _ = ticker.tick() => { @@ -1298,11 +1431,10 @@ async fn job_start( } }; ctx.job_watch_update(&last); - polls += 1; - if started && watch_done(polls, sleds, &last, &mut quiet) { + let rack = rack_sleds(client, &job_target).await; + if settling.done(&job_target, rack, &last) && started { break last; } - sleds = last.len(); } // While the job runs, an interrupt stops it but keeps @@ -1346,23 +1478,33 @@ async fn job_start( CommandError::JobStillRunning(job_id) }); } - for stream in [Stdout, Stderr] { - match with_login_via(ctx, client, Some(&target), async || { - client - .job_output() - .job_id(job_id) - .target(target.to_string()) - .stream(stream) - .send() - .await - }) - .await - { - Ok(byte_stream) => { - let output = byte_stream_to_vec(byte_stream.into_inner()).await?; - ctx.job_output(&job_id, stream, &output, binary); + // Show the output of every sled the job ran on. + let multiple = status.len() > 1; + for baseboard in status.keys() { + if multiple { + ctx.job_output_target(baseboard); + } + for stream in [Stdout, Stderr] { + match with_login_via(ctx, client, Some(baseboard), async || { + client + .job_output() + .job_id(job_id) + .target(baseboard.to_string()) + .stream(stream) + .send() + .await + }) + .await + { + Ok(byte_stream) => { + let output = byte_stream_to_vec(byte_stream.into_inner()).await?; + ctx.job_output(&job_id, stream, &output, binary); + } + Err(error) if multiple => { + let _ = ctx.job_error(error); + } + Err(error) => return Err(error), } - Err(error) => return Err(error), } } } else { @@ -1423,11 +1565,10 @@ async fn job_watch( ctx: &mut impl CommandContext, client: &Client, job_id: &JobId, + target: &Target, ) -> Result { ctx.job_watch_started(job_id); - let mut sleds = 0; - let mut quiet = 0; - let mut polls = 0; + let mut settling = Settling::default(); let mut sigint = signal(SignalKind::interrupt())?; let status = loop { let status = match job_status_map(ctx, client, job_id).await { @@ -1438,11 +1579,10 @@ async fn job_watch( } }; ctx.job_watch_update(&status); - polls += 1; - if watch_done(polls, sleds, &status, &mut quiet) { + let rack = rack_sleds(client, target).await; + if settling.done(target, rack, &status) { break status; } - sleds = status.len(); select! { _ = sleep(WATCH_INTERVAL) => {} _ = sigint.recv() => break status, @@ -1457,16 +1597,55 @@ async fn job_watch( /// is likely still missing sleds. const WATCH_MIN_POLLS: usize = 5; -/// A watched job has settled once the watch is old enough for gossip -/// to have named every sled, every known sled reports a terminal -/// status, and the set of sleds has been stable for a couple of polls. -fn watch_done(polls: usize, sleds: usize, status: &JobStatusMap, quiet: &mut usize) -> bool { - if !status.is_empty() && status.len() == sleds && status.values().all(JobStatus::is_terminal) { - *quiet += 1; - } else { - *quiet = 0; +/// Rolling settlement state for a watched job. +#[derive(Default)] +struct Settling { + polls: usize, + sleds: usize, + quiet: usize, + met: usize, +} + +impl Settling { + /// A watched job has settled once every sled it runs on reports a + /// terminal status. A target naming only baseboards is settled as + /// soon as they all report. Against the rack inventory, the count + /// must match on consecutive polls. Either way, the watch is done + /// once it is old enough for gossip to have named every sled and + /// the set of sleds has been stable for a couple of polls. + fn done(&mut self, target: &Target, rack: Option, status: &JobStatusMap) -> bool { + self.polls += 1; + let terminal = !status.is_empty() && status.values().all(JobStatus::is_terminal); + if terminal && status.len() == self.sleds { + self.quiet += 1; + } else { + self.quiet = 0; + } + self.sleds = status.len(); + if rack.is_some_and(|sleds| terminal && status.len() == sleds) { + self.met += 1; + } else { + self.met = 0; + } + if let Some(named) = target.named_baseboards() + && named + .into_iter() + .all(|sled| status.get(sled).is_some_and(JobStatus::is_terminal)) + { + return true; + } + self.met >= 2 || (self.polls >= WATCH_MIN_POLLS && self.quiet >= 2) } - polls >= WATCH_MIN_POLLS && *quiet >= 2 +} + +/// The number of sleds in the rack's own inventory, where available. +/// Skips the fetch for targets whose sleds are named. +async fn rack_sleds(client: &Client, target: &Target) -> Option { + if target.named_baseboards().is_some() { + return None; + } + let sleds = client.versions().send().await.ok()?.into_inner().len(); + (sleds > 0).then_some(sleds) } /// Stream some bytes into a vector. @@ -1510,14 +1689,17 @@ async fn job_output( args: JobOutput, ) -> Result<(), CommandError> { let target = match target { - TargetArg::Serial(serial) => { - let baseboard = resolve_serial(ctx, client, &args.job_id, serial).await?; - return job_output_from(ctx, client, &baseboard, stream, args).await; + TargetArg::Target(target) => target.clone(), + TargetArg::Abbreviated(sleds) => { + if let [SledArg::Serial(serial)] = sleds.as_slice() { + let baseboard = resolve_serial(ctx, client, &args.job_id, serial).await?; + return job_output_from(ctx, client, &baseboard, stream, args).await; + } + resolve_target_arg(client, target).await? } - TargetArg::Target(target) => target, }; if !target.is_all() { - let baseboard = resolve_target(client, target).await?; + let baseboard = resolve_target(client, &target).await?; return job_output_from(ctx, client, &baseboard, stream, args).await; } if args.binary || args.file.is_some() { @@ -1778,6 +1960,70 @@ async fn job_attach( } } +/// Receive a streaming job's output into an open file, returning the +/// hasher fed by the received bytes. +async fn job_stream( + ctx: &mut impl CommandContext, + client: &Client, + job_id: &JobId, + target: &BaseboardId, + path: &Path, + mut file: File, +) -> Result { + let io_error = |error| CommandError::io(path, error); + let socket = with_login_via(ctx, client, Some(target), async || { + client + .job_attach() + .job_id(job_id) + .target(target.to_string()) + .send() + .await + }) + .await? + .into_inner(); + let mut stream = WebSocketStream::from_raw_socket(socket, Role::Client, None).await; + file.set_len(0).map_err(io_error)?; + let mut hasher = Hasher::new(); + let mut sigint = signal(SignalKind::interrupt())?; + ctx.job_output_started(job_id, Stdout, "Streaming", 0); + let result = loop { + select! { + message = stream.next() => { + let Some(message) = message else { break Ok(()) }; + let message = match message.map_err(InteractiveJobError::from) { + Ok(message) => message, + Err(error) => break Err(error.into()), + }; + match InteractiveJobMessage::try_from(message) { + Ok(InteractiveJobMessage::Data(bytes)) => { + hasher.update(&bytes); + if let Err(error) = file.write_all(&bytes).map_err(io_error) { + break Err(error); + } + ctx.job_output_update(job_id, Stdout, bytes.len() as u64); + } + Ok(InteractiveJobMessage::Close) => break Ok(()), + Ok(InteractiveJobMessage::Control(_) | InteractiveJobMessage::Ignore) => (), + Err(error) => break Err(error.into()), + } + } + _ = sigint.recv() => break Err(CommandError::Canceled), + } + }; + match result { + Ok(()) => { + file.flush().map_err(io_error)?; + ctx.job_output_finished(job_id, Stdout, Some("✅ Streamed")); + let _ = stream.close(None).await; + Ok(hasher) + } + Err(error) => { + ctx.job_output_finished(job_id, Stdout, Some("❌ Received")); + Err(error) + } + } +} + /// Resolve a target to the baseboard of one sled it names. A single /// baseboard resolves locally. Anything else asks `/target`, routed /// by the expression, so a proxy resolves cubbies and `*` means the @@ -1795,11 +2041,18 @@ async fn resolve_target(client: &Client, target: &Target) -> Result), +} + +/// One sled named in a target argument. +#[derive(Clone, Debug)] +pub enum SledArg { + Sled(SledId), Serial(String), } @@ -1807,12 +2060,54 @@ impl FromStr for TargetArg { type Err = ::Err; fn from_str(s: &str) -> Result { - match s.parse() { - Ok(target) => Ok(Self::Target(target)), - Err(_) if !s.is_empty() && s.chars().all(|c| c.is_ascii_alphanumeric()) => { - Ok(Self::Serial(s.to_owned())) + let error = match s.parse() { + Ok(target) => return Ok(Self::Target(target)), + Err(error) => error, + }; + let mut sleds = Vec::new(); + for piece in s.split(',') { + let piece = piece.trim(); + match piece.parse::() { + Ok(Target::Sleds(sled)) if sled.len() == 1 => { + sleds.push(SledArg::Sled(sled.into_iter().next().expect("one sled"))); + } + _ if !piece.is_empty() && piece.chars().all(|c| c.is_ascii_alphanumeric()) => { + sleds.push(SledArg::Serial(piece.to_owned())); + } + _ => return Err(error), } - Err(err) => Err(err), + } + Ok(Self::Abbreviated(sleds)) + } +} + +/// Resolve a target argument to a target, matching bare serial +/// numbers against the rack's sled inventory. +async fn resolve_target_arg(client: &Client, target: &TargetArg) -> Result { + match target { + TargetArg::Target(target) => Ok(target.clone()), + TargetArg::Abbreviated(sleds) => { + let inventory = client.versions().send().await?.into_inner(); + sleds + .iter() + .map(|sled| match sled { + SledArg::Sled(sled) => Ok(sled.clone()), + SledArg::Serial(serial) => { + let mut matches = inventory + .iter() + .map(|sled| &sled.baseboard) + .filter(|b| b.serial_number.eq_ignore_ascii_case(serial)); + match (matches.next(), matches.next()) { + (Some(baseboard), None) => Ok(SledId::Baseboard(baseboard.clone())), + (None, _) => Err(CommandError::UnknownRackSerial(serial.to_owned())), + (Some(_), Some(_)) => { + Err(CommandError::AmbiguousSerial(serial.to_owned())) + } + } + } + }) + .collect::, _>>() + .map(Target::Sleds) } } } @@ -1874,7 +2169,7 @@ pub enum CommandError { Interactive(#[from] InteractiveJobError), #[error("❌ Interactive jobs must target exactly one sled")] InteractiveTarget, - #[error("❌ I/O error accessing `{path}`: {error}")] + #[error("❌ Local I/O error accessing `{path}`: {error}")] Io { path: PathBuf, error: std::io::Error, @@ -1948,6 +2243,10 @@ pub enum CommandError { #[cfg(not(feature = "permslip"))] #[error("❌ Signing unavailable: built without the `permslip` feature")] SigningUnavailable, + #[error("❌ Streaming jobs need `--file`")] + StreamingNeedsFile, + #[error("❌ Job `{0}` did not exit cleanly, streamed output is unverified")] + StreamUnverified(JobId), #[error("❌ Can't parse target baseboard ID: {0}")] BaseboardIdParseError(sled_hardware_types::BaseboardIdParseError), #[error("❌ SSH key error: {0}")] @@ -1958,6 +2257,8 @@ pub enum CommandError { TooMuchOutput, #[error("❌ No sled with serial `{serial}` has a status for job `{job_id}`")] UnknownSerial { serial: String, job_id: JobId }, + #[error("❌ Serial `{0}` matches no sled in the rack inventory")] + UnknownRackSerial(String), #[error("❌ Chain root does not match any supplied root certificate")] UntrustedRoot, #[error("❌ Can't start interactive session: {0}")] diff --git a/client/src/lib.rs b/client/src/lib.rs index 025532a8..cb30c43d 100644 --- a/client/src/lib.rs +++ b/client/src/lib.rs @@ -80,6 +80,7 @@ progenitor::generate_api!( Signature = sush_common::keys::Signature, SignedForJobStartRequest = sush_common::jobs::SignedJob, SledVersion = sush_common::targets::SledVersion, + Streaming = sush_common::jobs::Streaming, VersionInfo = sush_common::version::VersionInfo, }, timeout = 600, diff --git a/common/src/jobs.rs b/common/src/jobs.rs index f2fc2e6c..9bbed4e2 100644 --- a/common/src/jobs.rs +++ b/common/src/jobs.rs @@ -164,6 +164,45 @@ impl Session { } } +/// Allow **unrecorded** streaming I/O. +#[derive( + BorshDeserialize, + BorshSerialize, + Clone, + Copy, + Debug, + Default, + Deserialize, + Eq, + JsonSchema, + PartialEq, + Serialize, +)] +pub enum Streaming { + #[default] + None, + Input, + Output, +} + +impl Streaming { + pub fn is_none(&self) -> bool { + matches!(self, Self::None) + } + + pub fn is_some(&self) -> bool { + !self.is_none() + } + + pub fn as_str(&self) -> &'static str { + match self { + Self::None => "none", + Self::Input => "input", + Self::Output => "output", + } + } +} + /// A request to run the given `command` as `job_id`. #[derive( BorshDeserialize, @@ -181,6 +220,8 @@ pub struct JobStartRequest { pub command: String, #[serde(default, skip_serializing_if = "is_false")] pub interactive: bool, + #[serde(default, skip_serializing_if = "Streaming::is_none")] + pub streaming: Streaming, /// The sleds this job runs on. #[borsh( serialize_with = "borsh_ser_target", @@ -201,12 +242,14 @@ impl JobStartRequest { job_id: JobId, command: S, interactive: bool, + streaming: Streaming, target: Target, ) -> Self { Self { job_id, command: command.as_ref().to_string(), interactive, + streaming, target, } } @@ -250,12 +293,18 @@ impl ToBeSigned for JobStartRequest { job_id, command, interactive, + streaming, target, } = self; hash_with_len(Self::TYPE_NAME); hash_with_len(&job_id.to_be_bytes()); hash_with_len(command.as_bytes()); hash_with_len(if *interactive { &[1] } else { &[0] }); + hash_with_len(match streaming { + Streaming::None => &[0], + Streaming::Input => &[1], + Streaming::Output => &[2], + }); hash_with_len(target.to_string().as_bytes()); hasher.finalize().as_bytes().to_vec() } @@ -736,7 +785,13 @@ mod test { let me = sled("me"); let cubbies = Cubbies::from([(14, me.clone())]); let job = |interactive, target| { - JobStartRequest::new(JobId::random(), "true", interactive, target) + JobStartRequest::new( + JobId::random(), + "true", + interactive, + Streaming::None, + target, + ) }; let just_me = Target::Sleds(vec![SledId::Baseboard(me.clone())]); assert!(job(true, just_me.clone()).runs_on(&me, &cubbies)); diff --git a/common/src/targets.rs b/common/src/targets.rs index 0a866f2b..8d865713 100644 --- a/common/src/targets.rs +++ b/common/src/targets.rs @@ -10,7 +10,7 @@ //! against a mapping the rack learns at runtime, so a target is //! evaluated lazily, against what is known when it is asked. -use std::collections::BTreeMap; +use std::collections::{BTreeMap, BTreeSet}; use std::fmt; use std::str::FromStr; @@ -85,6 +85,20 @@ impl Target { }, } } + + /// Every sled this target names, when it names only baseboards. + pub fn named_baseboards(&self) -> Option> { + match self { + Self::All => None, + Self::Sleds(sleds) => sleds + .iter() + .map(|sled| match sled { + SledId::Baseboard(baseboard) => Some(baseboard), + SledId::Cubby(_) => None, + }) + .collect(), + } + } } impl From for Target { diff --git a/server/src/executor.rs b/server/src/executor.rs index ba511276..b8b8b849 100644 --- a/server/src/executor.rs +++ b/server/src/executor.rs @@ -32,7 +32,7 @@ use tokio_util::sync::CancellationToken; use sush_api::JobStartParams; use sush_common::interactive::WindowSize; use sush_common::jobs::{ - JobId, JobOutputStream, JobStartRequest, ProcessError, SignedJob, VerifiedJob, + JobId, JobOutputStream, JobStartRequest, ProcessError, SignedJob, Streaming, VerifiedJob, }; use crate::io::JobIo; @@ -180,7 +180,8 @@ async fn job_spawn( job_id, command, interactive, - target: _, + streaming, + target, } = request.payload().clone(); let JobStartParams { limits: requested, @@ -210,6 +211,22 @@ async fn job_spawn( limits.max_fsize = dirs.max_fsize(); } + if interactive && streaming.is_some() { + let error = ProcessError::InvalidJob("interactive jobs cannot stream".to_string()); + send_error(&log, &job_id, &events, error).await; + return; + } + if matches!(streaming, Streaming::Input) { + let error = ProcessError::InvalidJob("streaming input is not implemented".to_string()); + send_error(&log, &job_id, &events, error).await; + return; + } + if streaming.is_some() && target.single_baseboard().is_none() { + let error = ProcessError::InvalidJob("streaming jobs must target one sled".to_string()); + send_error(&log, &job_id, &events, error).await; + return; + } + // Report all I/O errors as job events. let io_err = |what| move |err: io::Error| ProcessError::io(what, err); macro_rules! with_io_err { @@ -238,16 +255,23 @@ async fn job_spawn( ); let stdout_path = dirs.job_output_path(&job_id, Stdout); let stderr_path = dirs.job_output_path(&job_id, Stderr); - let stdout_file = with_io_err!( - OpenOptions::new() - .create_new(true) - .read(true) // needed for interactive job output playback - .write(true) - .mode(file_mode) - .open(&stdout_path) - .await, - format!("creating job stdout file `{}`", stdout_path.display()) - ); + let stdout_file = if matches!(streaming, Streaming::Output) { + with_io_err!( + OpenOptions::new().write(true).open("/dev/null").await, + "opening /dev/null for streamed output".to_string() + ) + } else { + with_io_err!( + OpenOptions::new() + .create_new(true) + .read(true) // needed for interactive job output playback + .write(true) + .mode(file_mode) + .open(&stdout_path) + .await, + format!("creating job stdout file `{}`", stdout_path.display()) + ) + }; let stderr_file = with_io_err!( OpenOptions::new() .create_new(true) @@ -305,7 +329,8 @@ async fn job_spawn( } cmd.env("SSH_CLIENT", "sush") // read bashrc .env("SUSH_JOB_ID", job_id.to_string()) - .env("SUSH_COMMAND", &command); + .env("SUSH_COMMAND", &command) + .env("SUSH_JOB_OUTPUT_DIR", &job_dir); // Set process limits. unsafe { @@ -382,6 +407,7 @@ async fn job_spawn( io, stdout_file, stderr_file, + streaming, stop, ) } else { @@ -397,12 +423,13 @@ async fn job_spawn( child.stderr.take().expect("batch job should have stderr"), ); Job::start( - log.new(o!("interactive" => interactive)), + log.new(o!("interactive" => interactive, "streaming" => streaming.as_str())), limits, child, io, stdout_file, stderr_file, + streaming, stop, ) }; diff --git a/server/src/io.rs b/server/src/io.rs index a7aaecaa..610def50 100644 --- a/server/src/io.rs +++ b/server/src/io.rs @@ -14,6 +14,7 @@ use std::io; use std::time::Duration; use bytes::{Bytes, BytesMut}; +use rustix::io::Errno; use tokio::io::{AsyncRead, AsyncReadExt as _, AsyncWriteExt as _}; use tokio::process::{ChildStderr, ChildStdout}; use tokio::sync::mpsc; @@ -112,10 +113,16 @@ impl JobIo { pub async fn read_output(&mut self) -> io::Result<(Bytes, JobOutputStream)> { match self { Self::Interactive { pty, .. } => loop { - match pty.read().await? { - Some(b) => return Ok((b, Stdout)), - None if pty.eof => return Ok((Bytes::new(), Stdout)), - None => (), // drained, keep reading + match pty.read().await { + Ok(Some(b)) => return Ok((b, Stdout)), + Ok(None) if pty.eof => return Ok((Bytes::new(), Stdout)), + Ok(None) => (), // drained, keep reading + Err(error) if Errno::from_io_error(&error) == Some(Errno::IO) => { + // A pty master reads EIO once its last slave closes, + // which is EOF here, not an error. + return Ok((Bytes::new(), Stdout)); + } + Err(error) => return Err(error), } }, Self::Batch { stdout, stderr } => loop { @@ -162,11 +169,11 @@ impl JobIo { /// How long to keep reading output after the child dies. /// - /// For a pty there is no EOF; a short quiet period is the - /// only way to know we've drained (OpenSSH does this too). - /// Pipes deliver EOF once every writer exits, so this is only - /// a backstop against a descendant that escaped the process - /// group while holding the inherited pipe open. + /// For a pty there is no portable EOF signal; a short quiet period + /// is the only way to know we've drained (OpenSSH does this too). + /// Pipes deliver EOF once every writer exits, so this is only a + /// backstop against a descendant that escaped the process group + /// while holding the inherited pipe open. pub fn drain_timeout(&self) -> Duration { match self { Self::Interactive { .. } => Duration::from_millis(10), diff --git a/server/src/job.rs b/server/src/job.rs index 8c69d5dd..9873f785 100644 --- a/server/src/job.rs +++ b/server/src/job.rs @@ -43,7 +43,9 @@ use tokio_tungstenite::tungstenite::protocol::Message as WebSocketMessage; use sush_common::interactive::{ INTERACTIVE_JOB_BUFFER_SIZE, InteractiveJobControl as Control, InteractiveJobMessage as Message, }; -use sush_common::jobs::{Access, JobLimits, JobOutputState, JobOutputStream::*, ProcessError}; +use sush_common::jobs::{ + Access, JobLimits, JobOutputState, JobOutputStream::*, ProcessError, Streaming, +}; use tokio_util::sync::CancellationToken; use crate::executor::kill_job; @@ -54,12 +56,17 @@ pub type SocketStream = WebSocketStream; pub type SocketSender = mpsc::Sender<(SocketStream, Access)>; pub type SocketReceiver = mpsc::Receiver<(SocketStream, Access)>; +/// How long a finished streaming job waits for its consumer to +/// attach or make progress. +const STREAMING_LINGER: Duration = Duration::from_secs(600); + pub struct Job { task: JoinHandle<(Result, JobOutputState)>, tx_client: SocketSender, } impl Job { + #[allow(clippy::too_many_arguments)] pub fn start( log: Logger, limits: JobLimits, @@ -67,11 +74,14 @@ impl Job { io: JobIo, stdout: File, stderr: File, + streaming: Streaming, stop: CancellationToken, ) -> Self { let (tx_client, rx_client) = mpsc::channel(1); Self { - task: spawn(job(log, limits, child, io, stdout, stderr, rx_client, stop)), + task: spawn(job( + log, limits, child, io, stdout, stderr, streaming, rx_client, stop, + )), tx_client, } } @@ -95,6 +105,7 @@ async fn job( mut io: JobIo, mut stdout_file: File, mut stderr_file: File, + streaming: Streaming, mut rx_client: SocketReceiver, stop: CancellationToken, ) -> (Result, JobOutputState) { @@ -104,22 +115,54 @@ async fn job( let mut fatal = Option::::None; let mut killed = false; let mut dead = false; + let streaming = matches!(streaming, Streaming::Output); + let mut pending = Option::::None; + let mut attached = false; + let (tx_output, rx_output) = mpsc::channel(1); + let (tx_socket, rx_socket) = mpsc::channel(1); + if streaming { + spawn(stream_output( + log.clone(), + rx_socket, + rx_output, + stop.clone(), + )); + } + let mut eof = false; let drain_timeout = sleep(Duration::default()); pin!(drain_timeout); + let linger = sleep(Duration::default()); + pin!(linger); loop { + // A finished stream is done once its consumer has taken every chunk. + if eof && attached && pending.is_none() { + break; + } select! { // Read available job output, record it, and relay it to the clients // if there are any. We try to read regardless of whether the process // is known to be dead; it is essential to drain output that may be // sent before it dies, but which arrives after detection of its death. - read = io.read_output() => { + read = io.read_output(), if pending.is_none() && !eof => { match read { Ok((buf, _)) if buf.is_empty() => { debug!(log, "EOF on all job output streams"); - break; + if !streaming { + break; + } + eof = true; + linger.as_mut().reset(Instant::now() + STREAMING_LINGER); } Ok((buf, stream)) => { + if dead && streaming { + drain_timeout.as_mut().reset(Instant::now() + io.drain_timeout()); + } match stream { + Stdout if streaming => { + stdout_hasher.update(&buf); + pending = Some(buf); + continue; + } Stdout => { stdout_hasher.update(&buf); if let Err(error) = stdout_file.write_all(&buf).await { @@ -167,6 +210,11 @@ async fn job( if !dead { error!(log, "error reading from job process"; "error" => %error); } + if streaming { + eof = true; + linger.as_mut().reset(Instant::now() + STREAMING_LINGER); + continue; + } if let Err(error) = clients.send(Message::Close.try_into().unwrap()) { error!(log, "failed to send close message to clients"; "error" => %error); } @@ -175,8 +223,34 @@ async fn job( } } + permit = tx_output.reserve(), if pending.is_some() => { + match permit { + Ok(permit) => { + if dead { + drain_timeout.as_mut().reset(Instant::now() + io.drain_timeout()); + } + if eof { + linger.as_mut().reset(Instant::now() + STREAMING_LINGER); + } + if let Some(buf) = pending.take() { + permit.send(buf); + } + } + Err(_) => pending = None, + } + } + // Attach a new client, send it the current window size, and play back the last buffer. - Some((mut client, access)) = rx_client.recv(), if !dead => { + // Streaming jobs accept one read-write consumer, even posthumously. + Some((mut client, access)) = rx_client.recv(), if !dead || streaming => { + if streaming { + if access != Access::ReadWrite || attached || tx_socket.try_send(client).is_err() { + debug!(log, "refused streaming client"); + } else { + attached = true; + } + continue; + } match io.get_window_size() { Err(error) => error!(log, "failed to get pseudoterminal window size"; "error" => %error), Ok(size) => { @@ -259,8 +333,14 @@ async fn job( dead = true; } + // Give up on a stream whose consumer never drained it. + _ = &mut linger, if eof => { + warn!(log, "streaming consumer never drained the output"); + break; + } + // Give output a chance to drain from a dead process. - _ = &mut drain_timeout, if dead => { + _ = &mut drain_timeout, if dead && !eof => { match io { JobIo::Interactive { .. } => debug!(log, "drained job output"), JobIo::Batch { .. } => warn!( @@ -274,6 +354,10 @@ async fn job( } } + if let Some(buf) = pending { + let _ = tx_output.try_send(buf); + } + // Reap the process. let exit_status = select! { status = child.wait() => status, @@ -303,6 +387,27 @@ async fn job( (result, output_state) } +/// Deliver streamed output to the single attached client. +async fn stream_output( + log: Logger, + mut rx_socket: mpsc::Receiver, + mut rx_output: mpsc::Receiver, + stop: CancellationToken, +) { + let Some(mut client) = rx_socket.recv().await else { + return; + }; + while let Some(buf) = rx_output.recv().await { + if let Err(error) = client.send(Message::Data(buf).try_into().unwrap()).await { + error!(log, "failed to stream output to client"; "error" => %error); + stop.cancel(); + return; + } + } + let _ = client.send(Message::Close.try_into().unwrap()).await; + let _ = client.close(None).await; +} + fn process_exit(exit_status: ExitStatus) -> Result { if let Some(code) = exit_status.code() { Ok(code) @@ -382,6 +487,7 @@ mod test { .open("/dev/null") .await .unwrap(), + Streaming::None, stop, ); assert_eq!(job.wait().await.unwrap().0.unwrap(), 0); diff --git a/server/src/messages.rs b/server/src/messages.rs index cfebefeb..6274df81 100644 --- a/server/src/messages.rs +++ b/server/src/messages.rs @@ -321,7 +321,7 @@ mod wire_format { #[test] fn job_start_request() { - use sush_common::jobs::JobStartRequest; + use sush_common::jobs::{JobStartRequest, Streaming}; use sush_common::keys::{EncodedSignature, Signed}; let request = JobStartRequest::new( @@ -330,6 +330,7 @@ mod wire_format { .unwrap(), "echo hello", false, + Streaming::None, "14,16".parse().unwrap(), ); let signed = Signed::new( diff --git a/server/tests/common/mod.rs b/server/tests/common/mod.rs index 1a62c123..8c909018 100644 --- a/server/tests/common/mod.rs +++ b/server/tests/common/mod.rs @@ -27,7 +27,7 @@ use sprockets_tls_test_utils::{ }; use sush_common::authn::{Challenge, ChallengeResponse, Identity, Nonce, RequestKey}; use sush_common::codephrases::Codephrase; -use sush_common::jobs::{JobId, JobStartRequest, SignedJob}; +use sush_common::jobs::{JobId, JobStartRequest, SignedJob, Streaming}; use sush_common::keys::{EphemeralKey, KeyType, Signer as _}; use sush_common::targets::Target; use sush_server::gossip::GossipConfig; @@ -154,6 +154,7 @@ pub async fn sign_job(root: &mut EphemeralKey, job_id: &JobId, command: &str) -> job_id.to_owned(), command, false, + Streaming::None, Target::All, )) .await diff --git a/server/tests/output/job-start-request.bin b/server/tests/output/job-start-request.bin index 3fa78f1f..a57b2e7f 100644 Binary files a/server/tests/output/job-start-request.bin and b/server/tests/output/job-start-request.bin differ diff --git a/sush.json b/sush.json index dcd49e8b..7e975112 100644 --- a/sush.json +++ b/sush.json @@ -1207,6 +1207,9 @@ "job_id": { "$ref": "#/components/schemas/JobId" }, + "streaming": { + "$ref": "#/components/schemas/Streaming" + }, "target": { "description": "The sleds this job runs on.", "type": "string" @@ -1635,6 +1638,15 @@ "baseboard" ] }, + "Streaming": { + "description": "Allow **unrecorded** streaming I/O.", + "type": "string", + "enum": [ + "None", + "Input", + "Output" + ] + }, "VersionInfo": { "description": "One build's version and commit.", "type": "object", diff --git a/tests/Cargo.toml b/tests/Cargo.toml index 9d491150..03713cab 100644 --- a/tests/Cargo.toml +++ b/tests/Cargo.toml @@ -8,6 +8,7 @@ repository = "https://github.com/oxidecomputer/sush" publish = false [dependencies] +blake3.workspace = true bytes.workspace = true camino.workspace = true chrono.workspace = true diff --git a/tests/src/integration_tests.rs b/tests/src/integration_tests.rs index 385204bb..777b0547 100644 --- a/tests/src/integration_tests.rs +++ b/tests/src/integration_tests.rs @@ -18,7 +18,7 @@ use futures::{SinkExt as _, StreamExt as _}; use tokio::io::{AsyncRead, AsyncWrite}; use tokio::sync::watch; use tokio::test; -use tokio::time::timeout; +use tokio::time::{sleep, timeout}; use tokio_tungstenite::WebSocketStream; use tokio_tungstenite::tungstenite::protocol::Role; use tokio_util::sync::CancellationToken; @@ -41,7 +41,10 @@ use sush_api::sush_api_mod::api_description; use sush_client::tls::client as tls_client; use sush_client::{AuthzSigner, Client, Error as ClientError}; use sush_common::interactive::{InteractiveJobControl, InteractiveJobMessage}; -use sush_common::jobs::{Access, JobLimits, JobOutputStream, Session, SessionId}; +use sush_common::jobs::{ + Access, JobLimits, JobOutputState, JobOutputStream, JobStatus, Session, SessionId, + job_status_try_from_json_map, +}; use sush_common::keys::{EphemeralKey, KeyType, pem_cert_chain}; use sush_common::targets::Cubbies; use sush_server::proxy::{Targets, platform_tls}; @@ -804,3 +807,268 @@ async fn interactive_job() { hello_again, ); } + +#[named] +#[test] +async fn streaming_job() { + const LEN: usize = 0x40000; + + // Spin up a server. + let log = test_logger(function_name!()); + let (mgr, mut root, _dir, _shutdown) = manager_and_test_root(log.clone()).await; + let api = api_description::().unwrap(); + let server = ServerBuilder::new(api, Arc::new(mgr), log) + .config(ConfigDropshot { + bind_address: local_addr(), + ..Default::default() + }) + .start() + .expect("failed to start server"); + + // Connect and authenticate to the server. + let addr = server.local_addr(); + let signer = AuthzSigner::default(); + let client = Client::new(&format!("http://{addr}"), signer.clone()); + let ClientError::ErrorResponse(unauthz) = client.iam().body(None).send().await.unwrap_err() + else { + panic!("expected error response") + }; + let (_identity, credentials) = authz(&client, unauthz, &mut root).await; + signer.set(Some(credentials)); + + // Start a streaming job too big for its pipe, so it blocks until we attach. + let session = Session::new(SessionId::random()); + client + .session_start() + .session_id(session.session_id()) + .send() + .await + .expect("can't start session"); + let job_id = session.next_job_id(); + let job = root + .sign_streaming_job_request( + &job_id, + "dd if=/dev/zero bs=4096 count=64 2>/dev/null", + test_baseboard_id().into(), + ) + .await; + let JobLimits { + max_cpu, + max_mem, + max_fsize, + } = JobLimits::default(); + client + .job_start() + .job_id(job_id) + .target(test_baseboard_id().to_string()) + .max_cpu(max_cpu) + .max_mem(max_mem) + .max_fsize(max_fsize) + .wait(JobWait::Start) + .body(job.into_signed()) + .send() + .await + .expect("can't start job"); + + // Attach and collect the whole stream. + let socket = client + .job_attach() + .job_id(job_id) + .target(test_baseboard_id().to_string()) + .send() + .await + .expect("can't attach to job") + .into_inner(); + let mut stream = WebSocketStream::from_raw_socket(socket, Role::Client, None).await; + let mut streamed = BytesMut::new(); + loop { + let recvd = timeout(TIMEOUT, stream.next()) + .await + .expect("timed out waiting for streamed output") + .expect("stream ended without close") + .expect("can't get next stream item"); + match InteractiveJobMessage::try_from(recvd).expect("can't decode message") { + InteractiveJobMessage::Data(bytes) => streamed.extend(bytes), + InteractiveJobMessage::Close => break, + message => panic!("unexpected message: {message:?}"), + } + } + assert_eq!(streamed.len(), LEN); + assert!(streamed.iter().all(|byte| *byte == 0)); + + // The status records the length and hash of the unrecorded stream. + let status = timeout(TIMEOUT, async { + loop { + let status = job_status_try_from_json_map( + client + .job_status() + .job_id(job_id) + .send() + .await + .expect("can't get job status") + .into_inner(), + ) + .expect("can't parse job status")[&test_baseboard_id()] + .clone(); + if status.is_terminal() { + break status; + } + sleep(Duration::from_millis(50)).await; + } + }) + .await + .expect("timed out waiting for job stop"); + let JobStatus::Stopped { + result, + output: + JobOutputState { + stdout_len, + stderr_len, + stdout_hash, + .. + }, + .. + } = status + else { + panic!("expected stopped status"); + }; + assert_eq!(result, Ok(0)); + assert_eq!(stdout_len, LEN as u64); + assert_eq!(stderr_len, 0); + assert_eq!(stdout_hash, blake3::hash(&streamed).into()); + + // Streamed output is never stored, so it cannot be fetched. + let Err(error) = client + .job_output() + .job_id(job_id) + .stream(JobOutputStream::Stdout) + .target(test_baseboard_id().to_string()) + .send() + .await + else { + panic!("streamed output should not be fetchable") + }; + assert_eq!(error.status().map(|status| status.as_u16()), Some(404)); +} + +#[named] +#[test] +async fn streaming_job_linger() { + // Spin up a server. + let log = test_logger(function_name!()); + let (mgr, mut root, _dir, _shutdown) = manager_and_test_root(log.clone()).await; + let api = api_description::().unwrap(); + let server = ServerBuilder::new(api, Arc::new(mgr), log) + .config(ConfigDropshot { + bind_address: local_addr(), + ..Default::default() + }) + .start() + .expect("failed to start server"); + + // Connect and authenticate to the server. + let addr = server.local_addr(); + let signer = AuthzSigner::default(); + let client = Client::new(&format!("http://{addr}"), signer.clone()); + let ClientError::ErrorResponse(unauthz) = client.iam().body(None).send().await.unwrap_err() + else { + panic!("expected error response") + }; + let (_identity, credentials) = authz(&client, unauthz, &mut root).await; + signer.set(Some(credentials)); + + // Run a streaming job small enough to finish before anyone attaches. + let session = Session::new(SessionId::random()); + client + .session_start() + .session_id(session.session_id()) + .send() + .await + .expect("can't start session"); + let job_id = session.next_job_id(); + let job = root + .sign_streaming_job_request(&job_id, "echo -n hello", test_baseboard_id().into()) + .await; + let JobLimits { + max_cpu, + max_mem, + max_fsize, + } = JobLimits::default(); + client + .job_start() + .job_id(job_id) + .target(test_baseboard_id().to_string()) + .max_cpu(max_cpu) + .max_mem(max_mem) + .max_fsize(max_fsize) + .wait(JobWait::Start) + .body(job.into_signed()) + .send() + .await + .expect("can't start job"); + sleep(Duration::from_millis(300)).await; + + // A posthumous attach still collects the whole stream. + let socket = client + .job_attach() + .job_id(job_id) + .target(test_baseboard_id().to_string()) + .send() + .await + .expect("can't attach to job") + .into_inner(); + let mut stream = WebSocketStream::from_raw_socket(socket, Role::Client, None).await; + let mut streamed = BytesMut::new(); + loop { + let recvd = timeout(TIMEOUT, stream.next()) + .await + .expect("timed out waiting for streamed output") + .expect("stream ended without close") + .expect("can't get next stream item"); + match InteractiveJobMessage::try_from(recvd).expect("can't decode message") { + InteractiveJobMessage::Data(bytes) => streamed.extend(bytes), + InteractiveJobMessage::Close => break, + message => panic!("unexpected message: {message:?}"), + } + } + assert_eq!(streamed, "hello"); + + // The status records the length and hash of the unrecorded stream. + let status = timeout(TIMEOUT, async { + loop { + let status = job_status_try_from_json_map( + client + .job_status() + .job_id(job_id) + .send() + .await + .expect("can't get job status") + .into_inner(), + ) + .expect("can't parse job status")[&test_baseboard_id()] + .clone(); + if status.is_terminal() { + break status; + } + sleep(Duration::from_millis(50)).await; + } + }) + .await + .expect("timed out waiting for job stop"); + let JobStatus::Stopped { + result, + output: + JobOutputState { + stdout_len, + stdout_hash, + .. + }, + .. + } = status + else { + panic!("expected stopped status"); + }; + assert_eq!(result, Ok(0)); + assert_eq!(stdout_len, streamed.len() as u64); + assert_eq!(stdout_hash, blake3::hash(&streamed).into()); +} diff --git a/tests/src/manager_tests.rs b/tests/src/manager_tests.rs index 57f120ad..55d0e99b 100644 --- a/tests/src/manager_tests.rs +++ b/tests/src/manager_tests.rs @@ -30,7 +30,7 @@ use sush_client::context::Authz; use sush_common::authn::{Challenge, ChallengeResponse, Credentials, Identity, Nonce, RequestKey}; use sush_common::jobs::{ Access, JobId, JobLimits, JobOutputState, JobOutputStream::*, JobStartRequest, JobStatus, - ProcessError, Session, SessionId, SignedJob, + ProcessError, Session, SessionId, SignedJob, Streaming, }; use sush_common::keys::{EphemeralKey, KeyError, KeyId, KeyType, Signer as _, pem_cert_chain}; use sush_common::targets::{Cubbies, Target}; @@ -100,7 +100,7 @@ fn check_status_stopped( #[tokio::test] async fn jobs() { let log = test_logger(function_name!()); - let (mgr, mut root, _dir, _shutdown) = manager_and_test_root(log).await; + let (mgr, mut root, dir, _shutdown) = manager_and_test_root(log).await; let baseboard_id = mgr.own_baseboard(); let authn = fake_identity(&mut root).await; let session_id = SessionId::random(); @@ -218,6 +218,38 @@ async fn jobs() { .await .is_empty() ); + + let job_id = session.next_job_id(); + let output = dir + .path() + .join("jobs") + .join(job_id.to_string()) + .display() + .to_string(); + let job = root + .sign_job_request(&job_id, "printf %s \"$SUSH_JOB_OUTPUT_DIR\"", false) + .await; + mgr.job_start( + &authn, + job.clone().into_signed(), + JobStartParams { + wait: JobWait::Stop, + ..Default::default() + }, + ) + .await + .unwrap(); + session.job_started(job.clone().into_signed()); + let status = mgr.job_status(&authn, &job_id).await.unwrap()[baseboard_id].clone(); + check_status_stopped(status, &job_id, Ok(0), Some(output.len() as u64), Some(0)); + assert_eq!( + mgr.job_output(&authn, &job_id, baseboard_id, Stdout, None) + .await + .unwrap() + .into_bytes() + .await, + output.as_bytes(), + ); } #[named] @@ -1124,7 +1156,13 @@ async fn job_targets() { let mut start = async |command: &str, target: &str| { let job_id = session.next_job_id(); - let request = JobStartRequest::new(job_id, command, false, target.parse().unwrap()); + let request = JobStartRequest::new( + job_id, + command, + false, + Streaming::None, + target.parse().unwrap(), + ); let job = root .sign(request) .await @@ -2129,3 +2167,47 @@ async fn job_json() { let recorded: SignedJob = serde_json::from_slice(&read(&path).await.unwrap()).unwrap(); assert_eq!(recorded, job.into_signed()); } + +#[named] +#[tokio::test] +async fn streaming_validation() { + let log = test_logger(function_name!()); + let (mgr, mut root, _dir, _shutdown) = manager_and_test_root(log).await; + let authn = fake_identity(&mut root).await; + let session_id = SessionId::random(); + let mut session = Session::new(session_id); + mgr.session_start(&authn, session_id, true).await.unwrap(); + + let mut expect_invalid = async |streaming, target: Target| { + let job_id = session.next_job_id(); + let job = root + .sign_full_job_request(&job_id, "true", false, streaming, target) + .await; + session.job_started(job.clone().into_signed()); + mgr.job_start(&authn, job.into_signed(), JobStartParams::default()) + .await + .unwrap(); + let status = timeout(Duration::from_secs(5), async { + loop { + if let Ok(status) = mgr.job_status(&authn, &job_id).await + && let Some(status) = status.get(mgr.own_baseboard()) + && status.is_terminal() + { + break status.clone(); + } + sleep(Duration::from_millis(10)).await; + } + }) + .await + .unwrap(); + assert!(matches!(status, JobStatus::Error { .. }), "{status:?}"); + }; + + expect_invalid(Streaming::Output, Target::All).await; + expect_invalid( + Streaming::Output, + format!("{},14", test_baseboard_id()).parse().unwrap(), + ) + .await; + expect_invalid(Streaming::Input, test_baseboard_id().into()).await; +} diff --git a/tests/src/test_utils.rs b/tests/src/test_utils.rs index 7f7a87e9..8e0ccc78 100644 --- a/tests/src/test_utils.rs +++ b/tests/src/test_utils.rs @@ -26,7 +26,7 @@ use sush_client::context::Authz; use sush_client::{Client, ResponseValue}; use sush_common::authn::{Challenge, ChallengeResponse, Credentials, Identity, Nonce, RequestKey}; use sush_common::codephrases::Codephrase; -use sush_common::jobs::{JobId, JobStartRequest, VerifiedJob}; +use sush_common::jobs::{JobId, JobStartRequest, Streaming, VerifiedJob}; use sush_common::keys::{EphemeralKey, KeyType, Signer}; use sush_common::targets::{Cubbies, Target}; use sush_server::executor::PathIsolation; @@ -82,21 +82,47 @@ pub trait SignJobRequest { command: S, interactive: bool, target: Target, + ) -> VerifiedJob { + self.sign_full_job_request(job_id, command, interactive, Streaming::None, target) + .await + } + + /// Sign a batch job request with unrecorded output streaming. + async fn sign_streaming_job_request>( + &mut self, + job_id: &JobId, + command: S, + target: Target, + ) -> VerifiedJob { + self.sign_full_job_request(job_id, command, false, Streaming::Output, target) + .await + } + + /// Sign a job request with every field specified. + async fn sign_full_job_request>( + &mut self, + job_id: &JobId, + command: S, + interactive: bool, + streaming: Streaming, + target: Target, ) -> VerifiedJob; } impl SignJobRequest for EphemeralKey { - async fn sign_job_request_for>( + async fn sign_full_job_request>( &mut self, job_id: &JobId, command: S, interactive: bool, + streaming: Streaming, target: Target, ) -> VerifiedJob { self.sign(JobStartRequest::new( job_id.to_owned(), command, interactive, + streaming, target, )) .await